"""Verify SOFIA v7.1 model scales: param count, memory footprint, forward pass.

Run:
    python -m sofia_v71.verify
    python -m sofia_v71.verify --preset L
    python -m sofia_v71.verify --preset M --benchmark
"""

from __future__ import annotations

import argparse
import time
from typing import Optional

import torch

from .blocks import is_mamba_ssm_available, is_mamba2_available
from .config import sofia_m_config, sofia_l_config, sofia_tiny_config
from .model import SOFIAv71


def count_parameters(model: torch.nn.Module) -> dict:
    """Count parameters per named child."""
    counts = {}
    for name, param in model.named_parameters():
        top = name.split(".")[0]
        counts[top] = counts.get(top, 0) + param.numel()
    counts["_total"] = sum(counts.values())
    return counts


def estimate_quantized_size(n_params: int, precision: str = "int8") -> float:
    """Estimate on-disk / VRAM weight storage in MB."""
    bytes_per_param = {"fp32": 4, "fp16": 2, "int8": 1, "int4": 0.5}[precision]
    return n_params * bytes_per_param / (1024 ** 2)


def test_forward(model: SOFIAv71, device: str = "cpu", batch_size: int = 1) -> dict:
    """Dry forward pass, return output shapes."""
    model = model.to(device).eval()
    cfg = model.cfg
    sat = torch.randn(batch_size, 3, cfg.input_size, cfg.input_size, device=device)
    uav = torch.randn(batch_size, 3, cfg.input_size, cfg.input_size, device=device)
    altitude = torch.rand(batch_size, device=device) * 300.0 + 50.0

    shapes = {}
    with torch.no_grad():
        out = model(sat=sat, uav=uav, altitude=altitude)
    for k, v in out.items():
        if isinstance(v, torch.Tensor):
            shapes[k] = tuple(v.shape)
        elif isinstance(v, list):
            shapes[k] = f"list[{len(v)} × {tuple(v[0].shape)}]"
        elif isinstance(v, dict):
            for kk, vv in v.items():
                if isinstance(vv, torch.Tensor):
                    shapes[f"{k}.{kk}"] = tuple(vv.shape)
    return shapes


def benchmark(model: SOFIAv71, device: str = "cpu", n_warmup: int = 3, n_runs: int = 20) -> dict:
    """Measure forward pass latency."""
    model = model.to(device).eval()
    cfg = model.cfg
    sat = torch.randn(1, 3, cfg.input_size, cfg.input_size, device=device)
    uav = torch.randn(1, 3, cfg.input_size, cfg.input_size, device=device)
    altitude = torch.tensor([150.0], device=device)

    # Warmup
    with torch.no_grad():
        for _ in range(n_warmup):
            _ = model(sat=sat, uav=uav, altitude=altitude)
        if device.startswith("cuda"):
            torch.cuda.synchronize()

        # Measured runs
        times = []
        for _ in range(n_runs):
            if device.startswith("cuda"):
                torch.cuda.synchronize()
            t0 = time.perf_counter()
            _ = model(sat=sat, uav=uav, altitude=altitude)
            if device.startswith("cuda"):
                torch.cuda.synchronize()
            times.append(time.perf_counter() - t0)

    times_ms = [t * 1000 for t in times]
    return {
        "mean_ms": sum(times_ms) / len(times_ms),
        "min_ms": min(times_ms),
        "max_ms": max(times_ms),
        "std_ms": (sum((t - sum(times_ms) / len(times_ms)) ** 2 for t in times_ms) / len(times_ms)) ** 0.5,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Verify SOFIA v7.1 model")
    parser.add_argument("--preset", choices=["Tiny", "M", "L"], default="Tiny",
                        help="Model preset (default: Tiny for fast smoke test)")
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--benchmark", action="store_true")
    parser.add_argument("--batch-size", type=int, default=1)
    args = parser.parse_args()

    cfg_fn = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config}[args.preset]
    cfg = cfg_fn()

    # Report mamba_ssm availability
    print(f"\n=== Environment ===")
    print(f"  mamba_ssm v1 (Mamba-1 selective scan): "
          f"{'available' if is_mamba_ssm_available() else 'NOT available'}")
    print(f"  mamba_ssm Mamba-2 (SSD dual form):     "
          f"{'available' if is_mamba2_available() else 'NOT available'}")
    print(f"  configured variant: {cfg.mamba_variant}, backend: {cfg.mamba_backend}")

    print(f"\n=== SOFIA-{args.preset} configuration ===")
    print(cfg.summary())

    print("\n=== Building model ===")
    model = SOFIAv71(cfg)
    counts = count_parameters(model)
    total = counts["_total"]

    print(f"\nTotal params: {total / 1e6:.2f} M")
    print("\nPer-module breakdown:")
    for k, v in sorted(counts.items(), key=lambda x: -x[1]):
        if k == "_total":
            continue
        pct = 100 * v / total
        print(f"  {k:30s} {v / 1e6:>8.3f} M  ({pct:5.1f}%)")

    print("\n=== Memory footprint estimates ===")
    for prec in ["fp32", "fp16", "int8", "int4"]:
        mb = estimate_quantized_size(total, prec)
        print(f"  {prec.upper():6s} weights: {mb:>8.1f} MB  ({mb / 1024:.2f} GB)")

    print("\n=== Forward pass test ===")
    try:
        shapes = test_forward(model, device=args.device, batch_size=args.batch_size)
        for k, v in shapes.items():
            print(f"  out[{k}]: {v}")
    except Exception as e:
        print(f"  ERROR during forward: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()

    if args.benchmark:
        print("\n=== Latency benchmark ===")
        try:
            stats = benchmark(model, device=args.device)
            print(f"  mean: {stats['mean_ms']:.2f} ms  min: {stats['min_ms']:.2f}  "
                  f"max: {stats['max_ms']:.2f}  std: {stats['std_ms']:.2f}")
        except Exception as e:
            print(f"  ERROR during benchmark: {type(e).__name__}: {e}")


if __name__ == "__main__":
    main()