#!/usr/bin/env python3
"""
Generate markdown comparison report from Python and Rust benchmark results.

Supports separate CPU and GPU benchmark results for fair comparison.

Usage:
    python generate_report.py \
        --python-cpu-results results/python_cpu_results.json \
        --python-gpu-results results/python_gpu_results.json \
        --rust-cpu-output results/rust_cpu_output.txt \
        --rust-cuda-output results/rust_cuda_output.txt \
        --output results/comparison_report.md
"""

import json
import re
import argparse
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple


def load_python_results(path: Path) -> Dict:
    """Load Python benchmark results from JSON."""
    if not path or not path.exists():
        return {"benchmarks": [], "device": "unknown", "version": "unknown"}
    with open(path) as f:
        return json.load(f)


def parse_rust_criterion_output(path: Path) -> Dict[str, Dict]:
    """Parse Rust Criterion benchmark output."""
    if not path or not path.exists():
        return {}

    results = {}
    content = path.read_text()

    # Pattern: benchmark_name    time:   [lower mean upper]
    # Example: forward_pass/lffn_mlp/200    time:   [5.2345 ms 5.3456 ms 5.4567 ms]
    pattern = r'(\S+)\s+time:\s+\[(\d+\.?\d*)\s*(ns|µs|us|ms|s)\s+(\d+\.?\d*)\s*(ns|µs|us|ms|s)\s+(\d+\.?\d*)\s*(ns|µs|us|ms|s)\]'

    for match in re.finditer(pattern, content):
        name = match.group(1)
        mean_val = float(match.group(4))
        mean_unit = match.group(5)

        # Convert to ms
        multipliers = {'ns': 1e-6, 'µs': 1e-3, 'us': 1e-3, 'ms': 1.0, 's': 1000.0}
        mean_ms = mean_val * multipliers.get(mean_unit, 1.0)

        # Extract n_points from name if present
        parts = name.split('/')
        n_points = 0
        for part in parts:
            try:
                n_points = int(part)
                break
            except ValueError:
                continue

        results[name] = {
            'mean_ms': mean_ms,
            'n_points': n_points,
            'raw_name': name
        }

    return results


def normalize_category(name: str) -> str:
    """Normalize benchmark category names for matching."""
    # Map various naming conventions to standard names
    name_lower = name.lower().replace('_', ' ').replace('/', ' ')

    if 'forward' in name_lower or 'lffn' in name_lower:
        return 'forward_pass'
    elif 'pde' in name_lower and 'residual' in name_lower:
        return 'pde_residual'
    elif 'training' in name_lower and '100' in name_lower:
        return 'training_100_epochs'
    elif 'training' in name_lower and 'step' in name_lower:
        return 'training_step'
    elif 'data' in name_lower and 'gen' in name_lower:
        return 'data_generation'
    elif 'wave' in name_lower or 'calculate_k' in name_lower:
        return 'calculate_k'
    elif 'mse' in name_lower:
        return 'mse_loss'
    else:
        # Fallback: use first part of name
        parts = name.replace('/', ' ').replace('_', ' ').split()
        return parts[0] if parts else name


def extract_benchmark_info(name: str) -> Tuple[str, int]:
    """Extract normalized category and n_points from benchmark name."""
    category = normalize_category(name)

    # Extract n_points
    n_points = 0
    parts = name.replace('/', ' ').replace('_', ' ').split()
    for part in parts:
        try:
            n_points = int(part)
            break
        except ValueError:
            continue

    return category, n_points


