"""
ONNX model adapter — load and use ONNX Runtime models.

Supports:
- Loading ONNX models from files
- Converting PyTorch models to ONNX
- Wrapping ONNX sessions as PyTorch-compatible modules

Usage::

    from model_osmosis.adapters.onnx import ONNXAdapter

    adapter = ONNXAdapter()

    # Load an ONNX model
    model = adapter.load("model.onnx")

    # Convert PyTorch model to ONNX
    expert = adapter.from_pytorch(torch_model, input_shape=(1, 3, 224, 224))

    # Run inference
    output = expert(input_tensor)
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn


class ONNXExpert(nn.Module):
    """Wrapped ONNX model as PyTorch-compatible expert.

    This allows ONNX models to be used in ModelOsmosis stacking
    and other operations that expect nn.Module.

    Args:
        session: ONNX Runtime inference session.
        input_names: Names of input tensors.
        output_names: Names of output tensors.
        feature_dim: Output feature dimension.
        name: Human-readable name.
    """

    def __init__(
        self,
        session: Any,
        input_names: List[str],
        output_names: List[str],
        feature_dim: int,
        name: str = "onnx_expert",
    ) -> None:
        super().__init__()
        self.session = session
        self.input_names = input_names
        self.output_names = output_names
        self.feature_dim = feature_dim
        self.hidden_size = feature_dim
        self.name = name

        # Store input info for shape inference
        inputs = session.get_inputs()
        self.input_shapes = {inp.name: inp.shape for inp in inputs}
        self.input_types = {inp.name: inp.type for inp in inputs}

    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
        """Run ONNX inference and return as tensor.

        Args:
            x: Input tensor.

        Returns:
            Output tensor.
        """
        # Convert to numpy
        x_np = x.detach().cpu().numpy()

        # Prepare inputs
        inputs = {self.input_names[0]: x_np}

        # Add any additional inputs from kwargs
        for name in self.input_names[1:]:
            if name in kwargs:
                val = kwargs[name]
                if isinstance(val, torch.Tensor):
                    val = val.detach().cpu().numpy()
                inputs[name] = val

        # Run inference
        outputs = self.session.run(self.output_names, inputs)

        # Convert back to tensor
        return torch.from_numpy(outputs[0]).to(x.device)

    def extra_repr(self) -> str:
        return (
            f"name={self.name}, feature_dim={self.feature_dim}, "
            f"inputs={self.input_names}, outputs={self.output_names}"
        )


class ONNXAdapter:
    """Load and convert ONNX models.

    Args:
        providers: ONNX Runtime execution providers.
        verbose: Print loading progress.
    """

    def __init__(
        self,
        providers: Optional[List[str]] = None,
        verbose: bool = True,
    ) -> None:
        self.verbose = verbose

        # Default providers
        if providers is None:
            providers = ["CPUExecutionProvider"]
            try:
                import onnxruntime as ort
                available = ort.get_available_providers()
                if "CUDAExecutionProvider" in available:
                    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
            except ImportError:
                pass

        self.providers = providers

        # Check ONNX Runtime availability
        try:
            import onnxruntime as ort
            self._ort = ort
        except ImportError:
            raise ImportError(
                "ONNX adapter requires `onnxruntime`. "
                "Install with: pip install onnxruntime"
            )

    def load(
        self,
        model_path: Union[str, Path],
        name: Optional[str] = None,
    ) -> ONNXExpert:
        """Load an ONNX model from file.

        Args:
            model_path: Path to .onnx file.
            name: Human-readable name.

        Returns:
            ONNXExpert ready for inference.
        """
        model_path = Path(model_path)

        if not model_path.exists():
            raise FileNotFoundError(f"ONNX model not found: {model_path}")

        if self.verbose:
            print(f"[ONNX] Loading '{model_path}'...")

        # Create session
        session = self._ort.InferenceSession(
            str(model_path),
            providers=self.providers,
        )

        # Get input/output names
        input_names = [inp.name for inp in session.get_inputs()]
        output_names = [out.name for out in session.get_outputs()]

        # Infer feature dim from output shape
        output_shape = session.get_outputs()[0].shape
        if isinstance(output_shape, list) and len(output_shape) > 1:
            feature_dim = output_shape[-1] if output_shape[-1] is not None else 1
        else:
            feature_dim = 1

        if self.verbose:
            print(f"[ONNX] Loaded: inputs={input_names}, outputs={output_names}")

        return ONNXExpert(
            session=session,
            input_names=input_names,
            output_names=output_names,
            feature_dim=feature_dim,
            name=name or model_path.stem,
        )

    def from_pytorch(
        self,
        model: nn.Module,
        input_shape: Tuple[int, ...],
        output_path: Optional[Union[str, Path]] = None,
        name: Optional[str] = None,
        opset_version: int = 14,
        dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None,
    ) -> ONNXExpert:
        """Convert a PyTorch model to ONNX and wrap as expert.

        Args:
            model: PyTorch model to convert.
            input_shape: Shape of input tensor (including batch).
            output_path: Where to save the .onnx file (temp file if None).
            name: Human-readable name.
            opset_version: ONNX opset version.
            dynamic_axes: Dynamic axes for variable-length inputs.

        Returns:
            ONNXExpert wrapping the converted model.
        """
        import tempfile

        model.eval()

        # Create dummy input
        dummy_input = torch.randn(input_shape)

        # Output path
        if output_path is None:
            output_path = tempfile.mktemp(suffix=".onnx")
        output_path = Path(output_path)

        # Default dynamic axes (batch dimension)
        if dynamic_axes is None:
            dynamic_axes = {
                "input": {0: "batch"},
                "output": {0: "batch"},
            }

        if self.verbose:
            print(f"[ONNX] Converting PyTorch model to ONNX...")

        # Export to ONNX
        torch.onnx.export(
            model,
            dummy_input,
            str(output_path),
            input_names=["input"],
            output_names=["output"],
            dynamic_axes=dynamic_axes,
            opset_version=opset_version,
            do_constant_folding=True,
        )

        if self.verbose:
            print(f"[ONNX] Saved to: {output_path}")

        # Load the converted model
        return self.load(output_path, name=name)

    def validate(
        self,
        onnx_expert: ONNXExpert,
        pytorch_model: nn.Module,
        input_shape: Tuple[int, ...],
        rtol: float = 1e-3,
        atol: float = 1e-5,
    ) -> bool:
        """Validate ONNX model matches PyTorch model.

        Args:
            onnx_expert: ONNX expert to validate.
            pytorch_model: Original PyTorch model.
            input_shape: Shape for test input.
            rtol: Relative tolerance.
            atol: Absolute tolerance.

        Returns:
            True if outputs match within tolerance.
        """
        pytorch_model.eval()

        # Create test input
        test_input = torch.randn(input_shape)

        # Get PyTorch output
        with torch.no_grad():
            pytorch_output = pytorch_model(test_input)

        # Get ONNX output
        onnx_output = onnx_expert(test_input)

        # Compare
        if isinstance(pytorch_output, tuple):
            pytorch_output = pytorch_output[0]

        match = torch.allclose(
            pytorch_output,
            onnx_output,
            rtol=rtol,
            atol=atol,
        )

        if self.verbose:
            if match:
                print("[ONNX] Validation passed: outputs match")
            else:
                diff = (pytorch_output - onnx_output).abs().max().item()
                print(f"[ONNX] Validation failed: max diff = {diff}")

        return match
