"""
ArcisVLM ONNX Export Script
Exports the VL-JEPA visual encoder (ViT) to ONNX format with optional INT8 quantization.

TensorRT Conversion Hint:
  After exporting to ONNX, convert to TensorRT engine with:
    trtexec --onnx=exports/arcisvlm_encoder.onnx \
            --saveEngine=exports/arcisvlm_encoder.trt \
            --fp16 \
            --minShapes=pixel_values:1x3x448x448 \
            --optShapes=pixel_values:4x3x448x448 \
            --maxShapes=pixel_values:16x3x448x448
  Or with INT8:
    trtexec --onnx=exports/arcisvlm_encoder_int8.onnx \
            --saveEngine=exports/arcisvlm_encoder_int8.trt \
            --int8 \
            --calib=calibration_cache.bin \
            --minShapes=pixel_values:1x3x448x448 \
            --optShapes=pixel_values:4x3x448x448 \
            --maxShapes=pixel_values:16x3x448x448
"""

import argparse
import os
import sys
import time
from pathlib import Path

import numpy as np
import torch
import yaml

# ---------------------------------------------------------------------------
# Resolve project root so that model.* imports work regardless of cwd
# ---------------------------------------------------------------------------
SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parent  # adjust if needed
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from model.vit import ViTEncoder          # noqa: E402
from model.vlm import VLJEPAModel         # noqa: E402


# ---------------------------------------------------------------------------
# Constants (must match training config)
# ---------------------------------------------------------------------------
IMAGE_SIZE = 448
PATCH_SIZE = 14
EMBED_DIM = 2048
DEFAULT_CONFIG = "configs/scale_1.3b.yaml"
DEFAULT_OUTPUT = "exports/arcisvlm_encoder.onnx"
WARMUP_ITERS = 10
BENCH_ITERS = 100


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def load_config(config_path: str) -> dict:
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


def build_encoder_from_config(config: dict) -> ViTEncoder:
    """Instantiate a bare ViTEncoder using values from the yaml config."""
    enc_cfg = config.get("encoder", {})
    encoder = ViTEncoder(
        image_size=enc_cfg.get("image_size", IMAGE_SIZE),
        patch_size=enc_cfg.get("patch_size", PATCH_SIZE),
        embed_dim=enc_cfg.get("embed_dim", EMBED_DIM),
        depth=enc_cfg.get("depth", 48),
        num_heads=enc_cfg.get("num_heads", 32),
        mlp_ratio=enc_cfg.get("mlp_ratio", 4.0),
        drop_path_rate=0.0,   # inference — disable stochastic depth
    )
    return encoder


def extract_x_encoder_weights(checkpoint_path: str) -> dict:
    """
    Load a VLJEPAModel checkpoint and strip the 'x_encoder.' prefix so the
    weights can be loaded directly into a standalone ViTEncoder.
    """
    print(f"[export] Loading checkpoint: {checkpoint_path}")
    ckpt = torch.load(checkpoint_path, map_location="cpu")
    state = ckpt.get("model_state_dict", ckpt)

    encoder_state = {}
    prefix = "x_encoder."
    for k, v in state.items():
        if k.startswith(prefix):
            encoder_state[k[len(prefix):]] = v

    if not encoder_state:
        raise RuntimeError(
            "No keys with prefix 'x_encoder.' found in checkpoint. "
            "Available top-level keys: " + str(list(state.keys())[:20])
        )
    print(f"[export] Extracted {len(encoder_state)} x_encoder keys.")
    return encoder_state


def load_encoder(checkpoint_path: str, config: dict, device: torch.device) -> ViTEncoder:
    encoder = build_encoder_from_config(config)
    weights = extract_x_encoder_weights(checkpoint_path)
    missing, unexpected = encoder.load_state_dict(weights, strict=False)
    if missing:
        print(f"[warn] Missing keys ({len(missing)}): {missing[:5]} ...")
    if unexpected:
        print(f"[warn] Unexpected keys ({len(unexpected)}): {unexpected[:5]} ...")
    encoder.eval()
    encoder.to(device)
    print(f"[export] Encoder loaded on {device}.")
    return encoder


# ---------------------------------------------------------------------------
# ONNX export
# ---------------------------------------------------------------------------

