"""
ONNX export utilities for personality classifier.

Exports trained PyTorch models to ONNX format for cross-platform inference.
The exported model can be loaded with onnxruntime-node in the TypeScript API.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING

import torch
import numpy as np

if TYPE_CHECKING:
    from .model import PersonalityClassifier


logger = logging.getLogger(__name__)


def export_to_onnx(
    model: "PersonalityClassifier",
    output_path: str | Path,
    input_dim: int = 384,
    opset_version: int = 14,
    verify: bool = True,
    verbose: bool = False,
) -> Path:
    """
    Export trained PyTorch model to ONNX format.

    Args:
        model: Trained PersonalityClassifier model
        output_path: Path for the output ONNX file
        input_dim: Input embedding dimension (default: 384 for MiniLM)
        opset_version: ONNX opset version (default: 14)
        verify: If True, verify exported model matches PyTorch output
        verbose: If True, print ONNX export warnings

    Returns:
        Path to the exported ONNX file
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Prepare model for export
    model.eval()
    model.cpu()

    # Create dummy input for tracing
    dummy_input = torch.randn(1, input_dim, dtype=torch.float32)

    # Export to ONNX
    logger.info(f"Exporting model to ONNX: {output_path}")

    torch.onnx.export(
        model,
        dummy_input,
        str(output_path),
        input_names=["embedding"],
        output_names=["personality_scores"],
        dynamic_axes={
            "embedding": {0: "batch_size"},
            "personality_scores": {0: "batch_size"},
        },
        opset_version=opset_version,
        do_constant_folding=True,
        verbose=verbose,
    )

    logger.info(f"Model exported successfully: {output_path}")
    logger.info(f"File size: {output_path.stat().st_size / 1024:.1f} KB")

    # Verify the exported model
    if verify:
        verify_onnx_model(model, output_path, input_dim=input_dim)

    return output_path


def verify_onnx_model(
    pytorch_model: "PersonalityClassifier",
    onnx_path: str | Path,
    input_dim: int = 384,
    rtol: float = 1e-4,
    atol: float = 1e-4,
) -> bool:
    """
    Verify that ONNX model produces same output as PyTorch model.

    Args:
        pytorch_model: Original PyTorch model
        onnx_path: Path to exported ONNX model
        input_dim: Input dimension
        rtol: Relative tolerance for comparison
        atol: Absolute tolerance for comparison

    Returns:
        True if verification passes

    Raises:
        AssertionError: If outputs don't match
        ImportError: If onnx or onnxruntime not installed
    """
    try:
        import onnx
        from onnx import checker
    except ImportError:
        raise ImportError(
            "onnx is required for verification. Install with: pip install onnx"
        )

    try:
        import onnxruntime as ort
    except ImportError:
        raise ImportError(
            "onnxruntime is required for verification. "
            "Install with: pip install onnxruntime"
        )

    onnx_path = Path(onnx_path)

    # Check ONNX model is valid
    logger.info("Checking ONNX model validity...")
    onnx_model = onnx.load(str(onnx_path))
    checker.check_model(onnx_model)
    logger.info("ONNX model is valid")

    # Create test inputs
    test_inputs = [
        torch.randn(1, input_dim),  # Single sample
        torch.randn(5, input_dim),  # Batch of 5
        torch.randn(32, input_dim),  # Larger batch
    ]

    # Get PyTorch outputs
    pytorch_model.eval()
    pytorch_model.cpu()

    # Create ONNX session
    session = ort.InferenceSession(
        str(onnx_path),
        providers=["CPUExecutionProvider"],
    )

    logger.info("Comparing PyTorch and ONNX outputs...")

    for i, test_input in enumerate(test_inputs):
        # PyTorch forward pass
        with torch.no_grad():
            pytorch_output = pytorch_model(test_input).numpy()

        # ONNX forward pass
        onnx_output = session.run(
            None,
            {"embedding": test_input.numpy()},
        )[0]

        # Compare outputs
        try:
            np.testing.assert_allclose(
                pytorch_output,
                onnx_output,
                rtol=rtol,
                atol=atol,
            )
            logger.info(f"  Test {i + 1}: PASSED (batch_size={test_input.shape[0]})")
        except AssertionError as e:
            logger.error(f"  Test {i + 1}: FAILED")
            logger.error(f"  Max difference: {np.abs(pytorch_output - onnx_output).max()}")
            raise

    logger.info("All verification tests passed!")
    return True


def get_onnx_model_info(onnx_path: str | Path) -> dict:
    """
    Get information about an ONNX model.

    Args:
        onnx_path: Path to ONNX model file

    Returns:
        Dictionary with model information
    """
    try:
        import onnx
    except ImportError:
        raise ImportError("onnx is required. Install with: pip install onnx")

    onnx_path = Path(onnx_path)
    model = onnx.load(str(onnx_path))

    # Get input/output info
    inputs = []
    for inp in model.graph.input:
        shape = [dim.dim_value for dim in inp.type.tensor_type.shape.dim]
        inputs.append({
            "name": inp.name,
            "shape": shape,
            "dtype": onnx.TensorProto.DataType.Name(inp.type.tensor_type.elem_type),
        })

    outputs = []
    for out in model.graph.output:
        shape = [dim.dim_value for dim in out.type.tensor_type.shape.dim]
        outputs.append({
            "name": out.name,
            "shape": shape,
            "dtype": onnx.TensorProto.DataType.Name(out.type.tensor_type.elem_type),
        })

    return {
        "path": str(onnx_path),
        "size_bytes": onnx_path.stat().st_size,
        "size_kb": onnx_path.stat().st_size / 1024,
        "opset_version": model.opset_import[0].version,
        "ir_version": model.ir_version,
        "producer_name": model.producer_name,
        "inputs": inputs,
        "outputs": outputs,
        "num_nodes": len(model.graph.node),
    }
