import torch
import torch.nn as nn
from pathlib import Path


class SparseAutoencoder(nn.Module):
    def __init__(self, d_model: int = 2048, n_features: int = 16384):
        super().__init__()
        self.d_model = d_model
        self.n_features = n_features

        self.b_pre = nn.Parameter(torch.zeros(d_model))
        self.W_enc = nn.Parameter(torch.nn.init.kaiming_uniform_(torch.empty(d_model, n_features)))
        self.b_enc = nn.Parameter(torch.zeros(n_features))
        self.W_dec = nn.Parameter(torch.nn.init.kaiming_uniform_(torch.empty(n_features, d_model)))
        self.b_dec = nn.Parameter(torch.zeros(d_model))

        self._normalise_decoder()

    def _normalise_decoder(self):
        with torch.no_grad():
            norms = self.W_dec.norm(dim=-1, keepdim=True).clamp(min=1e-8)
            self.W_dec.data = self.W_dec.data / norms

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        return torch.relu((x - self.b_pre) @ self.W_enc + self.b_enc)

    def decode(self, f: torch.Tensor) -> torch.Tensor:
        return f @ self.W_dec + self.b_dec

    def forward(self, x: torch.Tensor):
        f = self.encode(x)
        return f, self.decode(f)

    @torch.no_grad()
    def get_top_features(self, x: torch.Tensor, k: int = 10):
        f = self.encode(x.unsqueeze(0)).squeeze(0)
        topk = f.topk(k)
        return [(int(i), float(v)) for i, v in zip(topk.indices, topk.values)]

    def save(self, path):
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"d_model": self.d_model, "n_features": self.n_features, "state_dict": self.state_dict()}, path)
        print(f"[SAE] saved to {path}", flush=True)

    @classmethod
    def load(cls, path, device="cuda"):
        ckpt = torch.load(path, map_location=device)
        sae = cls(d_model=ckpt["d_model"], n_features=ckpt["n_features"])
        sae.load_state_dict(ckpt["state_dict"])
        sae.to(device)
        sae.eval()
        print(f"[SAE] loaded from {path}", flush=True)
        return sae
