#!/usr/bin/env python3
"""
Benchmark runner for PyTorch PINN - mirrors Rust benchmark structure.

This script benchmarks the PyTorch PINN implementation to compare with RustyTorch++.

Usage:
    python benchmark_runner.py [--pytorch-pinn-path PATH] [--output OUTPUT]

The script expects a PyTorch PINN implementation with the following interface:
    - Config class with physics/network/training parameters
    - Mre1DPinnSolver class with:
        - forward() method for network evaluation
        - compute_pde_residual() method for PDE computation
        - training_step() method for single training iteration
        - train() method for full training loop

If no --pytorch-pinn-path is provided, a minimal reference implementation is used.
"""

import json
import time
import gc
import argparse
import sys
from dataclasses import dataclass, asdict, field
from typing import List, Dict, Optional, Tuple
from pathlib import Path

# Optional imports with fallbacks
try:
    import tracemalloc
    HAS_TRACEMALLOC = True
except ImportError:
    HAS_TRACEMALLOC = False

try:
    import numpy as np
    HAS_NUMPY = True
except ImportError:
    HAS_NUMPY = False
    print("Warning: numpy not available, some benchmarks will be skipped")

try:
    import torch
    import torch.nn as nn
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    print("Warning: torch not available, PyTorch benchmarks will be skipped")


@dataclass
class BenchmarkResult:
    """Result of a single benchmark."""
    name: str
    n_points: int
    mean_time_ms: float
    std_time_ms: float
    min_time_ms: float
    max_time_ms: float
    peak_memory_mb: float
    iterations: int
    throughput_elements_per_sec: float = 0.0


@dataclass
class BenchmarkSuite:
    """Collection of benchmark results."""
    framework: str
    device: str
    version: str
    benchmarks: List[BenchmarkResult] = field(default_factory=list)


def warmup(func, n_warmup: int = 10):
    """Run warmup iterations."""
    for _ in range(n_warmup):
        _ = func()


def measure_time(func, n_iterations: int = 100) -> Tuple[List[float], float]:
    """Measure execution time and peak memory."""
    times = []
    peak_memory = 0.0

    if HAS_TRACEMALLOC:
        tracemalloc.start()

    if HAS_TORCH and torch.cuda.is_available():
        torch.cuda.synchronize()

    for _ in range(n_iterations):
        start = time.perf_counter()
        _ = func()
        if HAS_TORCH and torch.cuda.is_available():
            torch.cuda.synchronize()
        times.append((time.perf_counter() - start) * 1000)

    if HAS_TRACEMALLOC:
        _, peak_memory = tracemalloc.get_traced_memory()
        tracemalloc.stop()
        peak_memory = peak_memory / 1024 / 1024  # Convert to MB

    return times, peak_memory


def create_benchmark_result(
    name: str,
    n_points: int,
    times: List[float],
    peak_memory: float,
    n_iterations: int
) -> BenchmarkResult:
    """Create a BenchmarkResult from timing data."""
    import statistics
    mean_time = statistics.mean(times)
    return BenchmarkResult(
        name=name,
        n_points=n_points,
        mean_time_ms=mean_time,
        std_time_ms=statistics.stdev(times) if len(times) > 1 else 0.0,
        min_time_ms=min(times),
        max_time_ms=max(times),
        peak_memory_mb=peak_memory,
        iterations=n_iterations,
        throughput_elements_per_sec=(n_points / mean_time * 1000) if mean_time > 0 else 0.0
    )


# =============================================================================
# MINIMAL PYTORCH PINN REFERENCE IMPLEMENTATION
# =============================================================================

