from __future__ import annotations

import argparse
import json
import time

import torch

from .config import apply_overrides, load_config
from .model import build_model
from .utils import trainable_parameter_count


def main() -> None:
    parser = argparse.ArgumentParser(description="Profile ObjectModel-v1 parameters and latency")
    parser.add_argument("--config", default="configs/objectmodel_v1.yaml")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--runs", type=int, default=50)
    parser.add_argument("--set", action="append", default=[])
    args = parser.parse_args()
    config = apply_overrides(load_config(args.config), args.set)
    device = torch.device(args.device)
    model = build_model(config).eval().to(device)
    size = model.spec.input_size
    sample = torch.randn(1, 3, size, size, device=device)
    with torch.inference_mode():
        for _ in range(args.warmup):
            model(sample)
        if device.type == "cuda":
            torch.cuda.synchronize()
        durations = []
        for _ in range(args.runs):
            start = time.perf_counter()
            model(sample)
            if device.type == "cuda":
                torch.cuda.synchronize()
            durations.append((time.perf_counter() - start) * 1000)
    durations.sort()
    parameters = trainable_parameter_count(model)
    report = {
        "parameters": parameters,
        "parameters_millions": round(parameters / 1e6, 3),
        "fp32_weight_megabytes": round(parameters * 4 / 1024**2, 2),
        "input_size": size,
        "device": str(device),
        "latency_ms_median": round(durations[len(durations) // 2], 3),
        "latency_ms_p95": round(durations[min(int(len(durations) * 0.95), len(durations) - 1)], 3),
    }
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()