#!/usr/bin/env python3
"""Parameter planning calculator for the bilateral Gated-DeltaNet-2 architecture.

Counts the design envelope, not an implementation-exact GDN2 checkpoint. The
official recurrent block must replace the provisional projection count once M3a
exists; this script prevents the old Transformer/CfC count from masquerading as
the new architecture.

Run: python3 params.py [--d 2048 --layers 18 --ffn 5120 --vocab 32000]
"""

import argparse
from dataclasses import dataclass


@dataclass
class Config:
    vocab: int = 32_000
    d_model: int = 2_048
    n_layers: int = 18          # total recurrent blocks across both hemispheres
    d_ffn: int = 5_120
    mixer_matrices: int = 5     # provisional GDN2 projection/gate budget per block
    workspace_matrices: int = 3 # shared bilateral read/write/broadcast transforms
    mtp_horizons: int = 4       # shared unembedding + one dxd transform per extra horizon
    n_regions: int = 16
    tie_embeddings: bool = True
    plastic_tags: bool = False  # M5-gated; absent from M3/M4


def count(cfg: Config) -> dict:
    d, f = cfg.d_model, cfg.d_ffn

    embed = cfg.vocab * d
    unembed = 0 if cfg.tie_embeddings else cfg.vocab * d

    # Per recurrent block: provisional GDN2 projection/gate budget + SwiGLU + norms.
    # Regions partition this total capacity; they do not duplicate it per hemisphere.
    mixer = cfg.mixer_matrices * d * d
    ffn = 3 * d * f
    norms = 2 * d
    core_total = cfg.n_layers * (mixer + ffn + norms)
    workspace = cfg.workspace_matrices * d * d
    router = 2 * d * cfg.n_regions
    mtp = max(cfg.mtp_horizons - 1, 0) * d * d
    core = core_total + workspace + router + mtp
    total = embed + unembed + core

    return {
        "embedding": embed,
        "unembedding": unembed,
        "GDN2 blocks (x%d)" % cfg.n_layers: core_total,
        "  - recurrent mixer/block (est.)": mixer,
        "  - swiglu ffn/layer": ffn,
        "shared workspace (est.)": workspace,
        "region router": router,
        "MTP horizon transforms": mtp,
        "TOTAL": total,
        "non-embedding": core,
    }


def fmt(n: float) -> str:
    if n >= 1e9:
        return "%.3f B" % (n / 1e9)
    if n >= 1e6:
        return "%.2f M" % (n / 1e6)
    return "{:,}".format(int(n))


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--d", type=int, default=2048)
    p.add_argument("--layers", type=int, default=18)
    p.add_argument("--ffn", type=int, default=5120)
    p.add_argument("--vocab", type=int, default=32000)
    p.add_argument("--mixer-mats", type=int, default=5)
    p.add_argument("--workspace-mats", type=int, default=3)
    p.add_argument("--mtp-horizons", type=int, default=4)
    p.add_argument("--regions", type=int, default=16)
    p.add_argument("--untie", action="store_true", help="use a separate unembedding")
    a = p.parse_args()

    cfg = Config(vocab=a.vocab, d_model=a.d, n_layers=a.layers, d_ffn=a.ffn,
                 mixer_matrices=a.mixer_mats, workspace_matrices=a.workspace_mats,
                 mtp_horizons=a.mtp_horizons, n_regions=a.regions,
                 tie_embeddings=not a.untie)
    c = count(cfg)

    print("Mind-1B bilateral GDN2 planning envelope")
    print("  d_model=%d layers=%d ffn=%d vocab=%d tie=%s mixer_mats=%d"
          % (cfg.d_model, cfg.n_layers, cfg.d_ffn, cfg.vocab,
             cfg.tie_embeddings, cfg.mixer_matrices))
    print("-" * 46)
    for k, v in c.items():
        pad = 30 - len(k)
        print("  %s%s%12s" % (k, " " * max(pad, 1), fmt(v)))
    print("-" * 46)

    total = c["TOTAL"]
    if cfg.plastic_tags:
        print("  plastic tags (int8, not counted as params): +%s bytes"
              % fmt(c["non-embedding"]))
    print("  ternary core packed at 2 bits/weight: %.0f MB"
          % (c["non-embedding"] * 2 / 8 / 1e6))
    print("  fp16 shadow weights during training:  %.2f GB" % (total * 2 / 1e9))


if __name__ == "__main__":
    main()