def export_to_onnx(
    encoder: ViTEncoder,
    output_path: str,
    device: torch.device,
    opset: int = 17,
) -> None:
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    dummy = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE, device=device)

    dynamic_axes = {
        "pixel_values": {0: "batch_size"},
        "embeddings":   {0: "batch_size"},
    }

    print(f"[export] Exporting to ONNX (opset {opset}): {output_path}")
    with torch.no_grad():
        torch.onnx.export(
            encoder,
            dummy,
            output_path,
            export_params=True,
            opset_version=opset,
            do_constant_folding=True,
            input_names=["pixel_values"],
            output_names=["embeddings"],
            dynamic_axes=dynamic_axes,
            verbose=False,
        )
    size_mb = os.path.getsize(output_path) / (1024 ** 2)
    print(f"[export] Done. File size: {size_mb:.1f} MB")


# ---------------------------------------------------------------------------
# INT8 quantization (post-training, static)
# ---------------------------------------------------------------------------

def quantize_onnx(input_path: str, output_path: str) -> None:
    """Apply static INT8 quantization using onnxruntime's quantization toolkit."""
    try:
        from onnxruntime.quantization import quantize_static, QuantType, CalibrationDataReader
    except ImportError:
        raise ImportError(
            "onnxruntime-tools is required for quantization.\n"
            "Install with: pip install onnxruntime-tools"
        )

    class RandomCalibrationReader(CalibrationDataReader):
        """Generates random calibration data — replace with real images for production."""
        def __init__(self, n_samples: int = 64):
            self.samples = iter(
                [{"pixel_values": np.random.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE).astype(np.float32)}
                 for _ in range(n_samples)]
            )

        def get_next(self):
            return next(self.samples, None)

    print(f"[quantize] Running INT8 static quantization -> {output_path}")
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    quantize_static(
        model_input=input_path,
        model_output=output_path,
        calibration_data_reader=RandomCalibrationReader(),
        quant_type=QuantType.QInt8,
        per_channel=True,
        reduce_range=True,
    )
    size_mb = os.path.getsize(output_path) / (1024 ** 2)
    print(f"[quantize] Done. File size: {size_mb:.1f} MB")


# ---------------------------------------------------------------------------
# Benchmarking
# ---------------------------------------------------------------------------

def _percentile(arr, p):
    return float(np.percentile(arr, p))


def benchmark_pytorch(encoder: ViTEncoder, device: torch.device, batch_size: int = 1) -> dict:
    dummy = torch.randn(batch_size, 3, IMAGE_SIZE, IMAGE_SIZE, device=device)
    latencies = []

    with torch.no_grad():
        # warmup
        for _ in range(WARMUP_ITERS):
            _ = encoder(dummy)
        if device.type == "cuda":
            torch.cuda.synchronize()

        for _ in range(BENCH_ITERS):
            t0 = time.perf_counter()
            _ = encoder(dummy)
            if device.type == "cuda":
                torch.cuda.synchronize()
            latencies.append((time.perf_counter() - t0) * 1000)  # ms

    return {
        "mean_ms": float(np.mean(latencies)),
        "p50_ms":  _percentile(latencies, 50),
        "p95_ms":  _percentile(latencies, 95),
        "p99_ms":  _percentile(latencies, 99),
        "fps":     1000.0 / float(np.mean(latencies)) * batch_size,
    }


def benchmark_onnx(onnx_path: str, device: torch.device, batch_size: int = 1) -> dict:
    try:
        import onnxruntime as ort
    except ImportError:
        raise ImportError("onnxruntime not installed. Run: pip install onnxruntime-gpu")

    providers = (
        ["CUDAExecutionProvider", "CPUExecutionProvider"]
        if device.type == "cuda"
        else ["CPUExecutionProvider"]
    )
    sess_opts = ort.SessionOptions()
    sess_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    sess = ort.InferenceSession(onnx_path, sess_options=sess_opts, providers=providers)

    dummy = np.random.randn(batch_size, 3, IMAGE_SIZE, IMAGE_SIZE).astype(np.float32)
    input_name = sess.get_inputs()[0].name
    latencies = []

    # warmup
    for _ in range(WARMUP_ITERS):
        sess.run(None, {input_name: dummy})

    for _ in range(BENCH_ITERS):
        t0 = time.perf_counter()
        sess.run(None, {input_name: dummy})
        latencies.append((time.perf_counter() - t0) * 1000)

    return {
        "mean_ms": float(np.mean(latencies)),
        "p50_ms":  _percentile(latencies, 50),
        "p95_ms":  _percentile(latencies, 95),
        "p99_ms":  _percentile(latencies, 99),
        "fps":     1000.0 / float(np.mean(latencies)) * batch_size,
    }


