#!/usr/bin/env python3
"""
PyTorch Benchmark for PINN MRE Helmholtz solver.

This script benchmarks forward pass and training step performance on CPU and GPU
for direct comparison with the Rust implementation.

Run with: python benchmark_pytorch.py
"""

import time
import numpy as np
import torch
import torch.nn as nn
from dataclasses import dataclass
import statistics

# Disable gradient computation for forward-only benchmarks
torch.set_grad_enabled(True)

@dataclass
class Config:
    """Configuration matching the Rust implementation."""
    RHO: float = 1040.0
    FREQ: float = 50.0
    L: float = 0.1
    U0: float = 1e-6
    G_PRIME_TRUE: float = 3000.0
    G_DOUBLE_TRUE: float = 1500.0
    N_DATA: int = 200
    N_PDE: int = 200
    U_LAYERS: int = 4
    U_HIDDEN: int = 64
    U_FF_DIM: int = 64
    U_FF_SCALE: float = 10.0
    LR: float = 1e-3
    EPOCHS: int = 100


def calculate_k(cfg: Config) -> complex:
    """Calculates the complex wave number k."""
    omega = 2.0 * np.pi * cfg.FREQ
    G_complex = cfg.G_PRIME_TRUE + 1j * cfg.G_DOUBLE_TRUE
    k = np.sqrt(cfg.RHO * omega**2 / G_complex)
    if np.imag(k) < 0:
        k = -k
    return k


def synthesize_displacement(cfg: Config, device: torch.device) -> tuple:
    """Generates displacement data u(x)."""
    k = calculate_k(cfg)
    x = np.linspace(0.0, cfg.L, cfg.N_DATA)
    u_complex = cfg.U0 * np.exp(1j * k * x)

    x_norm = torch.from_numpy(x / cfg.L).float().reshape(-1, 1).to(device)
    u_data = torch.from_numpy(
        np.stack([np.real(u_complex), np.imag(u_complex)], axis=1)
    ).float().to(device)

    return x_norm, u_data, k


class LffnUNet1D(nn.Module):
    """Learnable Fourier Feature Network (LFFN-MLP) - matches Rust implementation."""
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg

        # Learnable Fourier features
        B_init = torch.randn(1, cfg.U_FF_DIM) * cfg.U_FF_SCALE
        self.B_learnable = nn.Parameter(B_init, requires_grad=True)

        # Build MLP
        in_dim = cfg.U_FF_DIM * 2
        layers = []
        dim = in_dim
        for _ in range(cfg.U_LAYERS):
            layers.append(nn.Linear(dim, cfg.U_HIDDEN))
            layers.append(nn.Tanh())
            dim = cfg.U_HIDDEN
        layers.append(nn.Linear(dim, 2))
        self.net = nn.Sequential(*layers)

    def forward(self, x_norm: torch.Tensor) -> torch.Tensor:
        y = 2.0 * np.pi * (x_norm @ self.B_learnable)
        feat = torch.cat([torch.sin(y), torch.cos(y)], dim=-1)
        return self.net(feat)


def benchmark_forward_pass(n_points: int, device: torch.device, n_warmup: int = 10, n_iter: int = 50) -> dict:
    """Benchmark forward pass only."""
    cfg = Config(N_DATA=n_points, N_PDE=n_points)

    # Setup
    x_norm, u_data, k = synthesize_displacement(cfg, device)
    x_norm.requires_grad_(True)

    model = LffnUNet1D(cfg).to(device)
    model.eval()

    # Warmup
    for _ in range(n_warmup):
        with torch.no_grad():
            _ = model(x_norm)

    if device.type == 'cuda':
        torch.cuda.synchronize()

    # Benchmark
    times = []
    for _ in range(n_iter):
        if device.type == 'cuda':
            torch.cuda.synchronize()

        start = time.perf_counter()
        with torch.no_grad():
            output = model(x_norm)

        if device.type == 'cuda':
            torch.cuda.synchronize()

        end = time.perf_counter()
        times.append((end - start) * 1000)  # Convert to ms

    return {
        'mean_ms': statistics.mean(times),
        'std_ms': statistics.stdev(times) if len(times) > 1 else 0,
        'min_ms': min(times),
        'max_ms': max(times),
        'throughput_kelem_s': (n_points / statistics.mean(times)),
    }


