#!/usr/bin/env python3
"""
PyTorch MPS Mixture of Experts (MoE) Benchmark - Comparison baseline for RustyTorch.

This script benchmarks MoE routing and expert computation on Apple MPS backend
and outputs JSON results for automated comparison with RustyTorch MoE.

Usage:
    python3 benchmarks/metal/bench_moe_mps.py
    python3 benchmarks/metal/bench_moe_mps.py --json > benchmarks/reports/pytorch_moe.json

Run Rust benchmark with:
    cargo bench -p rtx-transformers --bench metal_moe_bench
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import time
import sys
import json
import argparse
import platform
import statistics
from dataclasses import dataclass, asdict
from typing import List, Optional


@dataclass
class MoEBenchmarkResult:
    """Single MoE benchmark result."""
    name: str
    batch_size: int
    seq_len: int
    hidden_size: int
    num_experts: int
    top_k: int
    routing_time_ms: float
    expert_compute_time_ms: float
    total_time_ms: float
    std_time_ms: float
    throughput_tokens_per_sec: float
    expert_utilization: float
    memory_mb: Optional[float] = None


@dataclass
class MoEBenchmarkReport:
    """Complete MoE benchmark report."""
    framework: str
    backend: str
    pytorch_version: str
    python_version: str
    macos_version: str
    chip: str
    timestamp: str
    iterations: int
    warmup: int
    results: List[dict]


class TopKRouter(nn.Module):
    """Top-K expert router with softmax gating."""

    def __init__(self, hidden_size: int, num_experts: int, top_k: int = 2):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        self.gate = nn.Linear(hidden_size, num_experts, bias=False)

    def forward(self, x: torch.Tensor):
        """
        Route tokens to experts.

        Args:
            x: Input tensor [batch, seq, hidden]

        Returns:
            expert_indices: [batch, seq, top_k]
            expert_weights: [batch, seq, top_k]
        """
        # Compute gating scores
        scores = self.gate(x)  # [batch, seq, num_experts]
        probs = F.softmax(scores, dim=-1)

        # Select top-k experts
        weights, indices = torch.topk(probs, self.top_k, dim=-1)
        weights = weights / weights.sum(dim=-1, keepdim=True)  # Renormalize

        return indices, weights


class ExpertFFN(nn.Module):
    """Single expert FFN."""

    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class SimpleMoE(nn.Module):
    """Simple Mixture of Experts layer for benchmarking."""

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        num_experts: int,
        top_k: int = 2,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.router = TopKRouter(hidden_size, num_experts, top_k)
        self.experts = nn.ModuleList([
            ExpertFFN(hidden_size, intermediate_size) for _ in range(num_experts)
        ])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        MoE forward pass.

        Args:
            x: Input tensor [batch, seq, hidden]

        Returns:
            Output tensor [batch, seq, hidden]
        """
        batch_size, seq_len, hidden_size = x.shape

        # Get routing decisions
        expert_indices, expert_weights = self.router(x)

        # Compute expert outputs (simplified - not optimized for production)
        output = torch.zeros_like(x)

        for k in range(self.top_k):
            for e in range(self.num_experts):
                # Create mask for tokens routed to this expert
                mask = expert_indices[:, :, k] == e  # [batch, seq]
                if mask.any():
                    expert_input = x[mask]  # [num_tokens, hidden]
                    expert_output = self.experts[e](expert_input)
                    weight = expert_weights[:, :, k][mask].unsqueeze(-1)
                    output[mask] += weight * expert_output

        return output


def check_mps() -> bool:
    """Check if MPS backend is available."""
    if not torch.backends.mps.is_available():
        print("ERROR: MPS not available on this system", file=sys.stderr)
        return False
    return True


def benchmark_moe(
    batch_size: int,
    seq_len: int,
    hidden_size: int,
    intermediate_size: int,
    num_experts: int,
    top_k: int = 2,
    iterations: int = 500,
    warmup: int = 50,
) -> MoEBenchmarkResult:
    """Run MoE benchmark."""
    device = torch.device("mps")

    # Create model
    model = SimpleMoE(
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        num_experts=num_experts,
        top_k=top_k,
    ).to(device)
    model.eval()

    # Create input
    x = torch.randn(batch_size, seq_len, hidden_size, device=device)

    # Warmup
    with torch.no_grad():
        for _ in range(warmup):
            _ = model(x)
        torch.mps.synchronize()

    # Benchmark
    total_times = []
    routing_times = []
    expert_times = []

    with torch.no_grad():
        for _ in range(iterations):
            # Time routing
            start_routing = time.perf_counter()
            expert_indices, expert_weights = model.router(x)
            torch.mps.synchronize()
            end_routing = time.perf_counter()

            # Time expert computation
            start_expert = time.perf_counter()
            _ = model(x)
            torch.mps.synchronize()
            end_expert = time.perf_counter()

            routing_times.append((end_routing - start_routing) * 1000)
            expert_times.append((end_expert - start_expert) * 1000 - routing_times[-1])
            total_times.append((end_expert - start_routing) * 1000)

    # Calculate statistics
    avg_routing = statistics.mean(routing_times)
    avg_expert = statistics.mean(expert_times)
    avg_total = statistics.mean(total_times)
    std_total = statistics.stdev(total_times) if len(total_times) > 1 else 0.0

    # Calculate throughput
    total_tokens = batch_size * seq_len
    throughput = total_tokens / (avg_total / 1000)

    # Estimate expert utilization (uniform ideal = 100%)
    expert_utilization = 100.0  # Simplified - would need actual load tracking

    # Memory estimate
    memory_mb = None
    try:
        params = sum(p.numel() * p.element_size() for p in model.parameters())
        activations = x.numel() * x.element_size()
        memory_mb = (params + activations) / (1024 * 1024)
    except Exception:
        pass

    name = f"MoE_E{num_experts}_K{top_k}_BS{batch_size}"

    return MoEBenchmarkResult(
        name=name,
        batch_size=batch_size,
        seq_len=seq_len,
        hidden_size=hidden_size,
        num_experts=num_experts,
        top_k=top_k,
        routing_time_ms=avg_routing,
        expert_compute_time_ms=avg_expert,
        total_time_ms=avg_total,
        std_time_ms=std_total,
        throughput_tokens_per_sec=throughput,
        expert_utilization=expert_utilization,
        memory_mb=memory_mb,
    )


