# standalone inference for the public Hub repo (torch + safetensors + tokenizers only)
import json, torch
from safetensors.torch import load_file
from tokenizers import Tokenizer


class FlyBrainLM(torch.nn.Module):
    def __init__(self, vocab=4096, n_fly=48000, e_in=160, leak=0.3):
        super().__init__()
        self.n, self.leak = n_fly, leak
        self.register_buffer("Widx", torch.zeros(2, 1, dtype=torch.long))
        self.Wval = torch.nn.Parameter(torch.zeros(1))
        self.Win = torch.nn.Parameter(torch.zeros(n_fly, e_in), requires_grad=False)
        self.emb = torch.nn.Embedding(vocab, e_in)
        self.lens1 = torch.nn.Linear(n_fly, 1024)
        self.lens2 = torch.nn.Linear(1024, vocab)

    def forward(self, idx):
        W = torch.sparse_coo_tensor(self.Widx, self.Wval, (self.n, self.n)).coalesce()
        e = self.emb(idx)
        s = torch.zeros(idx.shape[0], self.n, device=idx.device)
        outs = []
        for t in range(idx.shape[1]):
            s = (1 - self.leak) * s + self.leak * torch.tanh(
                e[:, t, :] @ self.Win.T + torch.sparse.mm(W, s.T).T)
            outs.append(s)
        return self.lens2(torch.tanh(self.lens1(torch.stack(outs, 1))))


def load_model(repo_or_dir, device="cuda"):
    try:
        from huggingface_hub import snapshot_download
        import os
        d = repo_or_dir if os.path.isdir(repo_or_dir) else snapshot_download(repo_or_dir)
    except Exception:
        d = repo_or_dir
    with open(f"{d}/config.json") as f:
        cfg = json.load(f)
    tok = Tokenizer.from_file(f"{d}/tokenizer.json")
    m = FlyBrainLM(cfg["vocab"], cfg["n_fly"], cfg["e_in"], cfg.get("leak", 0.3))
    m.load_state_dict(load_file(f"{d}/model.safetensors", device="cpu"))
    return m.to(device).eval(), tok, cfg


def generate(model, tok, prompt, max_new=120, temp=0.8, ctx_win=160):
    dev = next(model.parameters()).device
    ids = tok.encode(prompt).ids[-ctx_win:]
    W = torch.sparse_coo_tensor(model.Widx, model.Wval, (model.n, model.n)).coalesce()
    with torch.no_grad():
        e = model.emb(torch.tensor([ids], device=dev))
        s = torch.zeros(1, model.n, device=dev)
        for t in range(len(ids)):
            s = 0.7 * s + 0.3 * torch.tanh(e[:, t, :] @ model.Win.T + torch.sparse.mm(W, s.T).T)
        last, out = ids[-1], []
        for _ in range(max_new):
            e1 = model.emb(torch.tensor([[last]], device=dev))
            s = 0.7 * s + 0.3 * torch.tanh(e1[:, 0, :] @ model.Win.T + torch.sparse.mm(W, s.T).T)
            p = torch.softmax(model.lens2(torch.tanh(model.lens1(s)))[0] / temp, -1)
            nx = int(torch.multinomial(p, 1))
            if nx == tok.encode("[EOS]").ids[0]:
                break
            out.append(nx)
            last = nx
    txt = tok.decode(out)
    cut = txt.find("\nHuman:")
    return txt[:cut].strip() if cut >= 0 else txt.strip()


if __name__ == "__main__":
    import sys
    repo = sys.argv[1] if len(sys.argv) > 1 else "."
    m, tok, cfg = load_model(repo)
    print(generate(m, tok, "Human: Hi Fly! Who are you?\nFly:"))