if HAS_TORCH and HAS_NUMPY:
    class MinimalConfig:
        """Minimal configuration matching Rust Config."""
        def __init__(self, n_data: int = 200):
            # Physics
            self.rho = 1040.0
            self.freq = 50.0
            self.l = 0.1
            self.u0 = 1e-6
            self.g_prime_true = 3000.0
            self.g_double_true = 1500.0
            # Grid
            self.n_data = n_data
            self.n_pde = n_data
            # Network
            self.u_layers = 4
            self.u_hidden = 64
            self.u_ff_dim = 64
            self.u_ff_scale = 10.0
            # Training
            self.lr = 1e-3
            self.epochs = 50000
            self.scheduler_patience = 500
            self.scheduler_factor = 0.5
            # Loss
            self.data_weight = 1.0
            self.pde_weight = 1e-6

    def calculate_k(cfg: MinimalConfig) -> complex:
        """Calculate complex wave number."""
        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 k.imag < 0:
            k = -k
        return k

    def synthesize_displacement(cfg: MinimalConfig) -> Tuple[np.ndarray, np.ndarray, np.ndarray, complex]:
        """Generate synthetic displacement data."""
        k = calculate_k(cfg)
        x = np.linspace(0.0, cfg.l, cfg.n_data)
        u_complex = cfg.u0 * np.exp(1j * k * x)
        return x, u_complex.real, u_complex.imag, k

    class LffnUNet1D(nn.Module):
        """Learnable Fourier Feature Network with MLP."""
        def __init__(self, cfg: MinimalConfig):
            super().__init__()
            self.B_learnable = nn.Parameter(
                torch.randn(1, cfg.u_ff_dim) * cfg.u_ff_scale
            )
            layers = []
            dim = cfg.u_ff_dim * 2
            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
            sin_feat = torch.sin(y)
            cos_feat = torch.cos(y)
            feat = torch.cat([sin_feat, cos_feat], dim=-1)
            return self.net(feat)

    class MinimalPinnSolver:
        """Minimal PINN solver for benchmarking."""
        def __init__(self, cfg: MinimalConfig, device: str = 'cpu'):
            self.cfg = cfg
            self.device = torch.device(device)

            # Physics
            self.omega = 2.0 * np.pi * cfg.freq
            self.k_true = calculate_k(cfg)

            # Generate data
            x, u_r, u_i, _ = synthesize_displacement(cfg)
            u_scale = np.max(np.sqrt(u_r**2 + u_i**2)) + 1e-16
            self.u_scale = u_scale

            # Normalize and convert to tensors
            x_norm = torch.tensor(x / cfg.l, dtype=torch.float32, device=self.device).reshape(-1, 1)
            u_data = torch.tensor(
                np.stack([u_r / u_scale, u_i / u_scale], axis=1),
                dtype=torch.float32, device=self.device
            )

            self.x_data = x_norm.requires_grad_(True)
            self.u_data_target = u_data

            # Network
            self.u_net = LffnUNet1D(cfg).to(self.device)

            # Optimizer
            self.optimizer = torch.optim.Adam(self.u_net.parameters(), lr=cfg.lr)
            self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer, mode='min', factor=cfg.scheduler_factor,
                patience=cfg.scheduler_patience
            )

        def forward(self) -> torch.Tensor:
            """Forward pass through network."""
            return self.u_net(self.x_data)

        def compute_data_loss(self, u_pred: torch.Tensor) -> torch.Tensor:
            """Compute MSE data loss."""
            return torch.mean((u_pred - self.u_data_target)**2)

        def compute_pde_residual(self) -> Tuple[float, float]:
            """Compute PDE residual using analytical derivatives."""
            x = np.linspace(0.0, self.cfg.l, self.cfg.n_data)
            u_complex = self.cfg.u0 * np.exp(1j * self.k_true * x)

            # Analytical d²u/dx²
            d2u = -(self.k_true**2) * u_complex

            # Helmholtz: d²u/dx² + k²u = 0
            k_sq = self.k_true**2
            residual = d2u + k_sq * u_complex

            mse_rr = np.mean(residual.real**2)
            mse_ri = np.mean(residual.imag**2)
            return mse_rr, mse_ri

        def training_step(self) -> Tuple[float, float, float]:
            """Single training step."""
            self.optimizer.zero_grad()
            u_pred = self.forward()
            loss_data = self.compute_data_loss(u_pred)
            mse_rr, mse_ri = self.compute_pde_residual()
            loss_pde = mse_rr + mse_ri
            total_loss = self.cfg.data_weight * loss_data.item() + self.cfg.pde_weight * loss_pde
            loss_data.backward()
            self.optimizer.step()
            self.scheduler.step(loss_data)
            return loss_data.item(), loss_pde, total_loss

        def train(self, epochs: int = None) -> float:
            """Full training loop."""
            epochs = epochs or self.cfg.epochs
            best_loss = float('inf')
            for _ in range(epochs):
                _, _, total_loss = self.training_step()
                if total_loss < best_loss:
                    best_loss = total_loss
            return best_loss


# =============================================================================
# BENCHMARK FUNCTIONS
# =============================================================================

def benchmark_forward_pass(suite: BenchmarkSuite, point_sizes: List[int] = [200, 1000, 10000], forced_device: str = "auto"):
    """Benchmark forward pass (network evaluation)."""
    if not HAS_TORCH:
        print("Skipping forward_pass benchmark: torch not available")
        return

    if forced_device == "auto":
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    elif forced_device == "cuda" and torch.cuda.is_available():
        device = 'cuda'
    else:
        device = 'cpu'

    for n_points in point_sizes:
        print(f"  forward_pass n={n_points}...", end=" ", flush=True)

        cfg = MinimalConfig(n_data=n_points)
        solver = MinimalPinnSolver(cfg, device=device)

        warmup(solver.forward)
        times, peak_memory = measure_time(solver.forward, n_iterations=100)

        result = create_benchmark_result(
            name="forward_pass/lffn_mlp",
            n_points=n_points,
            times=times,
            peak_memory=peak_memory,
            n_iterations=100
        )
        suite.benchmarks.append(result)
        print(f"{result.mean_time_ms:.3f}ms")

        del solver
        gc.collect()
        if HAS_TORCH and torch.cuda.is_available():
            torch.cuda.empty_cache()


