from __future__ import annotations

import argparse
import json
import statistics
import time
from pathlib import Path

import torch

from train import SpectrogramDataset, build_model, read_labels


def model_parameter_summary(model: torch.nn.Module) -> dict[str, float | int]:
    total_params = sum(param.numel() for param in model.parameters())
    trainable_params = sum(param.numel() for param in model.parameters() if param.requires_grad)
    param_bytes = sum(param.numel() * param.element_size() for param in model.parameters())
    return {
        "total_parameters": total_params,
        "trainable_parameters": trainable_params,
        "parameter_size_mb_float32": round(param_bytes / 1024**2, 3),
    }


def benchmark_latency(
    model: torch.nn.Module,
    sample: torch.Tensor,
    device: torch.device,
    warmup: int,
    repeats: int,
) -> dict[str, float | int | str]:
    use_cuda = device.type == "cuda"
    sample = sample.unsqueeze(0).to(device)
    if use_cuda:
        sample = sample.contiguous(memory_format=torch.channels_last)

    model.eval()
    with torch.inference_mode():
        for _ in range(warmup):
            with torch.amp.autocast("cuda", enabled=use_cuda):
                _ = model(sample)
        if use_cuda:
            torch.cuda.synchronize()

        timings_ms = []
        for _ in range(repeats):
            start = time.perf_counter()
            with torch.amp.autocast("cuda", enabled=use_cuda):
                _ = model(sample)
            if use_cuda:
                torch.cuda.synchronize()
            timings_ms.append((time.perf_counter() - start) * 1000.0)

    timings_ms = sorted(timings_ms)
    return {
        "device": torch.cuda.get_device_name(0) if use_cuda else "cpu",
        "input_shape": list(sample.shape),
        "warmup_runs": warmup,
        "measured_runs": repeats,
        "batch_size": int(sample.shape[0]),
        "mean_latency_ms": round(statistics.fmean(timings_ms), 4),
        "median_latency_ms": round(statistics.median(timings_ms), 4),
        "p95_latency_ms": round(timings_ms[int(0.95 * (len(timings_ms) - 1))], 4),
        "throughput_samples_per_second": round(1000.0 / statistics.fmean(timings_ms), 2),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect model metrics for an RFUAV classifier.")
    parser.add_argument("--processed-dir", default="/data/RFUAV_processed")
    parser.add_argument("--results-dir", default="/data/results")
    parser.add_argument("--checkpoint-dir", default="/data/checkpoints")
    parser.add_argument("--model", choices=["small_cnn", "resnet18", "efficientnet_b0"], default="resnet18")
    parser.add_argument("--warmup", type=int, default=30)
    parser.add_argument("--repeats", type=int, default=200)
    args = parser.parse_args()

    processed_dir = Path(args.processed_dir)
    results_dir = Path(args.results_dir) / args.model
    checkpoint_path = Path(args.checkpoint_dir) / args.model / "best_model.pt"
    metrics_path = results_dir / "metrics.json"
    out_path = results_dir / "model_summary.json"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    labels = read_labels(processed_dir)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = build_model(args.model, len(labels)).to(device)
    if device.type == "cuda":
        model = model.to(memory_format=torch.channels_last)
        torch.backends.cudnn.benchmark = True

    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint["model"])

    test_ds = SpectrogramDataset(processed_dir, "test", preload=False)
    sample, _ = test_ds[0]

    accuracy_metrics = json.loads(metrics_path.read_text()) if metrics_path.exists() else {}
    checkpoint_size_mb = round(checkpoint_path.stat().st_size / 1024**2, 3)

    summary = {
        "model": args.model,
        "classes_trained": labels,
        "num_classes": len(labels),
        "accuracy": {
            "best_validation_accuracy": accuracy_metrics.get("best_val_accuracy"),
            "test_accuracy": accuracy_metrics.get("test_accuracy"),
        },
        "model_size": {
            **model_parameter_summary(model),
            "checkpoint_size_mb": checkpoint_size_mb,
        },
        "inference_latency": benchmark_latency(model, sample, device, args.warmup, args.repeats),
        "source_files": {
            "checkpoint": str(checkpoint_path),
            "metrics": str(metrics_path),
            "processed_manifest": str(processed_dir / "manifest.csv"),
        },
    }

    out_path.write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2))
    print(f"Saved model summary: {out_path}")


if __name__ == "__main__":
    main()