#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

from transformers import AutoTokenizer


PRESETS = {
    "large": {"d": 1024, "layers": 24, "heads": 16, "rank": 128},
    "agillm4_floor": {"d": 1280, "layers": 28, "heads": 20, "rank": 160},
    "agillm4_main": {"d": 1536, "layers": 32, "heads": 24, "rank": 192},
    "agillm4_big": {"d": 1792, "layers": 36, "heads": 28, "rank": 224},
}


def estimate(vocab: int, d: int, layers: int, heads: int, rank: int, tie_weights: bool = False) -> dict:
    dk = d // heads
    embed = vocab * d
    attn = 3 * d * d + d * d + dk * rank
    ff = d * (4 * d) + (4 * d) + (4 * d) * d + d
    norms = 4 * d
    block = attn + ff + norms
    core = embed + layers * block + 2 * d
    ar = 0 if tie_weights else vocab * d + vocab
    nat = vocab * d + vocab
    sat = vocab * d + vocab + 2 * d + 2
    total = core + ar + sat + nat
    return {
        "vocab": vocab,
        "d_model": d,
        "layers": layers,
        "heads": heads,
        "rank": rank,
        "tie_weights": tie_weights,
        "embedding_params": embed,
        "block_params_each": block,
        "core_params": core,
        "ar_head_params": ar,
        "nat_head_params": nat,
        "sat_head_params": sat,
        "total_params": total,
        "tokens_at_100_to_1": total * 100,
    }


def fmt(n: int) -> str:
    if n >= 1_000_000_000:
        return f"{n / 1_000_000_000:.3f}B"
    if n >= 1_000_000:
        return f"{n / 1_000_000:.1f}M"
    return f"{n:,}"


def main() -> int:
    parser = argparse.ArgumentParser(description="Estimate AGILLM-4 parameter and token targets")
    parser.add_argument("--tokenizer", default="deepseek-ai/DeepSeek-V4-Pro")
    parser.add_argument("--preset", default="agillm4_main", choices=PRESETS)
    parser.add_argument("--tie_weights", action="store_true")
    parser.add_argument("--json_out", default="")
    args = parser.parse_args()

    tok = AutoTokenizer.from_pretrained(args.tokenizer, use_fast=True, trust_remote_code=True)
    vocab = max(tok.get_vocab().values()) + 1
    result = estimate(vocab=vocab, tie_weights=args.tie_weights, **PRESETS[args.preset])
    print(f"preset={args.preset}")
    print(f"total_params={fmt(result['total_params'])} ({result['total_params']:,})")
    print(f"target_tokens_100_to_1={fmt(result['tokens_at_100_to_1'])} ({result['tokens_at_100_to_1']:,})")
    print(json.dumps(result, indent=2, sort_keys=True))
    if args.json_out:
        Path(args.json_out).write_text(json.dumps(result, indent=2, sort_keys=True), encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())