"""Amos: encoder + typed decision heads. Every output is an option, a span of the state, or a P(yes).

All heads read the question's [DECIDE_t] vector:
  noul          softmax over two implicit options ["no: ...", "yes: ..."]; P(yes) is the second
  choice/score  softmax over bilinear(h_decide, mean-pooled option tokens)
  multi         independent sigmoid over the same pointer logits (+ bias)
  span          start/end softmax, bilinear(h_decide, state tokens); decoded to an exact substring
  region        softmax over the first state image's 49 cells, bilinear(h_decide, image tokens)

With a vision tower, image slots are filled with projected vision features before the first layer;
the tower itself stays frozen and is not part of saved checkpoints.
"""

import json
import math
from pathlib import Path

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

from .encoder import Config, Encoder
from .facts import date_facts
from .packing import (
    IMAGE_TOKENS,
    TYPES,
    Codec,
    Packed,
    char_span_to_tokens,
    collate,
    pack,
)
from .vision import Projector, VisionTower, box_target, cell_boxes

NOUL, CHOICE, SCORE, MULTI, SPAN, REGION = range(len(TYPES))
MAX_SPAN_TOKENS = 30


def _init(*mods):
    for m in mods:
        nn.init.normal_(m.weight, std=0.02)
        nn.init.zeros_(m.bias)


class Amos(nn.Module):
    def __init__(self, encoder: Encoder, vision: VisionTower | None = None):
        super().__init__()
        d = encoder.config.hidden_size
        self.encoder = encoder
        self.ptr_q, self.ptr_o = nn.Linear(d, d), nn.Linear(d, d)
        self.multi_bias = nn.Parameter(torch.zeros(()))
        self.span_q, self.span_k = nn.Linear(d, 2 * d), nn.Linear(d, 2 * d)
        _init(self.ptr_q, self.ptr_o, self.span_q, self.span_k)
        self.vision = vision
        if vision is not None:
            self.projector = Projector(vision.hidden_size, d)
            self.region_q, self.region_k = nn.Linear(d, d), nn.Linear(d, d)
            _init(self.region_q, self.region_k)

    @classmethod
    def from_encoder(cls, repo: str, codec: Codec, vision_repo: str | None = None) -> "Amos":
        enc = Encoder.from_pretrained(repo)
        init_markers(enc, codec)
        return cls(enc, VisionTower.from_pretrained(vision_repo) if vision_repo else None)

    def saved_state(self) -> dict:
        """Weights to checkpoint: everything except the frozen, separately downloadable vision tower."""
        return {k: v for k, v in self.state_dict().items() if not k.startswith("vision.")}

    def embed(self, b: dict) -> torch.Tensor:
        emb = self.encoder.tok_embeddings(b["ids"])
        if "images" not in b:
            return emb
        x = b["images"]
        if x.dim() == 4:  # pixels -> frozen tower; precomputed features skip it
            with torch.no_grad():
                x = self.vision(x)
        tokens = self.projector(x.to(emb.dtype))
        cols = b["img_at"][:, None] + torch.arange(IMAGE_TOKENS, device=emb.device)
        rows = b["img_row"][:, None].expand_as(cols)
        return emb.index_put((rows.reshape(-1), cols.reshape(-1)), tokens.reshape(-1, tokens.shape[-1]))

    def forward(self, b: dict) -> dict:
        return self.heads(self.encoder(b["ids"], b["pos"], b["seg"], b["opt"], inputs_embeds=self.embed(b)), b)

    def forward_cached(self, b: dict, state_len: int, cache: dict | None):
        """One pack (B = 1) with a reusable state. Returns (outputs, state cache); same result as forward().

        The state never attends to the questions, so its keys, values and hidden states in a full pass are
        exactly what a state-only pass would produce: a miss runs one normal pass and slices the cache out
        of it; a hit encodes only the question tokens on top of the cached state.
        """
        emb, s = self.embed(b), state_len
        if cache is None:
            h, kv = self.encoder.run(emb, b["pos"], b["seg"], b["opt"])
            cache = {"kv": [(k[:, :, :s], v[:, :, :s]) for k, v in kv["kv"]], "hidden": h[:, :s],
                     **{k: b[k][:, :s] for k in ("pos", "seg", "opt")}}
            return self.heads(h, b), cache
        h_q, _ = self.encoder.run(emb[:, s:], b["pos"][:, s:], b["seg"][:, s:], b["opt"][:, s:], cache)
        return self.heads(torch.cat([cache["hidden"], h_q], 1), b), cache

    def heads(self, h: torch.Tensor, b: dict) -> dict:
        d = h.shape[-1]
        hq = h[b["q_row"], b["q_tok"]]
        out = {}

        Q, O = len(hq), len(b["o_q"])
        opt_logits = h.new_full((Q, opt_width(b)), -torch.inf)
        if O:
            of = b["option_of"]
            valid = of >= 0
            pooled = h.new_zeros(O, d).index_add_(0, of[valid], h[valid])
            pooled = pooled / torch.bincount(of[valid], minlength=O).clamp(min=1)[:, None]
            ptr = (self.ptr_q(hq)[b["o_q"]] * self.ptr_o(pooled)).sum(-1) / math.sqrt(d)
            ptr = ptr + (b["q_type"][b["o_q"]] == MULTI) * self.multi_bias
            opt_logits = opt_logits.index_put((b["o_q"], b["o_rank"]), ptr)
        out["opt"] = opt_logits

        sq = (b["q_type"] == SPAN).nonzero().squeeze(-1)
        out["span_q"] = sq
        if len(sq):
            rows = b["q_row"][sq]
            qs = self.span_q(hq[sq]).view(len(sq), 2, d)
            ks = self.span_k(h[rows]).view(len(sq), -1, 2, d)
            logits = torch.einsum("qcd,qlcd->qcl", qs, ks) / math.sqrt(d)
            out["span"] = logits.masked_fill(~b["state_mask"][rows][:, None, :], -torch.inf)

        rq = (b["q_type"] == REGION).nonzero().squeeze(-1)
        out["region_q"] = rq
        if len(rq):
            rows = b["q_row"][rq]
            logits = torch.einsum("qd,qld->ql", self.region_q(hq[rq]), self.region_k(h[rows])) / math.sqrt(d)
            out["region"] = logits.masked_fill(~b["region_mask"][rows], -torch.inf)
        return out