def benchmark_training_step(n_points: int, device: torch.device, n_warmup: int = 5, n_iter: int = 30) -> dict:
    """Benchmark a single training step (forward + backward + optimizer step)."""
    cfg = Config(N_DATA=n_points, N_PDE=n_points)

    # Setup
    x_norm, u_data, k = synthesize_displacement(cfg, device)
    x_norm.requires_grad_(True)

    # Normalize data (matching Rust)
    u_scale = float(torch.abs(u_data).max().item() + 1e-16)
    u_data_norm = u_data / u_scale

    model = LffnUNet1D(cfg).to(device)
    model.train()

    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.LR)
    mse_loss = nn.MSELoss()

    # Warmup
    for _ in range(n_warmup):
        optimizer.zero_grad()
        output = model(x_norm)
        loss = mse_loss(output, u_data_norm)
        loss.backward()
        optimizer.step()

    if device.type == 'cuda':
        torch.cuda.synchronize()

    # Benchmark
    times = []
    for _ in range(n_iter):
        if device.type == 'cuda':
            torch.cuda.synchronize()

        start = time.perf_counter()

        optimizer.zero_grad()
        output = model(x_norm)
        loss = mse_loss(output, u_data_norm)
        loss.backward()
        optimizer.step()

        if device.type == 'cuda':
            torch.cuda.synchronize()

        end = time.perf_counter()
        times.append((end - start) * 1000)  # Convert to ms

    return {
        'mean_ms': statistics.mean(times),
        'std_ms': statistics.stdev(times) if len(times) > 1 else 0,
        'min_ms': min(times),
        'max_ms': max(times),
        'throughput_kelem_s': (n_points / statistics.mean(times)),
    }


def benchmark_training_100_epochs(n_points: int, device: torch.device, n_iter: int = 5) -> dict:
    """Benchmark 100 epochs of training."""
    cfg = Config(N_DATA=n_points, N_PDE=n_points, EPOCHS=100)

    times = []
    for _ in range(n_iter):
        # Fresh setup for each run
        x_norm, u_data, k = synthesize_displacement(cfg, device)
        x_norm.requires_grad_(True)

        u_scale = float(torch.abs(u_data).max().item() + 1e-16)
        u_data_norm = u_data / u_scale

        model = LffnUNet1D(cfg).to(device)
        model.train()

        optimizer = torch.optim.Adam(model.parameters(), lr=cfg.LR)
        mse_loss = nn.MSELoss()

        if device.type == 'cuda':
            torch.cuda.synchronize()

        start = time.perf_counter()

        for epoch in range(100):
            optimizer.zero_grad()
            output = model(x_norm)
            loss = mse_loss(output, u_data_norm)
            loss.backward()
            optimizer.step()

        if device.type == 'cuda':
            torch.cuda.synchronize()

        end = time.perf_counter()
        times.append((end - start) * 1000)  # Convert to ms

    return {
        'mean_ms': statistics.mean(times),
        'std_ms': statistics.stdev(times) if len(times) > 1 else 0,
        'min_ms': min(times),
        'max_ms': max(times),
        'epochs_per_sec': 100 / (statistics.mean(times) / 1000),
    }


def print_results(name: str, results: dict):
    """Print benchmark results in a formatted way."""
    print(f"  {name}:")
    print(f"    Time: {results['mean_ms']:.3f} ms (+/- {results['std_ms']:.3f} ms)")
    print(f"    Min: {results['min_ms']:.3f} ms, Max: {results['max_ms']:.3f} ms")
    if 'throughput_kelem_s' in results:
        print(f"    Throughput: {results['throughput_kelem_s']:.2f} Kelem/s")
    if 'epochs_per_sec' in results:
        print(f"    Epochs/sec: {results['epochs_per_sec']:.2f}")


def main():
    print("=" * 70)
    print("PyTorch PINN MRE Helmholtz Benchmark")
    print("=" * 70)

    # Check available devices
    devices = [('cpu', torch.device('cpu'))]
    if torch.cuda.is_available():
        devices.append(('cuda', torch.device('cuda')))
        print(f"CUDA available: {torch.cuda.get_device_name(0)}")
    else:
        print("CUDA not available, running CPU benchmarks only")

    print(f"PyTorch version: {torch.__version__}")
    print()

    batch_sizes = [200, 1000, 10000]

    for device_name, device in devices:
        print(f"\n{'=' * 70}")
        print(f"Device: {device_name.upper()}")
        print(f"{'=' * 70}")

        # Forward pass benchmarks
        print("\n--- Forward Pass (inference only) ---")
        for n_points in batch_sizes:
            results = benchmark_forward_pass(n_points, device)
            print_results(f"n_points={n_points}", results)

        # Training step benchmarks
        print("\n--- Single Training Step (forward + backward + optimizer) ---")
        for n_points in batch_sizes[:2]:  # Only 200 and 1000 for training
            results = benchmark_training_step(n_points, device)
            print_results(f"n_points={n_points}", results)

        # 100 epochs benchmark
        print("\n--- 100 Epochs Training ---")
        results = benchmark_training_100_epochs(200, device)
        print_results("n_points=200, epochs=100", results)

    print("\n" + "=" * 70)
    print("Benchmark complete!")
    print("=" * 70)


if __name__ == "__main__":
    main()