from __future__ import annotations
import argparse
import time
import torch
import torch.nn.functional as F

from bihmoe.models.dense import DenseModel
from bihmoe.models.structured import StructuredBiHMoE
from bihmoe.utils.misc import set_seed, count_params, fmt_bytes

def cuda_mem(label: str) -> None:
    if not torch.cuda.is_available():
        print(f"{label}: cuda not available")
        return
    alloc = torch.cuda.memory_allocated()
    reserv = torch.cuda.memory_reserved()
    peak = torch.cuda.max_memory_allocated()
    print(f"{label}: alloc={fmt_bytes(alloc)} reserved={fmt_bytes(reserv)} peak={fmt_bytes(peak)}")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seed", type=int, default=123)
    ap.add_argument("--vocab", type=int, default=2048)
    ap.add_argument("--d_model", type=int, default=384)
    ap.add_argument("--heads", type=int, default=6)
    ap.add_argument("--seq", type=int, default=128)
    ap.add_argument("--batch", type=int, default=8)

    # Dense baseline
    ap.add_argument("--dense_layers", type=int, default=6)
    ap.add_argument("--dense_dff", type=int, default=1536)
    ap.add_argument("--dense_pool", type=str, default="mean", choices=["mean","first"])

    # Structured
    ap.add_argument("--stem_layers", type=int, default=1)
    ap.add_argument("--hemi_layers", type=int, default=4)
    ap.add_argument("--expert_dff", type=int, default=1024)
    ap.add_argument("--experts", type=int, default=8)
    ap.add_argument("--topk", type=int, default=1)
    ap.add_argument("--workspace", type=int, default=4)
    ap.add_argument("--reconcile_every", type=int, default=2)

    ap.add_argument("--steps", type=int, default=1)
    ap.add_argument("--dtype", type=str, default="fp16", choices=["fp16","bf16","fp32"])
    args = ap.parse_args()

    set_seed(args.seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if torch.cuda.is_available():
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.cuda.reset_peak_memory_stats()

    def pick_dtype():
        if args.dtype == "fp16":
            return torch.float16
        if args.dtype == "bf16":
            return torch.bfloat16
        return torch.float32

    dt = pick_dtype()
    print("device:", device)
    if torch.cuda.is_available():
        print("gpu:", torch.cuda.get_device_name(0))
    print("dtype:", dt)

    dense = DenseModel(
        vocab_size=args.vocab,
        d_model=args.d_model,
        n_heads=args.heads,
        n_layers=args.dense_layers,
        d_ff=args.dense_dff,
        dropout=0.0,
        head_mode="cls",
        pool=args.dense_pool,
    ).to(device).to(dtype=dt)

    struct = StructuredBiHMoE(
        vocab_size=args.vocab,
        d_model=args.d_model,
        n_heads=args.heads,
        n_layers_stem=args.stem_layers,
        n_layers_hemi=args.hemi_layers,
        d_ff_dense=args.dense_dff,     # NOTE: for probe we reuse dense_dff here; compute_match will override later
        d_ff_expert=args.expert_dff,
        n_experts=args.experts,
        top_k=args.topk,
        workspace_tokens=args.workspace,
        reconcile_every=args.reconcile_every,
        dropout=0.0,
    ).to(device).to(dtype=dt)

    print("params_dense:", count_params(dense))
    print("params_struct:", count_params(struct))

    opt_d = torch.optim.AdamW(dense.parameters(), lr=1e-4)
    opt_s = torch.optim.AdamW(struct.parameters(), lr=1e-4)

    cuda_mem("after_init")

    inp = torch.randint(0, args.vocab, (args.batch, args.seq), device=device)
    tgt_cls = torch.randint(0, args.vocab, (args.batch,), device=device)

    for step in range(args.steps):
        t0 = time.time()
        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats()

        opt_s.zero_grad(set_to_none=True)
        logits_s = struct(inp)  # (B,V)
        loss_s = F.cross_entropy(logits_s.float(), tgt_cls)
        loss_s.backward()
        opt_s.step()
        cuda_mem(f"after_struct_step{step}")

        opt_d.zero_grad(set_to_none=True)
        logits_d = dense(inp)   # (B,V)
        loss_d = F.cross_entropy(logits_d.float(), tgt_cls)
        loss_d.backward()
        opt_d.step()
        cuda_mem(f"after_dense_step{step}")

        t1 = time.time()
        print(f"step{step}: loss_s={loss_s.item():.4f} loss_d={loss_d.item():.4f} dt={t1-t0:.3f}s")

if __name__ == "__main__":
    main()