def init_markers(enc: nn.Module, codec: Codec, markers=None):
    """Reserved markers start as copies of [CLS] (their pretrained rows were never trained)."""
    with torch.no_grad():
        for i in markers or (*codec.decide.values(), codec.opt, codec.img, codec.img_end, codec.patch):
            enc.tok_embeddings.weight[i] = enc.tok_embeddings.weight[codec.cls]


def add_targets(b: dict, packs: list[Packed], questions: list[list[dict]]) -> dict:
    """Attach target tensors to a collated batch. Targets are probabilities (hard labels are one-hot).

    noul: float P(yes). choice/score: distribution over options. multi: independent P per option.
    span: list of accepted [start, end) char spans in the state. region: normalized (x0, y0, x1, y1) box
    in the first state image, spread over the 49 cells by overlap.
    """
    Q, L = len(b["q_row"]), b["ids"].shape[1]
    t_opt = torch.zeros(Q, opt_width(b))
    t_span, t_region, has = [], [], torch.ones(Q, dtype=torch.bool)
    qi = 0
    for row, (p, qs) in enumerate(zip(packs, questions)):
        ok = b["state_mask"][row]
        for q in qs:
            tgt, typ = q["target"], q["type"]
            if typ == "noul":  # implicit options are ["no", "yes"]
                t_opt[qi, :2] = torch.tensor([1 - float(tgt), float(tgt)])
            elif typ == "region":
                dist = torch.zeros(L)
                dist[p.image_at[0] : p.image_at[0] + IMAGE_TOKENS] = box_target(tgt)
                t_region.append(dist)
            elif typ == "span":
                dist = torch.zeros(2, L)
                for s, e in tgt:
                    tok = char_span_to_tokens(p.offsets, s, e)
                    if tok and ok[tok[0]] and ok[tok[1]]:  # boundaries must be decodable (not truncated/blank)
                        dist[0, tok[0]] += 1
                        dist[1, tok[1]] += 1
                has[qi] = bool(dist.sum())
                t_span.append(dist / dist.sum(-1, keepdim=True).clamp(min=1))
            else:
                t_opt[qi, : len(tgt)] = torch.tensor(tgt, dtype=torch.float)
            qi += 1
    b["t_opt"], b["has_target"] = t_opt, has
    b["t_span"] = torch.stack(t_span) if t_span else torch.zeros(0, 2, L)
    b["t_region"] = torch.stack(t_region) if t_region else torch.zeros(0, L)
    return b


