"""Analytic operation counts (2 x multiply-adds = "ops") per generated token at context length T.

This is an ARITHMETIC model of the reference implementation, not a hardware timing model. Conventions:
  * one multiply-add = 2 ops; softmax/normalisation/activation costs are ignored (they are O(dim), not O(dim^2));
  * decode-step accounting: the backbone attends over the whole context T (KV cache), the LM head is included;
  * source-bank projections are incremental (only the new position is projected);
  * the editor re-reads the full bank (T keys) for every active query of every block, and refreshes the slot K/V
    of all M slots in every block (reference implementation; an optimised kernel could cache unchanged slots);
  * controller: one shared state MLP per step + one w x w projection per candidate action (A = M + M/2);
  * readout: two attention reads (final and initial state) + the U projection.
"""
from __future__ import annotations

from dataclasses import dataclass

from .config import PrismeConfig


@dataclass
class OpsBreakdown:
    backbone_linear: float
    backbone_attention: float
    lm_head: float
    bank_projection: float
    editor_source_attn: float
    editor_slot_attn: float
    editor_ffn: float
    editor_gate: float
    controller: float
    readout: float

    @property
    def backbone(self) -> float:
        return self.backbone_linear + self.backbone_attention + self.lm_head

    @property
    def workspace(self) -> float:
        return (self.bank_projection + self.editor_source_attn + self.editor_slot_attn + self.editor_ffn
                + self.editor_gate + self.controller + self.readout)

    @property
    def total(self) -> float:
        return self.backbone + self.workspace

    def as_dict(self) -> dict:
        d = self.__dict__.copy()
        d.update(backbone=self.backbone, workspace=self.workspace, total=self.total)
        return d


def backbone_nonembedding_params(c: PrismeConfig) -> int:
    d, f, L, kv = c.d_model, c.ffn_width, c.n_layers, c.kv_dim
    return L * (2 * d * d + 2 * d * kv + 3 * d * f + 2 * d) + d


def editor_block_ops(c: PrismeConfig, k: int, T: int) -> dict[str, float]:
    w, fw, M = c.w, c.ffn_w, c.n_slots
    return {
        "source_attn": 4 * k * w * w + 4 * k * T * w,          # q/o projections + QK^T + AV over T bank entries
        "slot_attn": 4 * k * w * w + 4 * M * w * w + 4 * k * M * w,  # q/o + K,V refresh of M slots + scores/AV
        "ffn": 6 * k * w * fw,
    }


def controller_step_ops(c: PrismeConfig) -> float:
    w, hidden, A = c.w, (c.controller_hidden or c.ffn_w), c.n_actions
    din = c.d_model if c.ctrl_input == "h" else w
    shared = 2 * din * w + 2 * 2 * w + 6 * w * hidden + (2 * 2 * w if c.ctrl_use_entropy else 0)
    per_action = 2 * w * w + 2 * 4 * w
    return shared + A * per_action


def controller_token_ops(c: PrismeConfig) -> float:
    """Once-per-token controller features: the backbone entropy needs one head evaluation (budget 0 would need it
    anyway only if the token stops immediately; we charge it fully, conservatively)."""
    return 2 * c.d_model * c.vocab_size if c.ctrl_use_entropy else 0.0


def readout_ops(c: PrismeConfig) -> float:
    w, d, M = c.w, c.d_model, c.n_slots
    return 2 * w * w + 2 * (2 * M * w * w + 4 * M * w) + 2 * w * d


def ops_per_token(c: PrismeConfig, T: int, n_actions: int, k: int, use_controller: bool = True,
                  controller_steps: int | None = None) -> OpsBreakdown:
    """Ops for one decode step with `n_actions` executed edits of group size k (k = M for dense writes).
    controller_steps defaults to n_actions + 1 (one evaluation per executed action plus the one that STOPs)."""
    d, V = c.d_model, c.vocab_size
    bb_lin = 2 * backbone_nonembedding_params(c)
    bb_att = 4 * c.n_layers * T * d
    head = 2 * d * V
    if n_actions == 0:
        return OpsBreakdown(bb_lin, bb_att, head, 0, 0, 0, 0, 0, 0, 0)
    bank = 2 * 3 * d * c.w
    blk = editor_block_ops(c, k, T)
    nb = c.n_editor_blocks * n_actions
    gate = n_actions * 2 * k * 2 * c.w
    steps = (n_actions + 1 if controller_steps is None else controller_steps) if use_controller else 0
    steps = min(steps, c.r_max)
    ctrl = steps * controller_step_ops(c) + controller_token_ops(c)
    return OpsBreakdown(bb_lin, bb_att, head, bank, nb * blk["source_attn"], nb * blk["slot_attn"], nb * blk["ffn"],
                        gate, ctrl, readout_ops(c))


def action_cost_proxy(c: PrismeConfig, T: int) -> list[float]:
    """Known analytic cost of each nonterminal action (ops of one edit, incl. the next controller evaluation),
    normalised by the backbone-only ops/token at context T. Used as c(a) by the controller."""
    base = ops_per_token(c, T, 0, 1).backbone
    costs = []
    for k in [1] * c.n_slots + [2] * c.n_pairs:
        blk = editor_block_ops(c, k, T)
        e = c.n_editor_blocks * sum(blk.values()) + 2 * k * 2 * c.w + controller_step_ops(c)
        costs.append(e / base)
    return costs


def global_recurrence_step_ops(c: PrismeConfig, T: int, ffn_w: int | None = None) -> float:
    """One refinement step of the global-recurrence adapter (baselines/global_recurrence.py) for one position at
    context T: q/k/v/o projections (8w^2), causal attention over T positions (4Tw), SwiGLU (6 w f), plus the
    down/up projections charged once per token (4dw)."""
    w, f = c.w, (ffn_w or c.ffn_w)
    return 8 * w * w + 4 * T * w + 6 * w * f


def global_recurrence_token_ops(c: PrismeConfig) -> float:
    return 2 * 2 * c.d_model * c.w