"""Quantization benchmarking and comparison report.

Generates a structured comparison of original (FP32), FP16, and INT8
quantized ONNX models covering model size, inference latency, and
numerical accuracy delta.
"""

from __future__ import annotations

import logging
import os
import time
from pathlib import Path
from typing import Any, Union

logger = logging.getLogger(__name__)


# --------------------------------------------------------------------------- #
# Internal helpers
# --------------------------------------------------------------------------- #

def _get_file_size_mb(path: Union[str, Path]) -> float:
    """Return file size in megabytes."""
    return os.path.getsize(str(path)) / (1024 * 1024)


def _run_ort_inference(
    model_path: str,
    input_data: Any,
) -> Any:
    """Run a single inference pass with ONNX Runtime on CPU.

    Args:
        model_path: Path to the ``.onnx`` model.
        input_data: ``numpy.ndarray`` matching the model's first input.

    Returns:
        The first output tensor as a ``numpy.ndarray``.
    """
    try:
        import onnxruntime as ort
    except ImportError as exc:
        raise ImportError(
            "Benchmarking requires 'onnxruntime'. "
            "Install it with:  pip install onnxruntime"
        ) from exc

    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = (
        ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    )
    session = ort.InferenceSession(
        model_path,
        sess_options,
        providers=["CPUExecutionProvider"],
    )

    input_name = session.get_inputs()[0].name
    outputs = session.run(None, {input_name: input_data})
    return outputs[0]


def _benchmark_latency(
    model_path: str,
    input_data: Any,
    n_runs: int = 50,
    warmup_runs: int = 5,
) -> dict[str, float]:
    """Benchmark inference latency for an ONNX model.

    Args:
        model_path: Path to the ``.onnx`` model.
        input_data: ``numpy.ndarray`` input.
        n_runs: Number of timed runs.
        warmup_runs: Number of warm-up runs.

    Returns:
        A dict with ``mean_ms``, ``median_ms``, ``min_ms``, ``max_ms``.
    """
    try:
        import numpy as np
        import onnxruntime as ort
    except ImportError as exc:
        raise ImportError(
            "Latency benchmarking requires 'onnxruntime' and 'numpy'. "
            "Install them with:  pip install onnxruntime numpy"
        ) from exc

    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = (
        ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    )
    session = ort.InferenceSession(
        model_path,
        sess_options,
        providers=["CPUExecutionProvider"],
    )

    input_name = session.get_inputs()[0].name
    feed = {input_name: input_data}

    # Warm-up
    for _ in range(warmup_runs):
        session.run(None, feed)

    # Timed runs
    latencies: list[float] = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        session.run(None, feed)
        t1 = time.perf_counter()
        latencies.append((t1 - t0) * 1000.0)

    arr = np.array(latencies)
    return {
        "mean_ms": float(np.mean(arr)),
        "median_ms": float(np.median(arr)),
        "min_ms": float(np.min(arr)),
        "max_ms": float(np.max(arr)),
    }


def _compute_accuracy_delta(
    output: Any,
    baseline: Any,
) -> dict[str, float]:
    """Compute accuracy metrics between a quantized output and a baseline.

    Args:
        output: ``numpy.ndarray`` from the quantized model.
        baseline: ``numpy.ndarray`` from the reference (FP32 or PyTorch) model.

    Returns:
        A dict with ``max_abs_diff``, ``mean_abs_diff``,
        ``cosine_similarity``.
    """
    try:
        import numpy as np
    except ImportError as exc:
        raise ImportError(
            "Accuracy comparison requires 'numpy'. "
            "Install it with:  pip install numpy"
        ) from exc

    output = np.asarray(output, dtype=np.float64).flatten()
    baseline = np.asarray(baseline, dtype=np.float64).flatten()

    diff = np.abs(output - baseline)

    # Cosine similarity
    dot = np.dot(output, baseline)
    norm_out = np.linalg.norm(output)
    norm_base = np.linalg.norm(baseline)
    cosine_sim = (
        float(dot / (norm_out * norm_base))
        if norm_out > 0 and norm_base > 0
        else 0.0
    )

    return {
        "max_abs_diff": float(np.max(diff)),
        "mean_abs_diff": float(np.mean(diff)),
        "cosine_similarity": cosine_sim,
    }


# --------------------------------------------------------------------------- #
# Public API
# --------------------------------------------------------------------------- #

