"""
daisychain.py -- self-contained inference for the DaisyChain genomic modular mind.

4 dense ~74M DNA/RNA specialists (eukaryote, prokaryote, mrna, mrna_splice), each
per-domain-distilled from Carbon-500M, behind a learned router (MLP on PCA(hidden)
+ per-specialist surprise). route() picks the home specialist; generate() / surprise()
expose the rest. No training/datasets dependency -- only model.py, specialist_presets.py,
spike_tokenizer.py, registry.py + the bundled tokenizer.json / *.safetensors / router2.pt.
"""
from __future__ import annotations
import os, math
import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
from model import SpikeWhaleLM
from specialist_presets import generic_specialist_config
from spike_tokenizer import SpikeTokenizer
import registry

TOK_JSON = os.path.join(HERE, "tokenizer.json")
_TRANS = str.maketrans({"U": "T", "u": "T", "a": "A", "c": "C", "g": "G", "t": "T"})
LN2 = math.log(2)


def clean(seq: str) -> str:
    seq = seq.translate(_TRANS).upper()
    return "".join(c if c in "ACGT" else "N" for c in seq)


class _RouterMLP(torch.nn.Module):
    def __init__(self, dim, h=64):
        super().__init__()
        self.net = torch.nn.Sequential(torch.nn.Linear(dim, h), torch.nn.ReLU(),
                                       torch.nn.Dropout(0.0), torch.nn.Linear(h, 4))
    def forward(self, x): return self.net(x)