def generate_report(
    python_cpu_results: Dict,
    python_gpu_results: Dict,
    rust_cpu_results: Dict,
    rust_cuda_results: Dict
) -> str:
    """Generate markdown comparison report with CPU vs CPU and GPU vs GPU sections."""

    lines = []
    lines.append("# PINN Benchmark Report: Python (PyTorch) vs Rust (RustyTorch++)")
    lines.append("")
    lines.append(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    lines.append("")

    # System info
    lines.append("## System Information")
    lines.append("")
    py_cpu_device = python_cpu_results.get('device', 'N/A')
    py_gpu_device = python_gpu_results.get('device', 'N/A')
    py_version = python_cpu_results.get('version', 'unknown') or python_gpu_results.get('version', 'unknown')
    lines.append(f"- **PyTorch Version**: {py_version}")
    lines.append(f"- **Python CPU Device**: {py_cpu_device}")
    lines.append(f"- **Python GPU Device**: {py_gpu_device}")
    lines.append(f"- **Rust Framework**: RustyTorch++ (rtx-tensor, rtx-autograd)")
    lines.append("")

    # Organize results by category and n_points
    cpu_data = {}  # category -> n_points -> {python: ms, rust: ms}
    gpu_data = {}  # category -> n_points -> {python: ms, rust: ms}

    # Process Python CPU results
    for bench in python_cpu_results.get('benchmarks', []):
        name = bench.get('name', '')
        category, n_points = extract_benchmark_info(name)
        n_points = n_points or bench.get('n_points', 0)

        if category not in cpu_data:
            cpu_data[category] = {}
        if n_points not in cpu_data[category]:
            cpu_data[category][n_points] = {}
        cpu_data[category][n_points]['python'] = bench.get('mean_time_ms', 0)

    # Process Python GPU results
    for bench in python_gpu_results.get('benchmarks', []):
        name = bench.get('name', '')
        category, n_points = extract_benchmark_info(name)
        n_points = n_points or bench.get('n_points', 0)

        if category not in gpu_data:
            gpu_data[category] = {}
        if n_points not in gpu_data[category]:
            gpu_data[category][n_points] = {}
        gpu_data[category][n_points]['python'] = bench.get('mean_time_ms', 0)

    # Process Rust CPU results
    for name, data in rust_cpu_results.items():
        category, n_points = extract_benchmark_info(name)
        n_points = n_points or data.get('n_points', 0)

        if category not in cpu_data:
            cpu_data[category] = {}
        if n_points not in cpu_data[category]:
            cpu_data[category][n_points] = {}
        cpu_data[category][n_points]['rust'] = data.get('mean_ms', 0)

    # Process Rust CUDA results
    for name, data in rust_cuda_results.items():
        category, n_points = extract_benchmark_info(name)
        n_points = n_points or data.get('n_points', 0)

        if category not in gpu_data:
            gpu_data[category] = {}
        if n_points not in gpu_data[category]:
            gpu_data[category][n_points] = {}
        gpu_data[category][n_points]['rust'] = data.get('mean_ms', 0)

    # Generate CPU comparison tables
    lines.append("## CPU Comparison (PyTorch CPU vs RustyTorch++ CPU)")
    lines.append("")

    if cpu_data:
        for category in sorted(cpu_data.keys()):
            if not category:
                continue

            lines.append(f"### {category.replace('_', ' ').title()}")
            lines.append("")
            lines.append("| Points | PyTorch CPU (ms) | Rust CPU (ms) | Speedup |")
            lines.append("|--------|------------------|---------------|---------|")

            for n_points in sorted(cpu_data[category].keys()):
                data = cpu_data[category][n_points]
                py_time = data.get('python', 0)
                rust_time = data.get('rust', 0)

                # Calculate speedup (Rust vs PyTorch)
                speedup = "-"
                if py_time > 0 and rust_time > 0:
                    ratio = py_time / rust_time
                    if ratio >= 1:
                        speedup = f"**Rust {ratio:.1f}x**"
                    else:
                        speedup = f"PyTorch {1/ratio:.1f}x"

                py_str = f"{py_time:.3f}" if py_time > 0 else "-"
                rust_str = f"{rust_time:.3f}" if rust_time > 0 else "-"
                lines.append(f"| {n_points} | {py_str} | {rust_str} | {speedup} |")

            lines.append("")
    else:
        lines.append("*No CPU benchmark data available*")
        lines.append("")

    # Generate GPU comparison tables
    lines.append("## GPU Comparison (PyTorch CUDA vs RustyTorch++ CUDA)")
    lines.append("")

    if gpu_data:
        for category in sorted(gpu_data.keys()):
            if not category:
                continue

            lines.append(f"### {category.replace('_', ' ').title()}")
            lines.append("")
            lines.append("| Points | PyTorch GPU (ms) | Rust CUDA (ms) | Speedup |")
            lines.append("|--------|------------------|----------------|---------|")

            for n_points in sorted(gpu_data[category].keys()):
                data = gpu_data[category][n_points]
                py_time = data.get('python', 0)
                rust_time = data.get('rust', 0)

                # Calculate speedup
                speedup = "-"
                if py_time > 0 and rust_time > 0:
                    ratio = py_time / rust_time
                    if ratio >= 1:
                        speedup = f"**Rust {ratio:.1f}x**"
                    else:
                        speedup = f"PyTorch {1/ratio:.1f}x"

                py_str = f"{py_time:.3f}" if py_time > 0 else "-"
                rust_str = f"{rust_time:.3f}" if rust_time > 0 else "-"
                lines.append(f"| {n_points} | {py_str} | {rust_str} | {speedup} |")

            lines.append("")
    else:
        lines.append("*No GPU benchmark data available*")
        lines.append("")

    # Summary section
    lines.append("## Summary")
    lines.append("")

    # Calculate CPU speedups
    cpu_speedups = []
    for category in cpu_data.values():
        for data in category.values():
            py_time = data.get('python', 0)
            rust_time = data.get('rust', 0)
            if py_time > 0 and rust_time > 0:
                cpu_speedups.append(py_time / rust_time)

    if cpu_speedups:
        avg_cpu_speedup = sum(cpu_speedups) / len(cpu_speedups)
        rust_faster_cpu = sum(1 for s in cpu_speedups if s > 1.0)
        lines.append("### CPU Results")
        lines.append(f"- **Average speedup (Rust vs PyTorch)**: {avg_cpu_speedup:.2f}x")
        lines.append(f"- **Benchmarks where Rust is faster**: {rust_faster_cpu}/{len(cpu_speedups)}")
        lines.append("")

    # Calculate GPU speedups
    gpu_speedups = []
    for category in gpu_data.values():
        for data in category.values():
            py_time = data.get('python', 0)
            rust_time = data.get('rust', 0)
            if py_time > 0 and rust_time > 0:
                gpu_speedups.append(py_time / rust_time)

    if gpu_speedups:
        avg_gpu_speedup = sum(gpu_speedups) / len(gpu_speedups)
        rust_faster_gpu = sum(1 for s in gpu_speedups if s > 1.0)
        lines.append("### GPU Results")
        lines.append(f"- **Average speedup (Rust vs PyTorch)**: {avg_gpu_speedup:.2f}x")
        lines.append(f"- **Benchmarks where Rust is faster**: {rust_faster_gpu}/{len(gpu_speedups)}")
        lines.append("")

    if not cpu_speedups and not gpu_speedups:
        lines.append("- Not enough matching data to calculate speedup statistics")
        lines.append("")

    # Notes
    lines.append("## Notes")
    lines.append("")
    lines.append("- All times are in milliseconds (ms)")
    lines.append("- Speedup > 1.0x means Rust is faster than PyTorch")
    lines.append("- Speedup < 1.0x means PyTorch is faster than Rust")
    lines.append("- '-' indicates missing data for that configuration")
    lines.append("")
    lines.append("## Methodology")
    lines.append("")
    lines.append("- **Python**: Uses `time.perf_counter()` with warmup iterations")
    lines.append("- **Rust**: Uses Criterion.rs with statistical analysis")
    lines.append("- Both implementations use the same LFFN-MLP architecture")
    lines.append("- PDE residual computed using analytical derivatives")
    lines.append("")

    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Generate PINN benchmark comparison report")
    parser.add_argument(
        "--python-cpu-results",
        type=str,
        default=None,
        help="Path to Python CPU benchmark results JSON"
    )
    parser.add_argument(
        "--python-gpu-results",
        type=str,
        default=None,
        help="Path to Python GPU benchmark results JSON"
    )
    parser.add_argument(
        "--python-results",
        type=str,
        default=None,
        help="Path to Python benchmark results JSON (legacy, used as CPU)"
    )
    parser.add_argument(
        "--rust-cpu-output",
        type=str,
        default=None,
        help="Path to Rust CPU Criterion output text"
    )
    parser.add_argument(
        "--rust-output",
        type=str,
        default=None,
        help="Path to Rust Criterion output text (legacy, used as CPU)"
    )
    parser.add_argument(
        "--rust-cuda-output",
        type=str,
        default=None,
        help="Path to Rust CUDA Criterion output text"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="results/comparison_report.md",
        help="Output markdown file"
    )
    args = parser.parse_args()

    # Load results with fallbacks for legacy argument names
    python_cpu_path = Path(args.python_cpu_results) if args.python_cpu_results else (
        Path(args.python_results) if args.python_results else None
    )
    python_gpu_path = Path(args.python_gpu_results) if args.python_gpu_results else None

    rust_cpu_path = Path(args.rust_cpu_output) if args.rust_cpu_output else (
        Path(args.rust_output) if args.rust_output else None
    )
    rust_cuda_path = Path(args.rust_cuda_output) if args.rust_cuda_output else None

    python_cpu_results = load_python_results(python_cpu_path)
    python_gpu_results = load_python_results(python_gpu_path)
    rust_cpu_results = parse_rust_criterion_output(rust_cpu_path)
    rust_cuda_results = parse_rust_criterion_output(rust_cuda_path)

    # Generate report
    report = generate_report(
        python_cpu_results,
        python_gpu_results,
        rust_cpu_results,
        rust_cuda_results
    )

    # Write output
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(report)

    print(f"Report generated: {output_path}")


if __name__ == "__main__":
    main()