"""Smoke test: load ArchonBrain + state_dict + forward pass V100."""
import sys, time, torch
sys.path.insert(0, 'source')
from config import ArchonBrainConfig
import model as model_module

print('=== model.py classes ===')
print([c for c in dir(model_module) if not c.startswith('_') and c[0].isupper()])

# Identify Brain class
ArchonBrain = None
for name in ['ArchonBrain', 'Brain', 'ArchonBrainModel', 'Model']:
    if hasattr(model_module, name):
        ArchonBrain = getattr(model_module, name)
        print(f'Found: {name}')
        break

cfg = ArchonBrainConfig()
print(f'\nConfig param_count estimate: {cfg.param_count_human}')

print('\n[1] Instantiating model...')
m = ArchonBrain(cfg)
n_params = sum(p.numel() for p in m.parameters())
print(f'  Instantiated: {n_params/1e6:.1f}M params')

print('\n[2] Loading ckpt state_dict...')
ckpt = torch.load('ckpts/step_259567/archon.pt', map_location='cpu', weights_only=False)
missing, unexpected = m.load_state_dict(ckpt['model'], strict=False)
print(f'  missing: {len(missing)} (first 5: {missing[:5]})')
print(f'  unexpected: {len(unexpected)} (first 5: {unexpected[:5]})')

print('\n[3] Moving to GPU bf16...')
m = m.to(device='cuda', dtype=torch.bfloat16)
torch.cuda.synchronize()
free, total = torch.cuda.mem_get_info()
print(f'  VRAM used: {(total-free)/1e9:.2f}GB / {total/1e9:.2f}GB')

print('\n[4] Forward pass batch=4 seq=4096...')
x = torch.randint(0, 32000, (4, 4096), device='cuda')
with torch.no_grad():
    torch.cuda.synchronize()
    t0 = time.time()
    out = m(x) if not hasattr(m, 'forward_impl') else m.forward(x)
    torch.cuda.synchronize()
    dt = time.time() - t0
print(f'  output type: {type(out).__name__}')
if isinstance(out, torch.Tensor):
    print(f'  output shape: {tuple(out.shape)} dtype: {out.dtype}')
elif isinstance(out, (tuple, list)):
    for i, o in enumerate(out):
        if isinstance(o, torch.Tensor):
            print(f'  out[{i}]: {tuple(o.shape)} {o.dtype}')
elif isinstance(out, dict):
    for k, v in out.items():
        if isinstance(v, torch.Tensor):
            print(f'  out[{k}]: {tuple(v.shape)} {v.dtype}')
print(f'  forward time: {dt*1000:.1f}ms')

print('\n[5] Throughput estimate (forward only, no grad):')
tokens = 4 * 4096
print(f'  {tokens/dt:.0f} tok/s forward bf16')
print(f'  (training est ~3x slower = ~{tokens/(dt*3):.0f} tok/s with backward+optim)')

print('\nSMOKE_TEST_OK')