"""Standalone PSSA model for the Hub: numpy-only inference, bf16 safetensors."""
import math
from dataclasses import dataclass

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F

try:
    from pssa_kernels import ScanFn
except Exception:  # no compiled kernel (CPU-only install): fall back below
    ScanFn = None


def _scan(x, delta, bp, cp, A, h0):
    """h_t = exp(delta_t·A) h_{t-1} + delta_t·b_t·x_t, y_t = <h_t, c_t>."""
    if ScanFn is not None and x.is_cuda:
        return ScanFn.apply(x, delta, bp, cp, A, h0)
    B, T, D = x.shape
    h = h0
    ys = []
    for t in range(T):
        h = torch.exp(delta[:, t, :, None] * A) * h + delta[:, t, :, None] * bp[:, t, None, :] * x[:, t, :, None]
        ys.append((h * cp[:, t, None, :]).sum(-1))
    return torch.stack(ys, 1), h


@dataclass
class PSSAConfig:
    vocab_size: int = 16384
    d_model: int = 768
    depth: int = 18
    d_state: int = 16
    d_key: int = 32
    mem_slots: int = 512
    adapter_rank: int = 16
    mlp_hidden: int = 1536
    chunk: int = 256
    seq_len: int = 512
    tau_mem: float = 0.1
    value_cap: float = 8.0


def _xavier(m):
    if isinstance(m, nn.Linear):
        bound = math.sqrt(6 / (m.in_features + m.out_features))
        nn.init.uniform_(m.weight, -bound, bound)
        if m.bias is not None:
            nn.init.zeros_(m.bias)


def _poincare(x):
    """q -> q / (1 + |q|): maps R^n into the open unit ball."""
    return x / (1.0 + x.norm(dim=-1, keepdim=True))


class PSSABlock(nn.Module):
    """SSM scan -> hyperbolic memory read -> gate -> plastic adapter -> MLP."""

    def __init__(self, cfg):
        super().__init__()
        d, ds, dk, r = cfg.d_model, cfg.d_state, cfg.d_key, cfg.adapter_rank
        self.cfg = cfg
        self.norm = nn.Parameter(torch.ones(d))
        self.norm_b = nn.Parameter(torch.zeros(d))
        self.w_delta = nn.Linear(d, d, bias=False)
        self.w_b = nn.Linear(d, ds, bias=False)
        self.w_c = nn.Linear(d, ds, bias=False)
        ratio = torch.arange(ds).float() / (ds - 1)
        tau = torch.exp(math.log(1.5) + ratio * (math.log(200.0) - math.log(1.5)))
        self.a_raw = nn.Parameter(torch.log(torch.expm1(1.0 / tau)).unsqueeze(0).repeat(d, 1))
        self.w_qx = nn.Linear(d, dk, bias=False)
        self.w_qh = nn.Linear(d, dk, bias=False)
        self.w_gate = nn.Linear(d, d, bias=False)
        self.w_proj = nn.Linear(d, d, bias=False)
        self.ad_down = nn.Linear(d, r, bias=False)
        self.ad_up = nn.Linear(r, d, bias=False)
        self.mlp1 = nn.Linear(d, cfg.mlp_hidden, bias=False)
        self.mlp2 = nn.Linear(cfg.mlp_hidden, d, bias=False)
        self.mk_raw = nn.Parameter(torch.randn(cfg.mem_slots, dk) * 0.05)
        self.mv = nn.Parameter(torch.randn(cfg.mem_slots, d) * 0.02)
        for m in (self.w_delta, self.w_b, self.w_c, self.w_qx, self.w_qh,
                  self.w_gate, self.ad_down, self.mlp1):
            _xavier(m)
        nn.init.zeros_(self.w_proj.weight)
        nn.init.zeros_(self.ad_up.weight)
        nn.init.zeros_(self.mlp2.weight)

    def forward(self, x, h0):
        cfg = self.cfg
        xn = self.norm * (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-5)) + self.norm_b
        delta = F.softplus(self.w_delta(xn).float())
        bp = self.w_b(xn).float()
        cp = self.w_c(xn).float()
        A = -F.softplus(self.a_raw)
        y, h = _scan(xn, delta, bp, cp, A, h0)

        q = _poincare((self.w_qx(xn) + self.w_qh(y)).float())
        mk = _poincare(self.mk_raw)
        qsq = (q * q).sum(-1, keepdim=True)
        mksq = (mk * mk).sum(-1)
        sq = (qsq + mksq.view(1, 1, -1) - 2 * (q @ mk.T)).clamp_min(0)
        den = ((1 - qsq) * (1 - mksq).view(1, 1, -1)).clamp_min(1e-12)
        dist = 2 * torch.asinh(torch.sqrt(sq / den))
        w = torch.softmax(-dist / cfg.tau_mem, -1)
        mv = cfg.value_cap * self.mv / (cfg.value_cap + self.mv.norm(dim=-1, keepdim=True))
        m = w @ mv

        g = torch.sigmoid(self.w_gate(xn))
        z = y / math.sqrt(cfg.d_state) + g * self.w_proj(m.to(x.dtype))
        z = z + self.ad_up(F.silu(self.ad_down(xn)))
        z = z + self.mlp2(F.silu(self.mlp1(z)))
        return z, h


class PSSAModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList([PSSABlock(cfg) for _ in range(cfg.depth)])
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.scale = 1.0 / math.sqrt(cfg.depth)

    def init_state(self, batch_size, device=None):
        dev = device or self.embed.weight.device
        return [torch.zeros(batch_size, self.cfg.d_model, self.cfg.d_state, device=dev)
                for _ in range(self.cfg.depth)]

    def forward_step(self, tokens, states):
        x = self.embed(tokens)
        for i, blk in enumerate(self.blocks):
            z, h = blk(x, states[i])
            x = z if i == 0 else x + z * self.scale
            states[i] = h
        return self.head(x) / math.sqrt(self.cfg.d_model), states

    def forward(self, tokens):
        states = self.init_state(tokens.shape[0], tokens.device)
        outs = []
        for s in range(0, tokens.shape[1], self.cfg.chunk):
            logits, states = self.forward_step(tokens[:, s:s + self.cfg.chunk], states)
            outs.append(logits)
        return torch.cat(outs, 1)


def from_pretrained(model_dir: str, device: str = "cpu"):
    """Load the Hub export: (model, tokenizer). Works on CPU with no Triton install."""
    from safetensors.torch import load_file
    from tokenizers import Tokenizer

    cfg = PSSAConfig()
    sd = load_file(f"{model_dir}/model.safetensors")
    sd = {k: v.float() for k, v in sd.items()}
    m = PSSAModel(cfg)
    m.load_state_dict(sd)
    m.eval().to(device)
    return m, Tokenizer.from_file(f"{model_dir}/tokenizer.json")


@torch.no_grad()
def generate(model, tokenizer, prompt, max_new=100, temp=0.7, top_k=40, eot=0, device="cpu"):
    ids = tokenizer.encode(prompt).ids
    states = model.init_state(1, device)
    logits, states = model.forward_step(torch.tensor([ids], device=device), states)
    out = []
    for _ in range(max_new):
        lg = logits[:, -1, :].float() / temp
        if top_k:
            kth = torch.topk(lg, top_k).values[:, [-1]]
            lg = lg.masked_fill(lg < kth, -float("inf"))
        nxt = torch.multinomial(torch.softmax(lg, -1), 1)
        if int(nxt) == eot:
            break
        out.append(int(nxt))
        logits, states = model.forward_step(nxt, states)
    text = tokenizer.decode([i for i in out if i != eot])
    return prompt + text