def run_benchmark_suite(iterations: int = 500, warmup: int = 50) -> MoEBenchmarkReport:
    """Run the complete MoE benchmark suite."""
    import datetime

    results = []

    # MoE configurations
    # (batch_size, seq_len, hidden_size, intermediate_size, num_experts, top_k)
    scenarios = [
        # Small model (GPT-2 like)
        (1, 128, 768, 3072, 4, 2),
        (8, 128, 768, 3072, 4, 2),
        (32, 128, 768, 3072, 4, 2),

        # Medium model with more experts
        (1, 128, 1024, 4096, 8, 2),
        (8, 128, 1024, 4096, 8, 2),
        (16, 128, 1024, 4096, 8, 2),

        # Large model (LLaMA-like)
        (1, 128, 2048, 5504, 8, 2),
        (4, 128, 2048, 5504, 8, 2),

        # DeepSeek-style (many experts)
        (1, 128, 1024, 2816, 16, 2),
        (4, 128, 1024, 2816, 16, 2),
    ]

    for batch_size, seq_len, hidden_size, intermediate_size, num_experts, top_k in scenarios:
        try:
            result = benchmark_moe(
                batch_size=batch_size,
                seq_len=seq_len,
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                num_experts=num_experts,
                top_k=top_k,
                iterations=iterations,
                warmup=warmup,
            )
            results.append(asdict(result))
        except Exception as e:
            print(f"Warning: Benchmark failed: {e}", file=sys.stderr)

    return MoEBenchmarkReport(
        framework="PyTorch",
        backend="MPS",
        pytorch_version=torch.__version__,
        python_version=platform.python_version(),
        macos_version=platform.mac_ver()[0],
        chip=platform.processor() or "Apple Silicon",
        timestamp=datetime.datetime.now().isoformat(),
        iterations=iterations,
        warmup=warmup,
        results=results,
    )


def print_table(report: MoEBenchmarkReport):
    """Print results as formatted table."""
    print("=" * 100)
    print("PyTorch MPS Mixture of Experts (MoE) Benchmark")
    print("=" * 100)
    print(f"PyTorch version:  {report.pytorch_version}")
    print(f"Backend:          {report.backend}")
    print(f"macOS version:    {report.macos_version}")
    print(f"Chip:             {report.chip}")
    print(f"Iterations:       {report.iterations}")
    print()

    print("-" * 100)
    print(f"| {'Scenario':<25} | {'Routing':>10} | {'Expert':>10} | {'Total':>10} | {'Throughput':>18} |")
    print(f"| {'':<25} | {'(ms)':>10} | {'(ms)':>10} | {'(ms)':>10} | {'(tokens/s)':>18} |")
    print("-" * 100)

    for r in report.results:
        throughput_str = f"{r['throughput_tokens_per_sec']:.0f}"
        print(f"| {r['name']:<25} | {r['routing_time_ms']:>10.4f} | {r['expert_compute_time_ms']:>10.4f} | {r['total_time_ms']:>10.4f} | {throughput_str:>18} |")

    print("-" * 100)
    print()
    print("To compare with RustyTorch MoE:")
    print("  cargo bench -p rtx-transformers --bench metal_moe_bench")


def main():
    parser = argparse.ArgumentParser(description="PyTorch MPS MoE Benchmark")
    parser.add_argument("--json", action="store_true", help="Output JSON format")
    parser.add_argument("--iterations", type=int, default=500, help="Number of iterations")
    parser.add_argument("--warmup", type=int, default=50, help="Number of warmup iterations")
    args = parser.parse_args()

    if not check_mps():
        sys.exit(1)

    report = run_benchmark_suite(iterations=args.iterations, warmup=args.warmup)

    if args.json:
        print(json.dumps(asdict(report), indent=2))
    else:
        print_table(report)


if __name__ == "__main__":
    main()