class DaisyChain:
    DESCRIPTIONS = {
        "eukaryote": "Eukaryotic genomic DNA",
        "prokaryote": "Bacterial / prokaryotic DNA",
        "mrna": "Mature mRNA (coding transcript)",
        "mrna_splice": "Pre-mRNA / splice-site regions",
    }

    def __init__(self, root=HERE, device="cpu"):
        self.dev = device
        self.tok = SpikeTokenizer(vocab_file=os.path.join(root, "tokenizer.json"))
        self.bos, self.eos = self.tok._vocab["<bos>"], self.tok._vocab["<eos>"]
        self.models = {}
        from safetensors.torch import load_file
        for d in registry.ACTIVE:
            ckpt = os.path.join(root, d, "model.safetensors")
            if not os.path.exists(ckpt):
                continue
            cfg = generic_specialist_config(self.tok.vocab_size, position=registry.spec(d)["position"])
            m = SpikeWhaleLM(cfg).to(device).eval()
            sd = load_file(ckpt, device=device)
            m.load_state_dict({k: (v.float() if v.is_floating_point() else v) for k, v in sd.items()})
            for p in m.parameters():
                p.requires_grad_(False)
            self.models[d] = m
        self.domains = list(self.models)
        # FNS base-pair tables: map our 4096 6-mer tokens to their six per-position bases, so we
        # can marginalize the 6-mer softmax into six 4-way nucleotide distributions and decode /
        # score at the BASE level — the same factorization Carbon's FNS branch uses.
        import itertools
        _b2i = {"A": 0, "T": 1, "C": 2, "G": 3}
        self._idx2base = "ATCG"
        kmers = ["".join(t) for t in itertools.product("ACGT", repeat=6)]
        self._kmer_ids = torch.tensor([self.tok._vocab[k] for k in kmers], device=device)  # [4096]
        self._base_at = torch.zeros(6, 4096, dtype=torch.long, device=device)
        for i, k in enumerate(kmers):
            for pos in range(6):
                self._base_at[pos, i] = _b2i[k[pos]]
        self._kmer_id_of = {k: self.tok._vocab[k] for k in kmers}
        self.router2 = None
        r2 = os.path.join(root, "router2.pt")
        if os.path.exists(r2):
            d = torch.load(r2, map_location="cpu")
            if all(x in self.models for x in d["domains"]):
                mlp = _RouterMLP(d["k"] + 4, d["h"]); mlp.load_state_dict(d["mlp"]); mlp.eval()
                d["net"] = mlp; self.router2 = d

    @torch.no_grad()
    def _scores_hidden(self, seq):
        ids = [self.bos] + self.tok.encode(clean(seq), add_special_tokens=False) + [self.eos]
        t = torch.tensor([ids], device=self.dev)
        scores, hids = {}, {}
        for d, m in self.models.items():
            hids[d] = m.model(input_ids=t)[0][0].mean(0)
            scores[d] = float(m(input_ids=t, labels=t).loss)
        return scores, hids

    def surprise(self, seq):
        """Per-specialist bits/base (lower = more 'at home')."""
        s, _ = self._scores_hidden(seq)
        return {d: s[d] / 6 / LN2 for d in self.domains}

    @torch.no_grad()
    def route(self, seq):
        """Return (home_domain, bits_per_base_dict). Uses the learned MLP router."""
        scores, hids = self._scores_hidden(seq)
        bpb = {d: scores[d] / 6 / LN2 for d in self.domains}
        if self.router2 is not None:
            r = self.router2
            hidden = torch.cat([hids[d] for d in r["domains"]])
            bits = torch.tensor([scores[d] for d in r["domains"]])
            z = (hidden - r["pca_mu"]) @ r["P"]
            feat = ((torch.cat([z, bits]) - r["mu"]) / r["sd"])
            best = r["domains"][int(r["net"](feat.unsqueeze(0)).argmax(1))]
        else:
            best = min(scores, key=scores.get)
        return best, bpb

    @torch.no_grad()
    def generate_stream(self, domain, length=180, temperature=1.0, top_k=40,
                        repetition_penalty=1.3, prompt="", greedy=False, top_p=0.9):
        """Yield the growing continuation base-by-base (for live streaming UIs).
        greedy=True takes the argmax 6-mer each step (deterministic — the model's single
        best guess, same decoding as the sequence-recovery metric; the coding domains
        produce in-frame ATG-start sequences, the low-complexity domains collapse to
        homopolymers). Sampling (greedy=False) trades that for variety; repetition_penalty
        then discourages the repeat loops these small specialists fall into."""
        m = self.models[domain]
        # frame-align + cap: our 6-mer tokens tile cleanly only when the context length is a
        # multiple of 6 (else the model generates out of phase). Trim the leading remainder and
        # cap to the context window — this is what makes generation match the offline tests.
        p = clean(prompt) if prompt else ""
        p = p[-1020:]                 # 1020 = 170*6, within the 1024-token window
        p = p[len(p) % 6:]            # drop leading remainder so the 6-mers align to the end
        ids = [self.bos] + (self.tok.encode(p, add_special_tokens=False) if p else [])
        t = torch.tensor([ids], device=self.dev)
        bases, emitted = [], []
        while sum(len(b) for b in bases) < length:
            logits = m(input_ids=t).logits[:, -1, :].float()
            logits[:, :4] = -1e9                                   # never emit specials
            # Repetition control applies to BOTH decoders. Greedy without this falls into a
            # self-reinforcing loop (argmax keeps re-picking the same 6-mer); the penalty plus
            # a hard block on the last few emitted tokens forces it onward to its next-best,
            # non-repeating guess — which is what actually de-degenerates the output.
            if greedy:
                # argmax with repeat penalty + a hard block on recent tokens (greedy alone loops)
                if repetition_penalty and repetition_penalty != 1.0:
                    for tid in set(emitted[-12:]):
                        v = logits[0, tid]
                        logits[0, tid] = v / repetition_penalty if v > 0 else v * repetition_penalty
                for tid in set(emitted[-8:]):
                    logits[0, tid] = -1e9
                ti = int(logits.argmax())
            else:
                # nucleus (top-p) sampling — the same decoder Carbon uses. Adapts the candidate
                # set to the distribution (keeps only tokens carrying mass), which escapes the
                # low-complexity GC/AT loops that a fixed top-k falls into. No repetition penalty.
                logits = logits / max(temperature, 1e-6)
                sl, si = torch.sort(logits[0], descending=True)
                cum = torch.cumsum(F.softmax(sl, dim=-1), dim=-1)
                rm = cum > top_p
                rm[1:] = rm[:-1].clone(); rm[0] = False
                logits[0, si[rm]] = -1e9
                ti = int(torch.multinomial(F.softmax(logits[0], dim=-1), 1))
            emitted.append(ti)
            t = torch.cat([t, torch.tensor([[ti]], device=self.dev)], dim=1)
            bases.append(self.tok._ids_to_tokens[ti])
            yield "".join(bases)[:length]

    def generate(self, domain, length=180, temperature=1.0, top_k=40,
                 repetition_penalty=1.3, prompt="", greedy=False):
        out = ""
        for out in self.generate_stream(domain, length, temperature, top_k, repetition_penalty, prompt, greedy):
            pass
        return out

    # ---- FNS base-pair-level decode + score (Carbon's factorized-nucleotide approach) ----
    def _bp_marginals(self, logits, top_p=1.0):
        """Marginalize a 6-mer logit vector into six 4-way per-position base distributions [6,4].
        Matching Carbon's _BPLogitsProcessor: any top-p filtering happens at the 6-MER level
        (before marginalizing), not per-base."""
        kl = logits[self._kmer_ids]                            # [4096] 6-mer logits
        if top_p < 1.0:                                        # nucleus on the 6-mers first
            sk, si = torch.sort(kl, descending=True)
            rm = torch.cumsum(F.softmax(sk, -1), -1) > top_p
            rm[1:] = rm[:-1].clone(); rm[0] = False
            kl = kl.clone(); kl[si[rm]] = float("-inf")
        p = F.softmax(kl, dim=-1)
        bp = torch.zeros(6, 4, device=logits.device)
        for pos in range(6):
            bp[pos].scatter_add_(0, self._base_at[pos], p)
        return bp

    @torch.no_grad()
    def generate_baselevel_stream(self, domain, length=180, temperature=1.0, top_p=0.9,
                                  prompt="", greedy=False):
        """Base-pair generation exactly as Carbon's FNS BpLogitsProcessor: apply temperature +
        top-p at the 6-mer level, marginalize the filtered 6-mer softmax to six 4-way base
        distributions, then pick each base by multinomial (sampling) or argmax (greedy)."""
        m = self.models[domain]
        p = clean(prompt) if prompt else ""
        p = p[-1020:]; p = p[len(p) % 6:]
        ids = [self.bos] + (self.tok.encode(p, add_special_tokens=False) if p else [])
        t = torch.tensor([ids], device=self.dev)
        out = ""
        while len(out) < length:
            logits = m(input_ids=t).logits[0, -1].float() / max(temperature, 1e-6)
            bp = self._bp_marginals(logits, top_p=(1.0 if greedy else top_p))   # [6,4]
            if greedy:
                six = "".join(self._idx2base[int(bp[pos].argmax())] for pos in range(6))
            else:
                six = "".join(self._idx2base[int(torch.multinomial(bp[pos], 1))] for pos in range(6))
            t = torch.cat([t, torch.tensor([[self._kmer_id_of[six]]], device=self.dev)], dim=1)
            out += six
            yield out[:length]

    @torch.no_grad()
    def score(self, domain, seq):
        """Carbon `score_sequence` equivalent: mean per-base log-prob of the observed sequence under
        the base-level (marginalized) distribution. Higher = more likely. bits/base = -score/ln2.
        Right-pads to a multiple of 6 with 'A' (Carbon's convention) and scores only real bases."""
        m = self.models[domain]
        s = clean(seq)
        orig = len(s)
        r = len(s) % 6
        if r:
            s = s + "A" * (6 - r)                              # right-pad like score_sequence
        if orig < 12:
            return float("nan")
        ids = [self.bos] + self.tok.encode(s, add_special_tokens=False)
        t = torch.tensor([ids], device=self.dev)
        logits = m(input_ids=t).logits[0]                      # [L,V]; position i predicts the next 6-mer
        b2i = {"A": 0, "T": 1, "C": 2, "G": 3}
        tot, n = 0.0, 0
        for i in range(len(ids) - 1):
            if i * 6 >= orig:                                  # don't score the padding
                break
            bp = self._bp_marginals(logits[i].float())         # [6,4] predicted bases of next 6-mer
            nxt = s[i * 6:(i + 1) * 6]
            for pos, ch in enumerate(nxt):
                if i * 6 + pos >= orig:                          # padding base — don't score
                    break
                if ch not in b2i:                               # skip N / ambiguous bases
                    continue
                tot += math.log(max(float(bp[pos, b2i[ch]]), 1e-12)); n += 1
        return tot / max(n, 1)                                  # mean per-base logp