"""PrismeModel: backbone + projected source bank + ephemeral workspace + sparse editor + controller + readout.

Contracts implemented here (see docs/architecture_spec.md):
  * explicit bypass: max_budget == 0 (or actions None) returns the backbone head without touching the workspace;
  * causality: the bank visible to position t is 1..t; nothing from t+1 onward enters S, the editor or the controller;
  * locality: only the slots of the chosen action are written (bitwise identity elsewhere);
  * the editor never modifies H, K, V or z (they are inputs only);
  * the deployed controller scores actions from target-free features and never executes candidate edits.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

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

from .attention import RMSNorm, LayerKVCache
from .backbone import Backbone
from .config import PrismeConfig
from .controller import ValueController
from .editor import SparseEditor
from .readout import DifferentialReadout
from .workspace import action_table, action_group_size, init_workspace


@dataclass
class EncodedContext:
    """Everything the workspace needs for a batch of query positions. All tensors are read-only for the editor."""
    H: torch.Tensor          # [B,Tq,d]   backbone residual stream at the query positions
    K: torch.Tensor          # [B,Tb,w]   source bank keys (causally readable)
    V: torch.Tensor          # [B,Tb,w]
    z: torch.Tensor          # [B,Tq,w]
    bank_mask: torch.Tensor  # [B,Tq,Tb] bool
    S0: torch.Tensor         # [B,Tq,M,w]
    ctrl_in: Optional[torch.Tensor] = None   # controller input ([B,Tq,w] z or [B,Tq,d] normalised H)
    entropy: Optional[torch.Tensor] = None   # [B,Tq] backbone predictive entropy at budget 0 (if cfg.ctrl_use_entropy)


class PrismeModel(nn.Module):
    def __init__(self, cfg: PrismeConfig):
        super().__init__()
        cfg.validate()
        self.cfg = cfg
        self.backbone = Backbone(cfg)
        if cfg.use_workspace:
            self.bank_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
            self.W_K = nn.Linear(cfg.d_model, cfg.w, bias=False)
            self.W_V = nn.Linear(cfg.d_model, cfg.w, bias=False)
            self.P = nn.Linear(cfg.d_model, cfg.w, bias=False)
            self.slot_embed = nn.Parameter(torch.randn(cfg.n_slots, cfg.w) * cfg.init_std)
            self.editor = SparseEditor(cfg)
            self.readout = DifferentialReadout(cfg)
            self.controller = ValueController(cfg)
            for lin in (self.W_K, self.W_V, self.P):
                nn.init.normal_(lin.weight, std=cfg.init_std)
            self.register_buffer("action_slots", action_table(cfg), persistent=False)
            self.register_buffer("action_k", action_group_size(cfg), persistent=False)

    # ------------------------------------------------------------------ parameter groups
    def backbone_parameters(self):
        return list(self.backbone.parameters())

    def workspace_parameters(self):
        names = ("bank_norm", "W_K", "W_V", "P", "slot_embed", "editor", "readout")
        return [p for n, p in self.named_parameters() if n.split(".")[0] in names]

    def controller_parameters(self):
        return list(self.controller.parameters())

    # ------------------------------------------------------------------ encoding
    def encode(self, tokens: torch.Tensor) -> EncodedContext:
        """Teacher-forced parallel encoding of all positions; bank mask is strictly causal (key j <= query t)."""
        H = self.backbone(tokens)
        return self.context_from_H(H, H)

    def context_from_H(self, H_bank: torch.Tensor, H_query: torch.Tensor, query_pos: Optional[torch.Tensor] = None,
                       bank_len: Optional[torch.Tensor] = None) -> EncodedContext:
        """H_bank: [B,Tb,d] history (incl. the query positions); H_query: [B,Tq,d].
        query_pos: [B,Tq] absolute positions of the queries inside the bank (default: arange).
        bank_len: [B] valid bank length per request (for padded batches)."""
        B, Tb, _ = H_bank.shape
        Tq = H_query.shape[1]
        hb = self.bank_norm(H_bank)
        K, V = self.W_K(hb), self.W_V(hb)
        z = self.P(self.bank_norm(H_query))
        if query_pos is None:
            query_pos = torch.arange(Tq, device=H_bank.device).view(1, Tq).expand(B, Tq)
        key_pos = torch.arange(Tb, device=H_bank.device)
        bank_mask = key_pos.view(1, 1, Tb) <= query_pos.unsqueeze(-1)
        if bank_len is not None:
            bank_mask = bank_mask & (key_pos.view(1, 1, Tb) < bank_len.view(B, 1, 1))
        S0 = init_workspace(self.slot_embed, z)
        hq = self.bank_norm(H_query)
        ctrl_in = hq if self.cfg.ctrl_input == "h" else z
        entropy = None
        if self.cfg.ctrl_use_entropy:
            p = torch.softmax(self.backbone.head(H_query).float(), -1)
            entropy = -(p * torch.log(p.clamp_min(1e-12))).sum(-1)
        return EncodedContext(H=H_query, K=K, V=V, z=z, bank_mask=bank_mask, S0=S0, ctrl_in=ctrl_in, entropy=entropy)

    # ------------------------------------------------------------------ one action
    def apply_action(self, S: torch.Tensor, ctx: EncodedContext, action: torch.Tensor, step, remaining):
        B, Tq = action.shape
        if not torch.is_tensor(step):
            step = torch.full((B, Tq), int(step), device=S.device, dtype=torch.long)
        if not torch.is_tensor(remaining):
            remaining = torch.full((B, Tq), int(remaining), device=S.device, dtype=torch.long)
        return self.editor(S, ctx.K, ctx.V, ctx.bank_mask, action, step, remaining)

    # ------------------------------------------------------------------ readout
    def logits_from_state(self, S: torch.Tensor, ctx: EncodedContext, edited: Optional[torch.Tensor] = None):
        """h_corrected = H + U(r(S) - r(S0)). If `edited` ([B,Tq] bool) is given, positions that were never
        edited use the backbone head directly (exact bypass)."""
        delta = self.readout(S, ctx.S0, ctx.z)
        if edited is not None:
            delta = torch.where(edited.unsqueeze(-1), delta, torch.zeros_like(delta))
        return self.backbone.head(ctx.H + delta), delta

    # ------------------------------------------------------------------ forced-action forward (training / tests)
    def forward(self, tokens: torch.Tensor, actions: Optional[torch.Tensor] = None, return_state: bool = False):
        """tokens: [B,T]; actions: [B,T,R] long (STOP for unused steps) or None (-> exact bypass).
        Returns dict(logits=[B,T,V], ...)."""
        cfg = self.cfg
        if actions is None or not cfg.use_workspace or actions.shape[-1] == 0:
            H = self.backbone(tokens)
            return {"logits": self.backbone.head(H), "H": H, "n_edits": None}
        R = actions.shape[-1]
        assert R <= cfg.r_max, f"hard action cap exceeded: {R} > r_max={cfg.r_max}"
        ctx = self.encode(tokens)
        S = ctx.S0
        n_edits = torch.zeros(actions.shape[:2], dtype=torch.long, device=tokens.device)
        stopped = torch.zeros(actions.shape[:2], dtype=torch.bool, device=tokens.device)
        infos = []
        for r in range(R):
            a = actions[:, :, r]
            a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a)  # STOP is absorbing
            nonstop = a != cfg.stop_action
            stopped = stopped | ~nonstop
            remaining = (R - r) - 1
            S, info = self.apply_action(S, ctx, a, r, remaining)
            n_edits = n_edits + nonstop.long()
            infos.append(info)
        edited = n_edits > 0
        logits, delta = self.logits_from_state(S, ctx, edited)
        out = {"logits": logits, "H": ctx.H, "n_edits": n_edits, "delta_out": delta}
        if return_state:
            out.update({"ctx": ctx, "S_final": S, "infos": infos})
        return out

    # ------------------------------------------------------------------ controller-driven forward (teacher forced)
    @torch.no_grad()
    def forward_policy(self, tokens: torch.Tensor, max_budget: int, cost: torch.Tensor,
                       lambda_cost: float | None = None, kappa_risk: float | None = None, return_state: bool = False):
        """Parallel teacher-forced rollout of the deployed controller at every position (evaluation only).
        cost: [A] analytic cost proxy per nonterminal action. Returns logits and the action trace [B,T,max_budget]."""
        cfg = self.cfg
        if max_budget == 0:
            H = self.backbone(tokens)
            return {"logits": self.backbone.head(H), "actions": None, "n_edits": torch.zeros(tokens.shape, dtype=torch.long, device=tokens.device)}
        max_budget = min(max_budget, cfg.r_max)
        ctx = self.encode(tokens)
        S = ctx.S0
        B, T = tokens.shape
        stopped = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)
        trace = torch.full((B, T, max_budget), cfg.stop_action, dtype=torch.long, device=tokens.device)
        n_edits = torch.zeros(B, T, dtype=torch.long, device=tokens.device)
        for r in range(max_budget):
            step = torch.full((B, T), r, device=tokens.device, dtype=torch.long)
            remaining = torch.full((B, T), max_budget - r - 1, device=tokens.device, dtype=torch.long)
            preds = self.controller(ctx.ctrl_in, S, step, remaining, ctx.entropy)
            a = self.controller.select(preds, cost, lambda_cost=lambda_cost, kappa_risk=kappa_risk)
            a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a)
            stopped = stopped | (a == cfg.stop_action)
            trace[:, :, r] = a
            S, _ = self.apply_action(S, ctx, a, step, remaining)
            n_edits = n_edits + (a != cfg.stop_action).long()
            if bool(stopped.all()):
                break
        logits, delta = self.logits_from_state(S, ctx, n_edits > 0)
        out = {"logits": logits, "actions": trace, "n_edits": n_edits}
        if return_state:
            out.update({"ctx": ctx, "S_final": S})
        return out