"""
PaddleOCR-VL Layer-12 Feature Extractor — ONNX Inference
=========================================================
Pure ONNX Runtime inference with NO PyTorch dependency.

Extracts 1152-dimensional intermediate features from layer 12 of the
PaddleOCR-VL vision encoder. These features capture mid-level visual
patterns (texture, edge structure, compression artifacts) that are
highly predictive of image quality for OCR applications.

Usage:
    from inference.onnx_inference import Layer12ONNXExtractor
    from PIL import Image

    extractor = Layer12ONNXExtractor("paddleocr_layer12.onnx")
    img = Image.open("document.jpg").convert("RGB")
    features = extractor.extract(img)  # → np.ndarray shape (1152,)
"""

from __future__ import annotations

import os
from typing import Dict, List, Optional, Tuple

import numpy as np
from PIL import Image

from .preprocessing import preprocess_for_onnx


class Layer12ONNXExtractor:
    """
    ONNX-based feature extractor for PaddleOCR-VL layer 12.

    Input:  PIL image (any size, RGB)
    Output: 1152D feature vector (mean-pooled across patches)

    Model architecture:
        PaddleOCR-VL vision encoder truncated after layer 12.
        237M params → 0.8 MB ONNX (constant-folded).

    Key properties:
        - Dynamic input resolution: any image size → smart_resize →
          variable num_patches
        - Batch inference: pass batch of images
        - No GPU required: CPU inference ~50-200ms/image
    """

    def __init__(
        self,
        model_path: str = "model.onnx",
        providers: Optional[List[str]] = None,
    ):
        """
        Args:
            model_path: Path to ONNX model file.
            providers: ONNX Runtime execution providers.
                       Default: ['CPUExecutionProvider']
                       For GPU: ['CUDAExecutionProvider', 'CPUExecutionProvider']
        """
        import onnxruntime as ort

        if providers is None:
            # Auto-detect available providers
            available = ort.get_available_providers()
            if "CUDAExecutionProvider" in available:
                providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
            else:
                providers = ["CPUExecutionProvider"]

        self._sess = ort.InferenceSession(model_path, providers=providers)
        self._providers = providers

        # Cache model metadata
        self.input_names = [inp.name for inp in self._sess.get_inputs()]
        self.output_names = [out.name for out in self._sess.get_outputs()]

        # Feature dimension
        output_shape = self._sess.get_outputs()[0].shape
        self.feature_dim = output_shape[-1]  # 1152

        print(f"[Layer12ONNX] Loaded: {os.path.basename(model_path)}")
        print(f"[Layer12ONNX]   Providers: {self._sess.get_providers()}")
        print(f"[Layer12ONNX]   Feature dim: {self.feature_dim}")
        print(f"[Layer12ONNX]   Input: {list(zip(self.input_names, [i.shape for i in self._sess.get_inputs()]))}")

    def extract(self, image: Image.Image) -> np.ndarray:
        """
        Extract layer-12 features from a single image.

        Args:
            image: PIL RGB image.

        Returns:
            features: (1152,) float32 — mean-pooled layer-12 hidden states.
        """
        pixel_values, position_ids = preprocess_for_onnx(image)

        outputs = self._sess.run(
            self.output_names,
            {
                "pixel_values": pixel_values,
                "position_ids": position_ids,
            },
        )

        return outputs[0].flatten().astype(np.float32)

    def extract_batch(
        self,
        images: List[Image.Image],
    ) -> np.ndarray:
        """
        Extract features from a batch of images.

        Note: Since each image produces a different num_patches (dynamic),
        batch inference processes images sequentially by default. For
        same-size images, you can stack them manually.

        Args:
            images: List of PIL images.

        Returns:
            features: (batch_size, 1152) float32
        """
        features = [self.extract(img) for img in images]
        return np.stack(features, axis=0)

    def distance_from_reference(
        self,
        image: Image.Image,
        reference: Image.Image,
    ) -> float:
        """
        Compute cosine distance from a pristine reference image.

        Higher distance → more degraded. Range: [0, 2] typically.

        Args:
            image: Test image.
            reference: Pristine reference image.

        Returns:
            cosine_distance: 1 - cosine_similarity(f_test, f_ref)
        """
        f_test = self.extract(image)
        f_ref = self.extract(reference)

        cos_sim = np.dot(f_test, f_ref) / (
            np.linalg.norm(f_test) * np.linalg.norm(f_ref) + 1e-12
        )
        return float(1.0 - cos_sim)

    def quality_score(
        self,
        image: Image.Image,
        reference: Optional[Image.Image] = None,
    ) -> float:
        """
        Convert feature distance to a quality score in [0, 1].

        If reference is provided: score = exp(-5 * cosine_distance)
        If no reference: returns raw feature norm (heuristic).

        Args:
            image: Test image.
            reference: Optional pristine reference.

        Returns:
            score: 1.0 = pristine, →0 = heavily degraded.
        """
        if reference is not None:
            dist = self.distance_from_reference(image, reference)
            return float(np.exp(-5.0 * dist))

        # Heuristic: feature magnitude as quality proxy
        features = self.extract(image)
        norm = np.linalg.norm(features) / np.sqrt(self.feature_dim)
        return float(np.clip(1.0 / (1.0 + np.exp(-5 * (0.5 - norm))), 0.0, 1.0))

    @property
    def provider(self) -> str:
        """Active execution provider."""
        return self._sess.get_providers()[0]