"""
Export utilities for BADAS models.

Provides helper functions for ONNX/TensorRT export compatibility.
"""

import torch
from pathlib import Path
from typing import Optional


def patch_vjepa2_rope() -> bool:
    """
    Patch V-JEPA2's rotate_queries_or_keys for ONNX compatibility.

    The original implementation has operations that don't trace correctly
    for ONNX export. This patches the function with an equivalent that
    works with ONNX's static shape requirements.

    Returns:
        True if patching succeeded, False otherwise.

    Example:
        >>> from badas.export.utils import patch_vjepa2_rope
        >>> patch_vjepa2_rope()
        True
    """
    try:
        import transformers.models.vjepa2.modeling_vjepa2 as vjepa2_module

        def _onnx_compatible_rotate(x, pos):
            B, num_heads, N, D = x.size()
            omega = torch.arange(D // 2, dtype=x.dtype, device=x.device)
            omega /= D / 2.0
            omega = 1.0 / 10000**omega
            freq = pos.unsqueeze(-1) * omega
            emb_sin = freq.sin().repeat(1, 1, 1, 2)
            emb_cos = freq.cos().repeat(1, 1, 1, 2)
            y = x.unflatten(-1, (-1, 2))
            y1, y2 = y.unbind(dim=-1)
            y = torch.stack((-y2, y1), dim=-1).flatten(-2)
            return (x * emb_cos) + (y * emb_sin)

        vjepa2_module.rotate_queries_or_keys = _onnx_compatible_rotate
        return True
    except ImportError:
        return False
    except Exception:
        return False


def embed_onnx_weights(onnx_path: str) -> None:
    """
    Embed external weights into a single ONNX file.

    ONNX export sometimes creates separate .onnx and .onnx.data files.
    This function merges them into a single self-contained ONNX file.

    Args:
        onnx_path: Path to the ONNX model file.

    Example:
        >>> from badas.export.utils import embed_onnx_weights
        >>> embed_onnx_weights("model.onnx")
    """
    import onnx
    from onnx.external_data_helper import load_external_data_for_model

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

    # Try to load external data if it exists
    try:
        load_external_data_for_model(model, str(onnx_path.parent))
    except Exception:
        # Weights might already be embedded
        pass

    # Save with embedded weights
    onnx.save_model(model, str(onnx_path), save_as_external_data=False)

    # Clean up external data file if it exists
    data_file = Path(str(onnx_path) + '.data')
    if data_file.exists():
        data_file.unlink()


def validate_onnx(
    onnx_path: str,
    input_data: Optional[torch.Tensor] = None,
    expected_output: Optional[torch.Tensor] = None,
    rtol: float = 1e-3,
    atol: float = 1e-5,
) -> dict:
    """
    Validate an ONNX model's structure and optionally its numerical output.

    Args:
        onnx_path: Path to the ONNX model.
        input_data: Optional input tensor to run inference with.
        expected_output: Optional expected output to compare against.
        rtol: Relative tolerance for numerical comparison.
        atol: Absolute tolerance for numerical comparison.

    Returns:
        Dict with validation results:
            - structure_valid: bool
            - numerical_valid: bool (if input/expected provided)
            - max_diff: float (if comparison performed)
            - onnx_output: ndarray (if inference performed)

    Example:
        >>> result = validate_onnx("model.onnx", dummy_input, pytorch_output)
        >>> print(f"Valid: {result['structure_valid']}, Max diff: {result['max_diff']}")
    """
    import onnx
    import numpy as np

    result = {
        'structure_valid': False,
        'numerical_valid': None,
        'max_diff': None,
        'onnx_output': None,
    }

    # Check structure
    try:
        model = onnx.load(onnx_path)
        onnx.checker.check_model(model)
        result['structure_valid'] = True
    except Exception as e:
        result['error'] = str(e)
        return result

    # Run inference if input provided
    if input_data is not None:
        try:
            import onnxruntime as ort

            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
            session = ort.InferenceSession(onnx_path, providers=providers)

            input_name = session.get_inputs()[0].name
            output_name = session.get_outputs()[0].name

            # Determine input dtype from model
            input_type = session.get_inputs()[0].type
            if 'float16' in input_type:
                np_input = input_data.cpu().half().numpy()
            else:
                np_input = input_data.cpu().float().numpy()

            onnx_output = session.run([output_name], {input_name: np_input})[0]
            result['onnx_output'] = onnx_output

            # Compare with expected if provided
            if expected_output is not None:
                expected_np = expected_output.cpu().float().numpy()
                max_diff = float(np.max(np.abs(onnx_output - expected_np)))
                result['max_diff'] = max_diff
                result['numerical_valid'] = np.allclose(
                    onnx_output, expected_np, rtol=rtol, atol=atol
                )

        except Exception as e:
            result['inference_error'] = str(e)

    return result


def get_trtexec_command(
    onnx_path: str,
    engine_path: str,
    fp16: bool = True,
    workspace_mb: int = 4096,
) -> str:
    """
    Generate a trtexec command for ONNX to TensorRT conversion.

    Useful when TensorRT Python bindings are not available.

    Args:
        onnx_path: Path to input ONNX model.
        engine_path: Path for output TensorRT engine.
        fp16: Enable FP16 precision.
        workspace_mb: Workspace size in MB.

    Returns:
        trtexec command string.

    Example:
        >>> cmd = get_trtexec_command("model.onnx", "model.trt")
        >>> print(cmd)
        trtexec --onnx=model.onnx --saveEngine=model.trt --fp16 ...
    """
    parts = [
        "trtexec",
        f"--onnx={onnx_path}",
        f"--saveEngine={engine_path}",
        f"--memPoolSize=workspace:{workspace_mb}MiB",
    ]
    if fp16:
        parts.append("--fp16")
    return " ".join(parts)


__all__ = [
    'patch_vjepa2_rope',
    'embed_onnx_weights',
    'validate_onnx',
    'get_trtexec_command',
]
