#!/usr/bin/env python3
"""REAL TRANSFORMER INFERENCE BENCHMARK"""
try:
    import torch
    import torch.nn as nn
except ImportError:
    print("\033[91m[ERROR] torch not installed. Run: pip install torch\033[0m")
    exit(1)

import time
import numpy as np

class SmallTransformer(nn.Module):
    def __init__(self, d_model=256, n_heads=4, n_layers=2, max_seq=1024):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(10000, d_model)
        self.pos_embed = nn.Parameter(torch.randn(1, max_seq, d_model) * 0.02)
        
        attn = nn.MultiheadAttention(d_model, n_heads, dropout=0, batch_first=True)
        self.attention = nn.ModuleList([attn] * n_layers)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model)
        )
        self.norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, 10000)
        
    def forward(self, x):
        seq_len = x.size(1)
        h = self.embedding(x) + self.pos_embed[:, :seq_len, :]
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool().to(x.device)
        
        for attn_layer in self.attention:
            attn, _ = attn_layer(h, h, h, attn_mask=mask)
            h = h + attn
        h = h + self.ffn(h)
        h = self.norm(h)
        return self.lm_head(h)

def benchmark():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"\033[36mDevice:\033[0m {device}")
    if device == 'cuda':
        print(f"\033[36mGPU:\033[0m {torch.cuda.get_device_name(0)}")
    
    config = {'d_model': 256, 'n_heads': 4, 'n_layers': 2, 'max_seq': 1024}
    model = SmallTransformer(**config).to(device).eval()
    
    input_ids = torch.randint(0, 10000, (1, 256), device=device)
    
    times = []
    print("\033[33mRunning 5 inference cycles...\033[0m")
    for i in range(5):
        if device == 'cuda': torch.cuda.synchronize()
        start = time.perf_counter()
        with torch.inference_mode():
            _ = model(input_ids)
        if device == 'cuda': torch.cuda.synchronize()
        times.append((time.perf_counter() - start) * 1000)
        print(f"  Cycle {i+1}: {times[-1]:.2f} ms")
    
    print(f"\n\033[1m=== TRANSFORMER BASELINE RESULTS ===\033[0m")
    print(f"TTFT (256 tokens): {np.mean(times):.2f} ± {np.std(times):.2f} ms")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
    print(f"Model: {config['n_layers']}-layer, {config['d_model']}d")

if __name__ == '__main__':
    benchmark()