#!/usr/bin/env python3
"""
WiFi-DensePose Model Benchmarking

Loads trained ONNX models, runs inference on test data, and reports
performance metrics: latency, throughput, PCK@0.2, model size, and
estimated FLOPs.

Can compare multiple models from a hyperparameter sweep.

Usage:
    # Benchmark a single model
    python scripts/benchmark-model.py --model checkpoints/best.onnx

    # Benchmark with recorded test data
    python scripts/benchmark-model.py --model best.onnx --test-data data/recordings/test.csi.jsonl

    # Compare models from a sweep
    python scripts/benchmark-model.py --sweep-dir training-results/wdp-train-a100-*/checkpoints/

    # Benchmark with synthetic data (no recordings needed)
    python scripts/benchmark-model.py --model best.onnx --synthetic --num-samples 200

    # Export results as JSON
    python scripts/benchmark-model.py --model best.onnx --output results.json

Prerequisites:
    pip install onnxruntime numpy
    Optional: pip install onnx  (for FLOPs estimation)
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional

import numpy as np

try:
    import onnxruntime as ort
except ImportError:
    print("ERROR: onnxruntime not installed. Run: pip install onnxruntime")
    sys.exit(1)


# ── Configuration ────────────────────────────────────────────────────────────

# Default model input shape (must match TrainingConfig defaults)
NUM_SUBCARRIERS = 56
NUM_ANTENNAS_TX = 3
NUM_ANTENNAS_RX = 3
WINDOW_FRAMES = 100
NUM_KEYPOINTS = 17
HEATMAP_SIZE = 56

# PCK threshold
PCK_THRESHOLD = 0.2


# ── Data classes ─────────────────────────────────────────────────────────────

@dataclass
class BenchmarkResult:
    model_path: str
    model_size_mb: float
    num_parameters: Optional[int] = None
    estimated_flops: Optional[int] = None

    # Latency
    warmup_runs: int = 10
    benchmark_runs: int = 100
    latency_mean_ms: float = 0.0
    latency_std_ms: float = 0.0
    latency_p50_ms: float = 0.0
    latency_p95_ms: float = 0.0
    latency_p99_ms: float = 0.0
    throughput_fps: float = 0.0

    # Accuracy (if ground truth available)
    pck_at_02: Optional[float] = None
    mean_per_joint_error: Optional[float] = None
    num_test_samples: int = 0

    # Input shape
    input_shape: list = field(default_factory=list)
    provider: str = ""


# ── ONNX model loading ──────────────────────────────────────────────────────

def load_model(model_path: str) -> ort.InferenceSession:
    """Load an ONNX model with the best available execution provider."""
    providers = []
    if "CUDAExecutionProvider" in ort.get_available_providers():
        providers.append("CUDAExecutionProvider")
    providers.append("CPUExecutionProvider")

    sess_opts = ort.SessionOptions()
    sess_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    sess_opts.intra_op_num_threads = os.cpu_count() or 4

    session = ort.InferenceSession(model_path, sess_opts, providers=providers)
    return session


def get_model_info(model_path: str) -> dict:
    """Extract model metadata: size, parameter count, FLOPs estimate."""
    path = Path(model_path)
    size_mb = path.stat().st_size / (1024 * 1024)

    info = {
        "size_mb": round(size_mb, 2),
        "num_parameters": None,
        "estimated_flops": None,
    }

    # Try to count parameters via onnx
    try:
        import onnx
        model = onnx.load(model_path)
        total_params = 0
        for initializer in model.graph.initializer:
            shape = list(initializer.dims)
            if shape:
                total_params += int(np.prod(shape))
        info["num_parameters"] = total_params

        # Rough FLOPs estimate: ~2 * params (multiply-accumulate)
        info["estimated_flops"] = total_params * 2
    except ImportError:
        pass
    except Exception as e:
        print(f"  Warning: Could not extract parameter count: {e}")

    return info


# ── Synthetic data generation ────────────────────────────────────────────────

def generate_synthetic_input(
    batch_size: int = 1,
    num_subcarriers: int = NUM_SUBCARRIERS,
    num_tx: int = NUM_ANTENNAS_TX,
    num_rx: int = NUM_ANTENNAS_RX,
    window_frames: int = WINDOW_FRAMES,
) -> np.ndarray:
    """Generate synthetic CSI input tensor matching the model's expected shape.

    The WiFi-DensePose model expects input shape:
      [batch, channels, height, width]
    where channels = num_tx * num_rx, height = window_frames, width = num_subcarriers.
    """
    channels = num_tx * num_rx  # 3x3 = 9 MIMO streams
    # Simulate CSI amplitude data with realistic distribution
    rng = np.random.default_rng(42)
    data = rng.normal(loc=0.0, scale=1.0, size=(batch_size, channels, window_frames, num_subcarriers))
    return data.astype(np.float32)


def generate_synthetic_keypoints(
    num_samples: int,
    num_keypoints: int = NUM_KEYPOINTS,
    heatmap_size: int = HEATMAP_SIZE,
) -> np.ndarray:
    """Generate synthetic ground truth keypoint coordinates for PCK evaluation."""
    rng = np.random.default_rng(123)
    # Keypoints as (x, y) in [0, heatmap_size) range
    return rng.uniform(0, heatmap_size, size=(num_samples, num_keypoints, 2)).astype(np.float32)


# ── Load test data from .csi.jsonl ──────────────────────────────────────────

def load_test_data(
    jsonl_path: str,
    window_frames: int = WINDOW_FRAMES,
    num_subcarriers: int = NUM_SUBCARRIERS,
    max_samples: int = 500,
) -> np.ndarray:
    """Load CSI frames from a .csi.jsonl file and window them into model inputs."""
    frames = []
    path = Path(jsonl_path)

    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
                subs = record.get("subcarriers", [])
                if len(subs) > 0:
                    frames.append(subs)
            except json.JSONDecodeError:
                continue

    if len(frames) < window_frames:
        print(f"  Warning: Only {len(frames)} frames, need {window_frames}. Padding with zeros.")
        while len(frames) < window_frames:
            frames.append([0.0] * num_subcarriers)

    # Normalize subcarrier count
    normalized = []
    for frame in frames:
        if len(frame) < num_subcarriers:
            frame = frame + [0.0] * (num_subcarriers - len(frame))
        elif len(frame) > num_subcarriers:
            # Downsample via linear interpolation
            indices = np.linspace(0, len(frame) - 1, num_subcarriers)
            frame = np.interp(indices, range(len(frame)), frame).tolist()
        normalized.append(frame)

    frames = normalized

    # Create sliding windows
    samples = []
    stride = max(1, window_frames // 2)
    for i in range(0, len(frames) - window_frames + 1, stride):
        window = frames[i : i + window_frames]
        # Shape: [channels=1, window_frames, num_subcarriers]
        # Expand single stream to 9 channels (repeat for MIMO)
        arr = np.array(window, dtype=np.float32)
        arr = np.expand_dims(arr, axis=0)  # [1, window_frames, num_subcarriers]
        arr = np.repeat(arr, NUM_ANTENNAS_TX * NUM_ANTENNAS_RX, axis=0)  # [9, window, subs]
        samples.append(arr)

        if len(samples) >= max_samples:
            break

    if not samples:
        return generate_synthetic_input(1)

    return np.stack(samples, axis=0)  # [N, 9, window_frames, num_subcarriers]


# ── Benchmarking ─────────────────────────────────────────────────────────────

def benchmark_latency(
    session: ort.InferenceSession,
    input_data: np.ndarray,
    warmup: int = 10,
    runs: int = 100,
) -> dict:
    """Measure inference latency over multiple runs."""
    input_name = session.get_inputs()[0].name

    # Warmup
    for _ in range(warmup):
        session.run(None, {input_name: input_data[:1]})

    # Timed runs
    latencies = []
    for _ in range(runs):
        start = time.perf_counter()
        session.run(None, {input_name: input_data[:1]})
        end = time.perf_counter()
        latencies.append((end - start) * 1000)  # ms

    latencies = np.array(latencies)
    return {
        "mean_ms": float(np.mean(latencies)),
        "std_ms": float(np.std(latencies)),
        "p50_ms": float(np.percentile(latencies, 50)),
        "p95_ms": float(np.percentile(latencies, 95)),
        "p99_ms": float(np.percentile(latencies, 99)),
        "throughput_fps": 1000.0 / float(np.mean(latencies)),
    }


def compute_pck(
    predictions: np.ndarray,
    ground_truth: np.ndarray,
    threshold: float = PCK_THRESHOLD,
    normalize_by: float = HEATMAP_SIZE,
) -> float:
    """Compute Percentage of Correct Keypoints at a given threshold.

    PCK@t = fraction of predicted keypoints within t * normalize_by of ground truth.
    """
    if predictions.shape != ground_truth.shape:
        return 0.0

    # Euclidean distance per keypoint
    distances = np.linalg.norm(predictions - ground_truth, axis=-1)  # [N, K]
    threshold_pixels = threshold * normalize_by
    correct = (distances < threshold_pixels).astype(float)
    return float(np.mean(correct))


def extract_keypoints_from_heatmaps(heatmaps: np.ndarray) -> np.ndarray:
    """Convert heatmap outputs [N, K, H, W] to keypoint coordinates [N, K, 2]."""
    n, k, h, w = heatmaps.shape
    flat = heatmaps.reshape(n, k, -1)
    max_idx = np.argmax(flat, axis=-1)  # [N, K]
    y = max_idx // w
    x = max_idx % w
    return np.stack([x, y], axis=-1).astype(np.float32)


def benchmark_model(
    model_path: str,
    test_data: Optional[np.ndarray] = None,
    gt_keypoints: Optional[np.ndarray] = None,
    warmup: int = 10,
    runs: int = 100,
) -> BenchmarkResult:
    """Run full benchmark on a single model."""
    print(f"\nBenchmarking: {model_path}")

    # Load model
    session = load_model(model_path)
    provider = session.get_providers()[0]
    print(f"  Provider: {provider}")

    # Model info
    model_info = get_model_info(model_path)
    print(f"  Size: {model_info['size_mb']} MB")
    if model_info["num_parameters"]:
        print(f"  Parameters: {model_info['num_parameters']:,}")
    if model_info["estimated_flops"]:
        print(f"  Estimated FLOPs: {model_info['estimated_flops']:,}")

    # Input shape
    input_meta = session.get_inputs()[0]
    input_shape = input_meta.shape
    print(f"  Input: {input_meta.name} {input_shape} ({input_meta.type})")

    # Output shapes
    for out in session.get_outputs():
        print(f"  Output: {out.name} {out.shape}")

    # Generate or use provided test data
    if test_data is None:
        # Infer shape from model
        if input_shape and all(isinstance(d, int) for d in input_shape):
            batch = max(1, input_shape[0] if input_shape[0] > 0 else 1)
            test_data = np.random.randn(*[batch if d <= 0 else d for d in input_shape]).astype(np.float32)
        else:
            test_data = generate_synthetic_input(1)

    # Latency benchmark
    print(f"  Running {warmup} warmup + {runs} benchmark iterations...")
    latency = benchmark_latency(session, test_data, warmup=warmup, runs=runs)
    print(f"  Latency: {latency['mean_ms']:.2f} +/- {latency['std_ms']:.2f} ms")
    print(f"  P50/P95/P99: {latency['p50_ms']:.2f} / {latency['p95_ms']:.2f} / {latency['p99_ms']:.2f} ms")
    print(f"  Throughput: {latency['throughput_fps']:.1f} fps")

    # Accuracy (if ground truth provided or we can do synthetic evaluation)
    pck = None
    mpjpe = None
    num_samples = 0

    if gt_keypoints is not None and test_data is not None:
        input_name = session.get_inputs()[0].name
        all_preds = []

        for i in range(len(test_data)):
            outputs = session.run(None, {input_name: test_data[i : i + 1]})
            # Assume first output is keypoint heatmaps [1, K, H, W]
            heatmaps = outputs[0]
            if heatmaps.ndim == 4:
                kp = extract_keypoints_from_heatmaps(heatmaps)
                all_preds.append(kp[0])

        if all_preds:
            predictions = np.stack(all_preds)
            gt = gt_keypoints[: len(predictions)]
            pck = compute_pck(predictions, gt)
            distances = np.linalg.norm(predictions - gt, axis=-1)
            mpjpe = float(np.mean(distances))
            num_samples = len(predictions)
            print(f"  PCK@{PCK_THRESHOLD}: {pck:.4f}")
            print(f"  MPJPE: {mpjpe:.2f} px")
            print(f"  Samples: {num_samples}")

    result = BenchmarkResult(
        model_path=model_path,
        model_size_mb=model_info["size_mb"],
        num_parameters=model_info["num_parameters"],
        estimated_flops=model_info["estimated_flops"],
        warmup_runs=warmup,
        benchmark_runs=runs,
        latency_mean_ms=round(latency["mean_ms"], 3),
        latency_std_ms=round(latency["std_ms"], 3),
        latency_p50_ms=round(latency["p50_ms"], 3),
        latency_p95_ms=round(latency["p95_ms"], 3),
        latency_p99_ms=round(latency["p99_ms"], 3),
        throughput_fps=round(latency["throughput_fps"], 1),
        pck_at_02=round(pck, 4) if pck is not None else None,
        mean_per_joint_error=round(mpjpe, 2) if mpjpe is not None else None,
        num_test_samples=num_samples,
        input_shape=list(input_shape) if input_shape else [],
        provider=provider,
    )

    return result


# ── Comparison table ─────────────────────────────────────────────────────────

def print_comparison_table(results: list[BenchmarkResult]):
    """Print a formatted comparison table of multiple models."""
    if not results:
        return

    print("\n" + "=" * 100)
    print("  Model Comparison")
    print("=" * 100)

    # Header
    print(
        f"{'Model':<35} {'Size(MB)':>8} {'Params':>10} "
        f"{'Lat(ms)':>8} {'P95(ms)':>8} {'FPS':>7} {'PCK@0.2':>8}"
    )
    print("-" * 100)

    for r in results:
        name = Path(r.model_path).stem[:33]
        params = f"{r.num_parameters:,}" if r.num_parameters else "?"
        pck = f"{r.pck_at_02:.4f}" if r.pck_at_02 is not None else "N/A"

        print(
            f"{name:<35} {r.model_size_mb:>8.2f} {params:>10} "
            f"{r.latency_mean_ms:>8.2f} {r.latency_p95_ms:>8.2f} "
            f"{r.throughput_fps:>7.1f} {pck:>8}"
        )

    print("=" * 100)

    # Best model by latency
    best_latency = min(results, key=lambda r: r.latency_mean_ms)
    print(f"\n  Fastest: {Path(best_latency.model_path).stem} ({best_latency.latency_mean_ms:.2f} ms)")

    # Best by PCK (if available)
    pck_results = [r for r in results if r.pck_at_02 is not None]
    if pck_results:
        best_pck = max(pck_results, key=lambda r: r.pck_at_02)
        print(f"  Best accuracy: {Path(best_pck.model_path).stem} (PCK@0.2={best_pck.pck_at_02:.4f})")

    # Smallest model
    smallest = min(results, key=lambda r: r.model_size_mb)
    print(f"  Smallest: {Path(smallest.model_path).stem} ({smallest.model_size_mb:.2f} MB)")


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Benchmark WiFi-DensePose ONNX models",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument("--model", type=str, help="Path to a single ONNX model")
    parser.add_argument("--sweep-dir", type=str, help="Directory containing multiple ONNX models to compare")
    parser.add_argument("--test-data", type=str, help="Path to .csi.jsonl test data file")
    parser.add_argument("--synthetic", action="store_true", help="Use synthetic test data")
    parser.add_argument("--num-samples", type=int, default=100, help="Number of synthetic samples (default: 100)")
    parser.add_argument("--warmup", type=int, default=10, help="Warmup iterations (default: 10)")
    parser.add_argument("--runs", type=int, default=100, help="Benchmark iterations (default: 100)")
    parser.add_argument("--output", type=str, help="Save results to JSON file")
    parser.add_argument("--gpu", action="store_true", help="Force GPU execution provider")

    args = parser.parse_args()

    if not args.model and not args.sweep_dir:
        parser.error("Specify --model or --sweep-dir")

    # Prepare test data
    test_data = None
    gt_keypoints = None

    if args.test_data:
        print(f"Loading test data from: {args.test_data}")
        test_data = load_test_data(args.test_data)
        print(f"  Loaded {len(test_data)} windowed samples")
    elif args.synthetic:
        print(f"Generating {args.num_samples} synthetic samples...")
        test_data = generate_synthetic_input(args.num_samples)
        gt_keypoints = generate_synthetic_keypoints(args.num_samples)
        print(f"  Input shape: {test_data.shape}")

    # Collect models
    model_paths = []
    if args.model:
        model_paths.append(args.model)
    if args.sweep_dir:
        sweep = Path(args.sweep_dir)
        if sweep.is_dir():
            model_paths.extend(sorted(str(p) for p in sweep.glob("**/*.onnx")))
        else:
            # Glob pattern
            from glob import glob
            model_paths.extend(sorted(glob(str(sweep))))

    if not model_paths:
        print("ERROR: No ONNX models found.")
        sys.exit(1)

    print(f"Found {len(model_paths)} model(s) to benchmark.")

    # Benchmark each model
    results = []
    for path in model_paths:
        if not Path(path).exists():
            print(f"  Skipping (not found): {path}")
            continue
        try:
            result = benchmark_model(
                path,
                test_data=test_data,
                gt_keypoints=gt_keypoints,
                warmup=args.warmup,
                runs=args.runs,
            )
            results.append(result)
        except Exception as e:
            print(f"  ERROR benchmarking {path}: {e}")

    # Comparison table
    if len(results) > 1:
        print_comparison_table(results)

    # Save results
    if args.output:
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "w") as f:
            json.dump(
                {
                    "benchmark_results": [asdict(r) for r in results],
                    "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                    "num_models": len(results),
                },
                f,
                indent=2,
            )
        print(f"\nResults saved to: {output_path}")

    if not results:
        print("No models were successfully benchmarked.")
        sys.exit(1)


if __name__ == "__main__":
    main()