def benchmark_quantization(
    original_path: Union[str, Path],
    fp16_path: Union[str, Path, None] = None,
    int8_path: Union[str, Path, None] = None,
    sample_input: Any = None,
    pytorch_baseline: Any = None,
    n_runs: int = 50,
    warmup_runs: int = 5,
) -> dict[str, dict[str, Any]]:
    """Generate a comparison report across quantization variants.

    Compares original (FP32), FP16, and INT8 ONNX models on three axes:

    1. **Model size** (MB)
    2. **Inference latency** (ms) -- benchmarked with ONNX Runtime on CPU
    3. **Accuracy delta** -- max and mean absolute difference vs. a baseline

    The *baseline* for accuracy comparison is determined as follows:

    * If *pytorch_baseline* is provided, it is used directly.
    * Otherwise, the original FP32 ONNX model's output is used as baseline.
    * If *sample_input* is ``None``, accuracy and latency columns are
      omitted.

    Args:
        original_path: Path to the original FP32 ``.onnx`` model.
        fp16_path: Path to the FP16-quantized ``.onnx`` model, or ``None``
            to skip FP16 comparison.
        int8_path: Path to the INT8-quantized ``.onnx`` model, or ``None``
            to skip INT8 comparison.
        sample_input: A ``numpy.ndarray`` (or ``torch.Tensor``) used for
            inference during benchmarking.  If ``None``, latency and
            accuracy measurements are skipped.
        pytorch_baseline: Optional ``numpy.ndarray`` (or ``torch.Tensor``)
            of the expected output from the PyTorch model.  Used as the
            accuracy reference.
        n_runs: Number of timed inference runs per model.
        warmup_runs: Number of warm-up runs per model.

    Returns:
        A dict keyed by variant name (``"original"``, ``"fp16"``,
        ``"int8"``) where each value is a dict containing:

        * ``path`` -- model path
        * ``size_mb`` -- file size in MB
        * ``latency`` -- latency dict (or ``None``)
        * ``accuracy`` -- accuracy-delta dict (or ``None``)

    Raises:
        ImportError: If ``onnxruntime`` or ``numpy`` is not installed.
        FileNotFoundError: If *original_path* does not exist.
    """
    try:
        import numpy as np
    except ImportError as exc:
        raise ImportError(
            "Quantization benchmarking requires 'numpy'. "
            "Install it with:  pip install numpy"
        ) from exc

    original_path = Path(original_path)
    if not original_path.exists():
        raise FileNotFoundError(
            f"Original ONNX model not found: {original_path}"
        )

    # Coerce sample_input to numpy
    input_np: Any = None
    if sample_input is not None:
        if hasattr(sample_input, "detach"):
            input_np = sample_input.detach().cpu().numpy()
        else:
            input_np = np.asarray(sample_input, dtype=np.float32)

    # Coerce pytorch_baseline to numpy
    baseline_np: Any = None
    if pytorch_baseline is not None:
        if hasattr(pytorch_baseline, "detach"):
            baseline_np = pytorch_baseline.detach().cpu().numpy()
        else:
            baseline_np = np.asarray(pytorch_baseline, dtype=np.float64)

    # ------------------------------------------------------------------ #
    # Collect entries to benchmark
    # ------------------------------------------------------------------ #
    variants: dict[str, Path] = {"original": original_path}
    if fp16_path is not None:
        fp16_path = Path(fp16_path)
        if fp16_path.exists():
            variants["fp16"] = fp16_path
        else:
            logger.warning("FP16 model not found at %s; skipping.", fp16_path)
    if int8_path is not None:
        int8_path = Path(int8_path)
        if int8_path.exists():
            variants["int8"] = int8_path
        else:
            logger.warning("INT8 model not found at %s; skipping.", int8_path)

    # ------------------------------------------------------------------ #
    # If no baseline provided, derive one from the original model
    # ------------------------------------------------------------------ #
    if baseline_np is None and input_np is not None:
        logger.info(
            "No pytorch_baseline provided; deriving baseline from original "
            "FP32 model ..."
        )
        baseline_np = _run_ort_inference(str(original_path), input_np)

    # ------------------------------------------------------------------ #
    # Benchmark each variant
    # ------------------------------------------------------------------ #
    report: dict[str, dict[str, Any]] = {}

    for name, path in variants.items():
        logger.info("Benchmarking variant '%s' at %s ...", name, path)

        entry: dict[str, Any] = {
            "path": str(path),
            "size_mb": round(_get_file_size_mb(path), 2),
            "latency": None,
            "accuracy": None,
        }

        if input_np is not None:
            # Latency
            try:
                entry["latency"] = _benchmark_latency(
                    str(path),
                    input_np,
                    n_runs=n_runs,
                    warmup_runs=warmup_runs,
                )
            except Exception as exc:
                logger.warning(
                    "Latency benchmark failed for '%s': %s", name, exc
                )
                entry["latency"] = {"error": str(exc)}

            # Accuracy (skip for 'original' when baseline comes from itself)
            if baseline_np is not None:
                try:
                    model_output = _run_ort_inference(str(path), input_np)
                    entry["accuracy"] = _compute_accuracy_delta(
                        model_output, baseline_np
                    )
                except Exception as exc:
                    logger.warning(
                        "Accuracy comparison failed for '%s': %s", name, exc
                    )
                    entry["accuracy"] = {"error": str(exc)}

        report[name] = entry

    # ------------------------------------------------------------------ #
    # Summary log
    # ------------------------------------------------------------------ #
    _log_report_summary(report)

    return report


def _log_report_summary(report: dict[str, dict[str, Any]]) -> None:
    """Log a human-readable summary of the benchmark report."""
    lines = ["Quantization benchmark report:"]
    lines.append(
        f"  {'Variant':<10} {'Size (MB)':>10} {'Latency (ms)':>14} "
        f"{'Max Diff':>10} {'Cosine Sim':>12}"
    )
    lines.append("  " + "-" * 60)

    for name, entry in report.items():
        size = f"{entry['size_mb']:.2f}"

        latency = "N/A"
        if entry.get("latency") and "mean_ms" in entry["latency"]:
            latency = f"{entry['latency']['mean_ms']:.2f}"

        max_diff = "N/A"
        cosine = "N/A"
        if entry.get("accuracy") and "max_abs_diff" in entry["accuracy"]:
            max_diff = f"{entry['accuracy']['max_abs_diff']:.6f}"
            cosine = f"{entry['accuracy']['cosine_similarity']:.6f}"

        lines.append(
            f"  {name:<10} {size:>10} {latency:>14} "
            f"{max_diff:>10} {cosine:>12}"
        )

    logger.info("\n".join(lines))