def print_benchmark_table(pytorch_stats: dict, onnx_stats: dict, onnx_int8_stats: dict | None = None):
    header = f"{'Backend':<22} {'Mean (ms)':>10} {'P50 (ms)':>10} {'P95 (ms)':>10} {'P99 (ms)':>10} {'FPS':>8}"
    sep = "-" * len(header)
    print("\n" + sep)
    print(header)
    print(sep)

    def row(name, s):
        print(f"{name:<22} {s['mean_ms']:>10.2f} {s['p50_ms']:>10.2f} {s['p95_ms']:>10.2f} {s['p99_ms']:>10.2f} {s['fps']:>8.1f}")

    row("PyTorch (fp32)", pytorch_stats)
    row("ONNX Runtime",   onnx_stats)
    if onnx_int8_stats:
        row("ONNX INT8",  onnx_int8_stats)

    speedup = pytorch_stats["mean_ms"] / onnx_stats["mean_ms"]
    print(sep)
    print(f"  ONNX speedup vs PyTorch: {speedup:.2f}x")
    if onnx_int8_stats:
        int8_speedup = pytorch_stats["mean_ms"] / onnx_int8_stats["mean_ms"]
        print(f"  ONNX INT8 speedup vs PyTorch: {int8_speedup:.2f}x")
    print(sep + "\n")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(
        description="Export ArcisVLM ViT encoder to ONNX and benchmark inference speed."
    )
    p.add_argument(
        "--checkpoint", "-c",
        required=True,
        help="Path to VLJEPAModel checkpoint (.pt / .pth).",
    )
    p.add_argument(
        "--config",
        default=DEFAULT_CONFIG,
        help=f"Path to model config yaml (default: {DEFAULT_CONFIG}).",
    )
    p.add_argument(
        "--output", "-o",
        default=DEFAULT_OUTPUT,
        help=f"Output ONNX path (default: {DEFAULT_OUTPUT}).",
    )
    p.add_argument(
        "--quantize",
        action="store_true",
        help="Also export an INT8-quantized ONNX model.",
    )
    p.add_argument(
        "--opset",
        type=int,
        default=17,
        help="ONNX opset version (default: 17).",
    )
    p.add_argument(
        "--batch-size",
        type=int,
        default=1,
        help="Batch size for benchmarking (default: 1).",
    )
    p.add_argument(
        "--no-benchmark",
        action="store_true",
        help="Skip inference benchmarking.",
    )
    p.add_argument(
        "--device",
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device for PyTorch inference (default: cuda if available).",
    )
    return p.parse_args()


def main():
    args = parse_args()
    device = torch.device(args.device)
    print(f"[export] Device: {device}")

    # --- Load config ---
    config_path = Path(args.config)
    if not config_path.is_absolute():
        config_path = PROJECT_ROOT / config_path
    config = load_config(str(config_path))

    # --- Load encoder ---
    encoder = load_encoder(args.checkpoint, config, device)

    # --- ONNX export ---
    onnx_path = args.output
    export_to_onnx(encoder, onnx_path, device, opset=args.opset)

    # --- INT8 quantization ---
    onnx_int8_path = None
    if args.quantize:
        stem = Path(onnx_path).stem
        onnx_int8_path = str(Path(onnx_path).parent / f"{stem}_int8.onnx")
        quantize_onnx(onnx_path, onnx_int8_path)

    # --- Benchmark ---
    if not args.no_benchmark:
        print(f"\n[bench] Running {BENCH_ITERS} iterations (batch={args.batch_size}) ...")

        pt_stats   = benchmark_pytorch(encoder, device, args.batch_size)
        ort_stats  = benchmark_onnx(onnx_path, device, args.batch_size)
        int8_stats = benchmark_onnx(onnx_int8_path, device, args.batch_size) if onnx_int8_path else None

        print_benchmark_table(pt_stats, ort_stats, int8_stats)

    print("[export] All done.")


if __name__ == "__main__":
    main()