"""Parameter counts and FLOP estimation for the matched-budget comparisons.

Convention (documented so the numbers are reproducible):
  - forward FLOPs per token per block application:
      matmuls  2*(4*d^2 + 2*d*d_ff)   (qkv+out proj: 4d^2; mlp: 2*d*d_ff)
      attention 4*T*d                  (qk^T and av, T = S+2, MACs x2)
    plus head 2*S*d once per forward.
  - training FLOPs = 3 x forward FLOPs (backward ~2x forward).
  - a recurrent model with `depth` shared blocks unrolled n_iters times pays
    n_iters*depth block applications per step; its parameter count has only
    `depth` block copies. Estimates ignore embeddings/layernorms/CE (<1%).
"""

from ptr.models import PtrTransformer


def count_params(model):
    total = sum(p.numel() for p in model.parameters())
    blocks = sum(p.numel() for b in model.blocks for p in b.parameters())
    return {"total": total, "blocks": blocks}


def forward_flops_per_token(model, S, n_iters):
    d, d_ff = model.d_model, model.blocks[0].mlp_in.out_features
    T = S + 2
    per_block = 2 * (4 * d * d + 2 * d * d_ff) + 4 * T * d
    head = 2 * S * d
    return n_iters * model.depth * per_block + head


def step_flops(model, S, batch, n_iters, training=True):
    fwd = batch * (S + 2) * forward_flops_per_token(model, S, n_iters)
    return 3 * fwd if training else fwd


def flops_budget_summary(model, S, batch, n_iters, steps):
    """Total training FLOPs for one run, plus params. `n_iters` may be an int
    or a (mode, lo, hi) tuple for per-batch sampled unrolls (mean used)."""
    p = count_params(model)
    if isinstance(n_iters, tuple):
        lo, hi = n_iters[1], n_iters[2]
        mean_iters = (lo + hi) / 2.0
        if len(n_iters) > 3:
            mean_iters += n_iters[3]
            desc = f"{n_iters[0]} U{{{lo}..{hi}}}+{n_iters[3]}"
        else:
            desc = f"{n_iters[0]} U{{{lo}..{hi}}}"
    else:
        mean_iters = n_iters
        desc = str(n_iters)
    total = step_flops(model, S, batch, mean_iters, training=True) * steps
    return {
        "params_total": p["total"],
        "params_blocks": p["blocks"] + sum(q.numel() for blk in model.stem
                                           for q in blk.parameters())
        if hasattr(model, "stem") else p["blocks"],
        "iters_per_step": desc,
        "block_apps_per_step": model.depth * mean_iters
        + getattr(model, "stem_depth", 0),
        "train_flops_total": float(total),
    }