def benchmark_pde_residual(suite: BenchmarkSuite, point_sizes: List[int] = [200, 1000, 10000], forced_device: str = "auto"):
    """Benchmark PDE residual computation."""
    if not HAS_TORCH:
        print("Skipping pde_residual benchmark: torch not available")
        return

    if forced_device == "auto":
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    elif forced_device == "cuda" and torch.cuda.is_available():
        device = 'cuda'
    else:
        device = 'cpu'

    for n_points in point_sizes:
        print(f"  pde_residual n={n_points}...", end=" ", flush=True)

        cfg = MinimalConfig(n_data=n_points)
        solver = MinimalPinnSolver(cfg, device=device)

        warmup(solver.compute_pde_residual)
        times, peak_memory = measure_time(solver.compute_pde_residual, n_iterations=100)

        result = create_benchmark_result(
            name="pde_residual/helmholtz",
            n_points=n_points,
            times=times,
            peak_memory=peak_memory,
            n_iterations=100
        )
        suite.benchmarks.append(result)
        print(f"{result.mean_time_ms:.3f}ms")

        del solver
        gc.collect()


def benchmark_training_step(suite: BenchmarkSuite, point_sizes: List[int] = [200, 1000], forced_device: str = "auto"):
    """Benchmark single training step."""
    if not HAS_TORCH:
        print("Skipping training_step benchmark: torch not available")
        return

    if forced_device == "auto":
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    elif forced_device == "cuda" and torch.cuda.is_available():
        device = 'cuda'
    else:
        device = 'cpu'

    for n_points in point_sizes:
        print(f"  training_step n={n_points}...", end=" ", flush=True)

        cfg = MinimalConfig(n_data=n_points)
        solver = MinimalPinnSolver(cfg, device=device)

        warmup(solver.training_step, n_warmup=5)
        times, peak_memory = measure_time(solver.training_step, n_iterations=50)

        result = create_benchmark_result(
            name="training_step/single_step",
            n_points=n_points,
            times=times,
            peak_memory=peak_memory,
            n_iterations=50
        )
        suite.benchmarks.append(result)
        print(f"{result.mean_time_ms:.3f}ms")

        del solver
        gc.collect()
        if HAS_TORCH and torch.cuda.is_available():
            torch.cuda.empty_cache()


def benchmark_data_generation(suite: BenchmarkSuite, point_sizes: List[int] = [200, 1000, 10000, 100000]):
    """Benchmark data generation."""
    if not HAS_NUMPY:
        print("Skipping data_generation benchmark: numpy not available")
        return

    for n_points in point_sizes:
        print(f"  data_generation n={n_points}...", end=" ", flush=True)

        cfg = MinimalConfig(n_data=n_points)

        def gen_data():
            return synthesize_displacement(cfg)

        warmup(gen_data)
        times, peak_memory = measure_time(gen_data, n_iterations=100)

        result = create_benchmark_result(
            name="data_generation/synthesize_displacement",
            n_points=n_points,
            times=times,
            peak_memory=peak_memory,
            n_iterations=100
        )
        suite.benchmarks.append(result)
        print(f"{result.mean_time_ms:.3f}ms")


def benchmark_wave_number(suite: BenchmarkSuite):
    """Benchmark wave number calculation."""
    if not HAS_NUMPY:
        print("Skipping wave_number benchmark: numpy not available")
        return

    print("  calculate_k...", end=" ", flush=True)
    cfg = MinimalConfig()

    def calc_k():
        return calculate_k(cfg)

    warmup(calc_k)
    times, peak_memory = measure_time(calc_k, n_iterations=1000)

    result = create_benchmark_result(
        name="wave_number/calculate_k",
        n_points=1,
        times=times,
        peak_memory=peak_memory,
        n_iterations=1000
    )
    suite.benchmarks.append(result)
    print(f"{result.mean_time_ms:.4f}ms")


