"""Exact parameter accounting.

The headline number is trainable weights PLUS frozen buffer values.
"""
from __future__ import annotations

from collections import OrderedDict

import torch
import torch.nn as nn

# Buffers that hold only derived constants would be excluded and listed here.
# The design deliberately has none; RoPE tables are plain Python state.
DERIVED_BUFFERS: tuple[str, ...] = ()


def _tied_targets(model: nn.Module) -> set[str]:
    """Parameter names that are copies of another parameter, so count once.

    Object identity is not enough: `from_pretrained` constructs the model before
    tying the weights, so during construction `lm_head.weight` is still a
    distinct tensor. Counting it would double the embedding and trip the budget
    check on an otherwise valid model.
    """
    if not getattr(model.config, "tie_word_embeddings", False):
        return set()
    keys = getattr(model, "_tied_weights_keys", None) or []
    # transformers 4.x: a list of target names. 5.x: {target: source}.
    return set(keys.keys()) if isinstance(keys, dict) else set(keys)


def account(model: nn.Module) -> dict:
    """Walk named_parameters() and named_buffers(). Returns a count breakdown."""
    trainable, frozen, excluded = 0, 0, 0
    by_param: "OrderedDict[str, int]" = OrderedDict()
    by_buffer: "OrderedDict[str, int]" = OrderedDict()
    excluded_names: list[str] = []

    tied = _tied_targets(model)
    seen: set[int] = set()
    for name, p in model.named_parameters():
        if id(p) in seen:          # tied weights counted once
            continue
        if name in tied or any(name.endswith("." + t) for t in tied):
            continue
        seen.add(id(p))
        n = p.numel()
        by_param[name] = n
        trainable += n

    for name, b in model.named_buffers():
        if b is None or id(b) in seen:
            continue
        seen.add(id(b))
        n = b.numel()
        if any(name.endswith(suffix) for suffix in DERIVED_BUFFERS):
            excluded += n
            excluded_names.append(name)
            continue
        by_buffer[name] = n
        frozen += n

    return {
        "trainable": trainable,
        "frozen": frozen,
        "total": trainable + frozen,
        "excluded": excluded,
        "excluded_names": excluded_names,
        "by_param": by_param,
        "by_buffer": by_buffer,
    }


def _group(model: nn.Module) -> "OrderedDict[str, int]":
    """Spec-shaped component table, computed from the live model."""
    cfg = model.config
    base = model.model if hasattr(model, "model") else model
    rows: "OrderedDict[str, int]" = OrderedDict()

    def count(mods) -> int:
        seen, total = set(), 0
        for m in mods:
            if m is None:
                continue
            for p in m.parameters(recurse=True):
                if id(p) not in seen:
                    seen.add(id(p))
                    total += p.numel()
        return total

    layer0 = base.layers[0]
    n_layer = len(base.layers)

    attn = count([layer0.attn.q_proj, layer0.attn.k_proj, layer0.attn.v_proj, layer0.attn.o_proj])
    qk = count([layer0.attn.q_norm, layer0.attn.k_norm])
    gate = count([layer0.attn.gate_proj])
    mlp = count([layer0.mlp])
    norms = count([layer0.attn_in_norm, layer0.attn_out_norm, layer0.mlp_in_norm, layer0.mlp_out_norm])
    canon = count([layer0.canon_a, layer0.canon_c])
    per_layer = attn + qk + gate + mlp + norms + canon

    rows["Attention q, k, v, o"] = attn
    rows["QK-norm"] = qk
    rows[f"{cfg.attn_gate.capitalize()} gate (W + b)"] = gate
    rows[f"SwiGLU MLP (3 x {cfg.d_model} x {cfg.ffn_hidden})"] = mlp
    rows[f"Peri-LN, 4 norms x {cfg.d_model}"] = norms
    rows["Canon A + C"] = canon
    rows["--- Per layer"] = per_layer
    rows[f"{n_layer} layers"] = per_layer * n_layer

    emb = base.embed_tokens.weight.numel()
    rows[f"Tied embedding ({cfg.vocab_size} x {cfg.d_model})"] = emb
    rows["Final norm"] = base.final_norm.weight.numel()

    info = account(model)
    rows["--- Trainable"] = info["trainable"]
    if info["frozen"]:
        rows["Frozen buffers"] = info["frozen"]
    rows["--- Total"] = info["total"]
    return rows


def table(model: nn.Module) -> str:
    rows = _group(model)
    info = account(model)
    width = max(len(k) for k in rows) + 2
    lines = [f"{'Component'.ljust(width)}{'Count':>12}", "-" * (width + 12)]
    for key, val in rows.items():
        if key.startswith("--- "):
            lines.append("-" * (width + 12))
            key = key[4:]
        lines.append(f"{key.ljust(width)}{val:>12,}")
    lines.append("-" * (width + 12))

    budget = getattr(model.config, "max_total_params", None)
    warn = getattr(model.config, "warn_total_params", None)
    if budget:
        ok = "OK" if info["total"] <= budget else "OVER BUDGET"
        lines.append(f"{'Budget (max_total_params)'.ljust(width)}{budget:>12,}  [{ok}]")
    if warn:
        flag = "ok" if info["total"] <= warn else "WARN"
        lines.append(f"{'Warn threshold'.ljust(width)}{warn:>12,}  [{flag}]")
    if info["excluded"]:
        lines.append(
            f"Excluded derived buffers ({info['excluded']:,}): "
            + ", ".join(info["excluded_names"])
        )
    return "\n".join(lines)


def count_from_config(cfg_or_model) -> dict:
    """Build the model if given a config dict, then account for it.

    Construction itself raises when the budget is exceeded, so this doubles as
    the budget check.
    """
    if isinstance(cfg_or_model, nn.Module):
        return account(cfg_or_model)
    from ethoric.model.configuration_ethoric import EthoricConfig
    from ethoric.model.modeling_ethoric import EthoricForCausalLM

    model_cfg = EthoricConfig.from_pipeline(cfg_or_model)
    with torch.device("meta"):
        model = EthoricForCausalLM(model_cfg)
    return account(model)