#!/usr/bin/env python3
"""Planning VRAM budget for bilateral Gated-DeltaNet-2 Mind-1B.

Training assumes BitNet-style quantization-aware training: bf16 master
weights, bf16 grads, fp32 Adam moments (m, v), optional fp32 master copy.
The GDN2 state allocation must be measured from the implementation and supplied
with --state; this script no longer pretends a Transformer KV cache is the state.

Run: python3 memory_budget.py [--params 1034559488 --core 969023488 --batch 8 --seq 2048]
"""

import argparse

GB = 1e9


def training_budget(n_params, n_core, batch, seq, d_model, n_layers,
                    checkpointing=True, fp32_master=False, plastic_tags=False):
    weights = n_params * 2                      # bf16
    grads = n_params * 2                        # bf16
    adam = n_params * 8                         # fp32 m + v
    master = n_params * 4 if fp32_master else 0
    plastic = n_core if plastic_tags else 0      # M5-only int8 Hebbian tag per core weight

    # Activation memory per token per layer, bf16. With full activation
    # checkpointing only layer inputs are stored (~2 tensors of d_model);
    # without it, attention + swiglu intermediates dominate (~14x d_model).
    per_tok_layer = (2 if checkpointing else 14) * d_model * 2
    activations = batch * seq * n_layers * per_tok_layer

    total = weights + grads + adam + master + plastic + activations
    return {
        "master/bf16 weights": weights + master,
        "gradients (bf16)": grads,
        "adam moments (fp32)": adam,
        "plastic tags (int8, M5 only)": plastic,
        "activations (b=%d, s=%d%s)" % (batch, seq, ", ckpt" if checkpointing else ""): activations,
        "TOTAL": total,
    }


def inference_budget(n_core, n_embed, d_model, d_vsa, cache_tokens,
                     recurrent_state, plastic_tags=False):
    core = n_core * 2 / 8                       # 2-bit packed ternary
    embeds = n_embed * 2                        # fp16 embed + unembed
    plastic = n_core if plastic_tags else 0      # M5-only tags, live during wake
    vsa = d_vsa * 4                             # holographic trace, int32 accumulators
    # Hippocampal ring buffer: token ids (4B) + one d_model int8 state per token
    hippo = cache_tokens * (4 + d_model)
    total = core + embeds + plastic + vsa + hippo + recurrent_state
    return {
        "ternary core (2-bit packed)": core,
        "embeddings (fp16)": embeds,
        "plastic tags (int8, M5 only)": plastic,
        "VSA trace (int32 x %d)" % d_vsa: vsa,
        "hippocampal cache (%dk tok)" % (cache_tokens // 1024): hippo,
        "GDN2 recurrent state (measured)": recurrent_state,
        "TOTAL": total,
    }


def show(title, d):
    print(title)
    print("-" * 52)
    for k, v in d.items():
        print("  %-38s %8.2f GB" % (k, v / GB))
    print()


def main():
    p = argparse.ArgumentParser()
    # Keep these aligned with calc/params.py's default Config. The split is
    # material: only the non-embedding core is 2-bit packed at inference.
    p.add_argument("--params", type=float, default=1_034_559_488)
    p.add_argument("--core", type=float, default=969_023_488, help="planning non-embedding params")
    p.add_argument("--batch", type=int, default=8, help="per-GPU micro-batch")
    p.add_argument("--seq", type=int, default=2048)
    p.add_argument("--d", type=int, default=2048)
    p.add_argument("--layers", type=int, default=18)
    p.add_argument("--vsa", type=int, default=0, help="M5-gated VSA dimension")
    p.add_argument("--cache", type=int, default=0, help="M5-gated hippocampal tokens")
    p.add_argument("--state", type=float, default=0,
                   help="measured GDN2 recurrent-state bytes per session")
    p.add_argument("--plastic", action="store_true", help="include M5 int8 plastic tags")
    a = p.parse_args()

    show("TRAINING per GPU (QAT, activation checkpointing on)",
         training_budget(a.params, a.core, a.batch, a.seq, a.d, a.layers,
                         plastic_tags=a.plastic))
    show("TRAINING per GPU (no checkpointing)",
         training_budget(a.params, a.core, a.batch, a.seq, a.d, a.layers,
                         checkpointing=False, plastic_tags=a.plastic))
    show("INFERENCE (single session, excluding unspecified system overhead)",
         inference_budget(a.core, a.params - a.core, a.d, a.vsa, a.cache,
                          a.state, plastic_tags=a.plastic))
    if not a.state:
        print("WARNING: --state=0; measure and supply the GDN2 recurrent state before using this as a deployment budget.")
    print("takeaway: this is a planning envelope, not an implementation-verified deployment claim.")


if __name__ == "__main__":
    main()