def benchmark_training_100_epochs(suite: BenchmarkSuite, point_sizes: List[int] = [200], forced_device: str = "auto"):
    """Benchmark 100 epochs of training."""
    if not HAS_TORCH:
        print("Skipping training_100_epochs benchmark: torch not available")
        return

    if forced_device == "auto":
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    elif forced_device == "cuda" and torch.cuda.is_available():
        device = 'cuda'
    else:
        device = 'cpu'

    for n_points in point_sizes:
        print(f"  training_100_epochs n={n_points}...", end=" ", flush=True)

        def run_training():
            cfg = MinimalConfig(n_data=n_points)
            solver = MinimalPinnSolver(cfg, device=device)
            return solver.train(epochs=100)

        # Only 5 iterations since this is slow
        times, peak_memory = measure_time(run_training, n_iterations=5)

        result = create_benchmark_result(
            name="training_100_epochs/train_100",
            n_points=n_points,
            times=times,
            peak_memory=peak_memory,
            n_iterations=5
        )
        suite.benchmarks.append(result)
        print(f"{result.mean_time_ms:.1f}ms")

        gc.collect()
        if HAS_TORCH and torch.cuda.is_available():
            torch.cuda.empty_cache()


def run_all_benchmarks(point_sizes: Optional[List[int]] = None, forced_device: str = "auto") -> BenchmarkSuite:
    """Run complete benchmark suite.

    Args:
        point_sizes: List of point sizes to benchmark
        forced_device: Device to use - "cpu", "cuda", or "auto" (detect)
    """
    point_sizes = point_sizes or [200, 1000, 10000]

    # Determine device based on forced_device argument
    if forced_device == "auto":
        if HAS_TORCH and torch.cuda.is_available():
            device = f"cuda:{torch.cuda.get_device_name(0)}"
        else:
            device = "cpu"
    elif forced_device == "cuda":
        if HAS_TORCH and torch.cuda.is_available():
            device = f"cuda:{torch.cuda.get_device_name(0)}"
        else:
            print("Warning: CUDA requested but not available, falling back to CPU")
            device = "cpu"
    else:
        device = "cpu"

    version = "unknown"
    if HAS_TORCH:
        version = torch.__version__

    suite = BenchmarkSuite(
        framework="pytorch",
        device=device,
        version=version
    )

    print(f"\n=== PyTorch PINN Benchmarks ===")
    print(f"Device: {device}")
    print(f"Version: {version}")
    print()

    print("Running benchmarks:")
    benchmark_wave_number(suite)
    benchmark_data_generation(suite, point_sizes + [100000])
    benchmark_forward_pass(suite, point_sizes, forced_device)
    benchmark_pde_residual(suite, point_sizes, forced_device)
    benchmark_training_step(suite, [200, 1000], forced_device)
    benchmark_training_100_epochs(suite, [200], forced_device)

    return suite


def main():
    parser = argparse.ArgumentParser(description="PyTorch PINN Benchmark Runner")
    parser.add_argument(
        "--pytorch-pinn-path",
        type=str,
        default=None,
        help="Path to PyTorch PINN implementation (uses minimal reference if not provided)"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="python_results.json",
        help="Output JSON file for results"
    )
    parser.add_argument(
        "--point-sizes",
        type=str,
        default="200,1000,10000",
        help="Comma-separated list of point sizes to benchmark"
    )
    parser.add_argument(
        "--device",
        type=str,
        choices=["cpu", "cuda", "auto"],
        default="auto",
        help="Device to use: cpu, cuda, or auto (detect)"
    )
    args = parser.parse_args()

    point_sizes = [int(x) for x in args.point_sizes.split(",")]
    forced_device = args.device

    if args.pytorch_pinn_path:
        print(f"Loading custom PyTorch PINN from: {args.pytorch_pinn_path}")
        # Import the reference implementation if available
        import importlib.util
        spec = importlib.util.spec_from_file_location("pytorch_pinn", args.pytorch_pinn_path)
        if spec and spec.loader:
            pytorch_pinn = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(pytorch_pinn)
            print(f"Loaded custom PINN from: {args.pytorch_pinn_path}")
            # Use the loaded implementation for benchmarks
            global MinimalConfig, MinimalPinnSolver, LffnUNet1D
            MinimalConfig = pytorch_pinn.Config
            # Note: Would need adapter for Mre1DFixedGPinn -> MinimalPinnSolver interface
        else:
            print("Failed to load custom PINN, using minimal reference")

    suite = run_all_benchmarks(point_sizes, forced_device)

    # Save results
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w") as f:
        json.dump(asdict(suite), f, indent=2)

    print(f"\nResults saved to: {output_path}")

    # Print summary table
    print("\n" + "=" * 70)
    print("BENCHMARK SUMMARY")
    print("=" * 70)
    print(f"{'Benchmark':<40} {'Points':>8} {'Mean (ms)':>12} {'Std':>10}")
    print("-" * 70)
    for result in suite.benchmarks:
        print(f"{result.name:<40} {result.n_points:>8} {result.mean_time_ms:>12.3f} {result.std_time_ms:>10.3f}")
    print("=" * 70)


if __name__ == "__main__":
    main()