"""Gate: the instantiated model must have exactly 20,453,760 parameters."""
import math
import sys

import torch
import torch.nn.functional as F

from common import EXPECTED_PARAMS, build_model


def main() -> int:
    torch.manual_seed(0)
    model = build_model()
    cfg = model.config
    m = model.model

    total = sum(p.numel() for p in model.parameters())
    emb = m.embed_tokens.weight.numel()
    per_layer = [sum(p.numel() for p in layer.parameters()) for layer in m.layers]
    final_norm = m.norm.weight.numel()
    tied = model.lm_head.weight.data_ptr() == m.embed_tokens.weight.data_ptr()

    print(f"params total        : {total:,}")
    print(f"  embedding (tied)  : {emb:,}")
    print(f"  per layer         : {per_layer[0]:,} x {len(per_layer)} = {sum(per_layer):,}")
    print(f"  final norm        : {final_norm:,}")
    print(f"lm_head tied        : {tied}")
    print(f"attn_implementation : {cfg._attn_implementation}")
    print(f"hidden_act          : {cfg.hidden_act}")
    print(f"rope                : {getattr(cfg, 'rope_parameters', None) or getattr(cfg, 'rope_theta', None)}")
    print(f"param dtype         : {next(model.parameters()).dtype}")

    layer0 = m.layers[0]
    print("init std            : "
          f"embed {m.embed_tokens.weight.std():.4f} | q_proj {layer0.self_attn.q_proj.weight.std():.4f} | "
          f"o_proj {layer0.self_attn.o_proj.weight.std():.4f} | gate_proj {layer0.mlp.gate_proj.weight.std():.4f} | "
          f"down_proj {layer0.mlp.down_proj.weight.std():.4f} (target {0.02 / math.sqrt(2 * cfg.num_hidden_layers):.4f}) | "
          f"norm mean {layer0.input_layernorm.weight.mean():.2f}")

    if torch.cuda.is_available():
        model.cuda()
        x = torch.randint(0, cfg.vocab_size, (4, 1025), device="cuda")
        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
            logits = model(input_ids=x[:, :-1]).logits
        loss = F.cross_entropy(logits.float().reshape(-1, cfg.vocab_size), x[:, 1:].reshape(-1))
        print(f"init loss (random)  : {loss.item():.4f}  (ln V = {math.log(cfg.vocab_size):.4f})")

    ok = total == EXPECTED_PARAMS
    print(f"GATE {'PASS' if ok else 'FAIL'}: {total:,} (esperado {EXPECTED_PARAMS:,})")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())