"""Minimal loader for the TopK SAEs in this repository.

Usage:
    from loader import load_sae
    sae, cfg = load_sae("layer6/sae_instruct_base_layer6.pt", device="cuda")
    x_hat, z = sae(x)   # x: (N, d_model=896)

The `sae(x)` forward returns:
    x_hat:    (N, d_model)   reconstruction
    z_sparse: (N, d_sae)     sparse code, exactly `k` non-zeros per row

When splicing into the base model's residual stream, only replace
real-token positions (see README for the rationale):
    patched = torch.where(mask.unsqueeze(-1).bool(), sae(h)[0], h)
"""

import torch
import torch.nn as nn


class TopKSAE(nn.Module):
    def __init__(self, d_model: int, d_sae: int, k: int):
        super().__init__()
        self.k = k
        self.d_model = d_model
        self.d_sae = d_sae
        self.b_pre = nn.Parameter(torch.zeros(d_model))
        self.encoder = nn.Linear(d_model, d_sae, bias=True)
        self.decoder = nn.Linear(d_sae, d_model, bias=True)

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        z = self.encoder(x - self.b_pre)
        topk_values, topk_indices = torch.topk(z, self.k, dim=-1)
        z_sparse = torch.zeros_like(z)
        z_sparse.scatter_(-1, topk_indices, topk_values)
        return z_sparse

    def forward(self, x: torch.Tensor):
        z_sparse = self.encode(x)
        return self.decoder(z_sparse), z_sparse


def load_sae(path: str, device: str = "cpu"):
    """Load a checkpoint saved by the training pipeline.

    Checkpoint format:
        {"state_dict": ..., "config": {"d_model", "d_sae", "k", "source"}}
    """
    ckpt = torch.load(path, map_location=device, weights_only=False)
    cfg = ckpt["config"]
    sae = TopKSAE(cfg["d_model"], cfg["d_sae"], cfg["k"])
    sae.load_state_dict(ckpt["state_dict"], strict=False)
    return sae.to(device).eval(), cfg