"""Compare Tier 1 vs simple neural network"""
import torch
import torch.nn as nn
import numpy as np
import time
from tier1_system import Tier1FastSystem

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Simple baseline network
class SimpleNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 32),
            nn.Tanh(),
            nn.Linear(32, 32),
            nn.Tanh(),
            nn.Linear(32, 1)
        )
    
    def forward(self, x):
        return self.net(x)

# Create both systems
tier1 = Tier1FastSystem(device=device)
baseline = SimpleNN().to(device)

# Test data
batch = 16
test_input = torch.randn(batch, 1).to(device) * 0.1

# Warmup
for _ in range(10):
    tier1(test_input)
    baseline(test_input)

# Benchmark
n_runs = 1000

# Tier 1
tier1.reset_all_states(batch)
start = time.time()
for _ in range(n_runs):
    result = tier1(test_input, return_resonance=True)
tier1_time = (time.time() - start) / n_runs * 1000

# Baseline
start = time.time()
for _ in range(n_runs):
    out = baseline(test_input)
baseline_time = (time.time() - start) / n_runs * 1000

print("="*60)
print("  ⚡ SPEED COMPARISON")
print("="*60)
print(f"Simple NN:     {baseline_time:.4f} ms")
print(f"Tier 1 (ESN):  {tier1_time:.4f} ms")
print(f"Speedup:       {baseline_time/tier1_time:.2f}x")
print()
print(f"Tier 1 includes:")
print(f"  • 3 parallel reservoir layers")
print(f"  • Resonance computation")
print(f"  • φ-weighted prediction")
print("="*60)