#!/usr/bin/env python3
"""
Benchmark adapter generation latency — measures HyperNetwork + ConditionEncoder speed.

Target: <10ms per adapter generation (real-time compatible).

Usage:
    python3 scripts/benchmark_adapter_gen.py --device cuda --n_trials 1000
"""

import argparse
import json
import os
import sys
import time

import torch
import yaml

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from model.hypernetwork import HyperNetwork
from model.condition_encoder import ConditionEncoder
from model.lora import LoRAConfig, LoRAInjector, compute_total_lora_params


def parse_args():
    p = argparse.ArgumentParser(description="Benchmark adapter generation latency")
    p.add_argument("--config", type=str, default="configs/hypernetwork.yaml")
    p.add_argument("--n_trials", type=int, default=1000)
    p.add_argument("--batch_sizes", type=str, default="1,4,8,16", help="Comma-separated batch sizes")
    p.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--output", type=str, default="adapter_gen_benchmark.json")
    return p.parse_args()


@torch.no_grad()
def benchmark(condition_encoder, hypernetwork, device, n_trials, batch_size):
    """Measure latency for generating adapters."""
    # Warmup
    for _ in range(10):
        cam_id = torch.randint(0, 100, (batch_size,), device=device)
        scene = torch.randn(batch_size, 2048, device=device)
        query = torch.randn(batch_size, 2048, device=device)
        cond = condition_encoder(cam_id, scene, query)
        params, sigma = hypernetwork(cond)
    if device == "cuda":
        torch.cuda.synchronize()

    # Benchmark
    latencies = []
    for _ in range(n_trials):
        cam_id = torch.randint(0, 100, (batch_size,), device=device)
        scene = torch.randn(batch_size, 2048, device=device)
        query = torch.randn(batch_size, 2048, device=device)

        if device == "cuda":
            torch.cuda.synchronize()
        start = time.perf_counter()

        cond = condition_encoder(cam_id, scene, query)
        params, sigma = hypernetwork(cond)

        if device == "cuda":
            torch.cuda.synchronize()
        latencies.append((time.perf_counter() - start) * 1000)

    latencies.sort()
    return {
        "batch_size": batch_size,
        "mean_ms": round(sum(latencies) / len(latencies), 4),
        "median_ms": round(latencies[len(latencies) // 2], 4),
        "p95_ms": round(latencies[int(len(latencies) * 0.95)], 4),
        "p99_ms": round(latencies[int(len(latencies) * 0.99)], 4),
        "min_ms": round(min(latencies), 4),
        "max_ms": round(max(latencies), 4),
        "throughput_per_sec": round(1000 / (sum(latencies) / len(latencies)) * batch_size, 1),
    }


def main():
    args = parse_args()
    device = torch.device(args.device)

    # Load config
    if os.path.exists(args.config):
        with open(args.config) as f:
            config = yaml.safe_load(f)
    else:
        config = {}

    # Build components
    lora_config = LoRAConfig(
        rank=config.get("lora", {}).get("rank", 16),
        alpha=config.get("lora", {}).get("alpha", 32.0),
        targets=tuple(config.get("lora", {}).get("targets", ["q", "v"])),
    )

    condition_encoder = ConditionEncoder(
        n_cameras=config.get("condition_encoder", {}).get("n_cameras", 2048),
    ).to(device).eval()

    hypernetwork = HyperNetwork(
        cond_dim=config.get("hypernetwork", {}).get("cond_dim", 256),
        hidden_dim=config.get("hypernetwork", {}).get("hidden_dim", 512),
        lora_config=lora_config,
        num_decoder_blocks=12,
        decoder_embed_dim=1024,
    ).to(device).eval()

    total_lora_params = compute_total_lora_params(12, 1024, lora_config.rank, lora_config.targets)

    print(f"Adapter Generation Benchmark")
    print(f"  Device: {device}")
    print(f"  LoRA config: rank={lora_config.rank}, targets={lora_config.targets}")
    print(f"  Total LoRA params per adapter: {total_lora_params:,}")
    print(f"  HyperNetwork params: {hypernetwork.num_own_params:,}")
    print(f"  ConditionEncoder params: {sum(p.numel() for p in condition_encoder.parameters()):,}")
    print(f"  Trials per batch size: {args.n_trials}")
    print()

    batch_sizes = [int(x) for x in args.batch_sizes.split(",")]
    all_results = []

    for bs in batch_sizes:
        result = benchmark(condition_encoder, hypernetwork, args.device, args.n_trials, bs)
        all_results.append(result)
        target_met = "PASS" if result["mean_ms"] < 10 else "FAIL"
        print(f"  BS={bs:2d}: mean={result['mean_ms']:.3f}ms  "
              f"p99={result['p99_ms']:.3f}ms  "
              f"throughput={result['throughput_per_sec']:.0f}/s  "
              f"[{target_met}]")

    # Save
    output = {
        "device": str(device),
        "lora_params_per_adapter": total_lora_params,
        "hypernetwork_params": hypernetwork.num_own_params,
        "n_trials": args.n_trials,
        "results": all_results,
    }
    with open(args.output, "w") as f:
        json.dump(output, f, indent=2)
    print(f"\nResults saved to {args.output}")


if __name__ == "__main__":
    main()