def opt_width(b: dict) -> int:
    return max(1, int(b["q_nopt"].max())) if len(b["q_nopt"]) else 1


def emd2(p: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """Squared earth mover's distance between ordinal distributions [N, K] (padded with zeros)."""
    return (p.cumsum(-1) - t.cumsum(-1)).pow(2).sum(-1)


def soft_ce(logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Cross-entropy against a (soft) target distribution; masked (-inf) logits must have zero target."""
    logp = logits.log_softmax(-1).masked_fill(target == 0, 0)
    return -(target * logp).sum(-1)


def brier(p: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """Half the squared distance between distributions [N, K]; padded options have p = t = 0."""
    return 0.5 * (p - t).pow(2).sum(-1)


def loss_fn(out: dict, b: dict, ordinal_weight: float = 1.0, brier_weight: float = 0.0,
            type_weights: dict | None = None) -> tuple[torch.Tensor, dict]:
    qt, has = b["q_type"], b["has_target"]
    parts = {}
    for typ in (NOUL, CHOICE, SCORE):
        m = (qt == typ) & has
        if m.any():
            ce = soft_ce(out["opt"][m], b["t_opt"][m])
            if brier_weight:
                ce = ce + brier_weight * brier(out["opt"][m].softmax(-1), b["t_opt"][m])
            if typ == SCORE:
                ce = ce + ordinal_weight * emd2(out["opt"][m].softmax(-1), b["t_opt"][m])
            parts[TYPES[typ]] = ce
    m = (qt == MULTI) & has
    if m.any():
        logits, valid = out["opt"][m], out["opt"][m].isfinite()
        bce = F.binary_cross_entropy_with_logits(logits.masked_fill(~valid, 0), b["t_opt"][m], reduction="none")
        parts["multi"] = (bce * valid).sum(-1) / valid.sum(-1)
    if len(out["span_q"]):
        m = has[out["span_q"]]
        if m.any():
            parts["span"] = soft_ce(out["span"][m], b["t_span"][m]).mean(-1)
    if len(out["region_q"]):
        parts["region"] = soft_ce(out["region"], b["t_region"])
    w = type_weights or {}
    total = torch.cat([v * w.get(k, 1.0) for k, v in parts.items()]).mean()
    return total, {k: v.mean().item() for k, v in parts.items()}


def to_device(b: dict, device) -> dict:
    return {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in b.items()}


def decode_span(start: torch.Tensor, end: torch.Tensor, p: Packed, length_penalty: float = 0.0,
                max_tokens: int = MAX_SPAN_TOKENS):
    """Best valid span under independent start/end softmax. Returns (char_start, char_end, text, prob).
    length_penalty: each token past the start scales a span's score by exp(-length_penalty). max_tokens: the longest
    span considered.

    The answer is always p.state[char_start:char_end]: it is cut from the input, never generated.
    """
    n = len(p.offsets)
    if not n or not start.isfinite().any():
        return 0, 0, "", 0.0
    ps, pe = start.softmax(-1)[1 : n + 1], end.softmax(-1)[1 : n + 1]  # state text tokens only
    K = min(max_tokens, n)
    pe_pad = F.pad(pe, (0, K))
    joint = torch.stack([ps * pe_pad[k : k + n] for k in range(K)], 1)  # [start, length - 1]
    joint = joint * torch.exp(-length_penalty * torch.arange(K, dtype=joint.dtype))
    joint = joint / joint.sum().clamp(min=1e-12)
    s, k = divmod(int(joint.argmax()), K)
    prob = float(joint[s, k])
    a, z = p.offsets[s][0], p.offsets[s + k][1]
    raw = p.state[a:z]
    a, z = a + len(raw) - len(raw.lstrip()), z - (len(raw) - len(raw.rstrip()))
    return a, max(a, z), p.state[a : max(a, z)], prob


def decode(out: dict, packs: list[Packed]) -> list[list[dict]]:
    """Turn a batch forward into one typed result per question, grouped per pack."""
    opt = out["opt"].float().cpu()
    span = out["span"].float().cpu() if "span" in out else None
    region = out["region"].float().cpu() if "region" in out else None
    res, qi, si, ri = [], 0, 0, 0
    for p in packs:
        rs = []
        for t, n in zip(p.types, p.n_options):
            if t == "noul":
                rs.append({"type": t, "p_yes": opt[qi, :2].softmax(-1)[1].item()})
            elif t in ("choice", "score"):
                probs = opt[qi, :n].softmax(-1)
                r = {"type": t, "probs": probs.tolist(), "choice": int(probs.argmax())}
                if t == "score":
                    r["expected"] = float((probs * torch.arange(n)).sum())
                rs.append(r)
            elif t == "multi":
                rs.append({"type": t, "probs": torch.sigmoid(opt[qi, :n]).tolist()})
            elif t == "region":
                cells = region[ri].softmax(-1)[p.image_at[0] : p.image_at[0] + IMAGE_TOKENS]
                k = int(cells.argmax())
                rs.append({"type": t, "cell": k, "box": cell_boxes()[k].tolist(), "p": float(cells[k]), "heatmap": cells.tolist()})
                ri += 1
            else:
                a, z, text, prob = decode_span(span[si, 0], span[si, 1], p)
                rs.append({"type": t, "start": a, "end": z, "text": text, "p": prob})
                si += 1
            qi += 1
        res.append(rs)
    return res


def load_run(run, device="cpu"):
    """A trained run directory (config.json + model.pt) -> (model in eval mode, its codec)."""
    cfg = json.loads((Path(run) / "config.json").read_text())
    enc = cfg["encoder"]
    if cfg.get("readout", "pointer") != "pointer" or enc.get("question_reads") or enc.get("model_type", "modernbert") != "modernbert":
        raise ValueError(f"{run}: trained with a removed mode (MLM readout, question reads or DeBERTa)")
    codec = Codec(cfg["repo"], max_state=cfg["max_state"], digits=cfg.get("digits", False),
                  score_rank=cfg.get("score_rank", False))
    vision = VisionTower.from_pretrained(cfg["vision"]) if cfg.get("vision") else None
    model = Amos(Encoder(Config.from_hf(enc)), vision)
    missing, unexpected = model.load_state_dict(torch.load(Path(run) / "model.pt", map_location="cpu", weights_only=True), strict=False)
    if unexpected or any(not k.startswith("vision.") for k in missing):
        raise ValueError(f"checkpoint mismatch: missing {missing}, unexpected {unexpected}")
    return model.to(device).eval(), codec


def pack_request(model: Amos, codec: Codec, state: str, questions: list[dict], images=(), *, strict=False) -> list[Packed]:
    """One pack holding every question (a list, so callers batch requests as rows)."""
    if model.encoder.config.date_facts:
        state = date_facts(state)
    return [pack(codec, state, questions, images, strict=strict)]


@torch.no_grad()
def predict(model: Amos, codec: Codec, state: str, questions: list[dict], device="cpu", images=()) -> list[dict]:
    """Run one state with its questions; returns one typed result per question."""
    ps = pack_request(model, codec, state, questions, images)
    return [r for rs in decode(model(to_device(collate(ps, codec.pad), device)), ps) for r in rs]