"""SpEED-QA — Spatial Efficient Entropic Differencing for Quality Assessment.

Bampis et al. 2017 — NR-VQA based on local entropy differences between
consecutive frames. Efficient blind quality predictor using spatial
entropic differencing without requiring motion estimation.

speedqa_score — higher = better quality
"""

import logging
import cv2
import numpy as np
from typing import Optional

from ayase.models import Sample, QualityMetrics
from ayase.pipeline import PipelineModule

logger = logging.getLogger(__name__)


def _local_entropy(gray: np.ndarray, kernel_size: int = 7) -> np.ndarray:
    """Compute local entropy map using histogram-based estimation."""
    h, w = gray.shape
    pad = kernel_size // 2
    gray_uint8 = np.clip(gray, 0, 255).astype(np.uint8)

    # Use block-based entropy for efficiency
    block_h = max(1, h // 16)
    block_w = max(1, w // 16)
    entropy_map = np.zeros((16, 16), dtype=np.float64)

    for i in range(16):
        for j in range(16):
            block = gray_uint8[
                i * block_h : (i + 1) * block_h,
                j * block_w : (j + 1) * block_w,
            ]
            if block.size == 0:
                continue
            hist = np.histogram(block, bins=64, range=(0, 256))[0]
            hist = hist / (hist.sum() + 1e-7)
            hist = hist[hist > 0]
            entropy_map[i, j] = -np.sum(hist * np.log2(hist + 1e-10))

    # Resize back to original resolution
    return cv2.resize(entropy_map, (w, h), interpolation=cv2.INTER_LINEAR)


def _entropic_difference(gray1: np.ndarray, gray2: np.ndarray) -> np.ndarray:
    """Compute local entropic difference between two frames."""
    ent1 = _local_entropy(gray1)
    ent2 = _local_entropy(gray2)
    return np.abs(ent2 - ent1)


class SpEEDQAModule(PipelineModule):
    name = "speedqa"
    description = "SpEED-QA spatial efficient entropic differencing NR-VQA (Bampis 2017)"
    default_config = {
        "subsample": 8,
    }

    def __init__(self, config=None):
        super().__init__(config)
        self.subsample = self.config.get("subsample", 8)
        self._model = None
        self._backend = "heuristic"

    def setup(self) -> None:
        # Tier 1: Try native SpEED-QA
        try:
            import speedqa
            self._model = speedqa
            self._backend = "native"
            logger.info("SpEED-QA (native) initialised")
            return
        except ImportError:
            pass

        # Tier 2: Heuristic fallback
        self._backend = "heuristic"
        logger.info("SpEED-QA (heuristic) initialised — install speedqa for full model")

    def process(self, sample: Sample) -> Sample:
        try:
            if self._backend == "native":
                score = float(self._model.predict(str(sample.path)))
            else:
                score = self._process_heuristic(sample)

            if score is not None:
                if sample.quality_metrics is None:
                    sample.quality_metrics = QualityMetrics()
                sample.quality_metrics.speedqa_score = score

        except Exception as e:
            logger.warning(f"SpEED-QA failed for {sample.path}: {e}")

        return sample

    def _process_heuristic(self, sample: Sample) -> Optional[float]:
        """Heuristic: local entropy differences between frames."""
        if not sample.is_video:
            # For images, compute spatial entropy quality
            img = cv2.imread(str(sample.path))
            if img is None:
                return None
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).astype(np.float64)
            ent = _local_entropy(gray)
            # Higher, more uniform entropy = richer content
            ent_mean = np.mean(ent)
            ent_uniformity = 1.0 / (1.0 + np.std(ent))
            return float(np.clip(
                0.6 * min(ent_mean / 6.0, 1.0) + 0.4 * ent_uniformity, 0.0, 1.0
            ))

        cap = cv2.VideoCapture(str(sample.path))
        try:
            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            if total <= 1:
                return None

            n_frames = min(self.subsample + 1, total)
            indices = np.linspace(0, total - 1, n_frames, dtype=int)

            frames_gray = []
            for idx in indices:
                cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
                ret, frame = cap.read()
                if ret:
                    frames_gray.append(
                        cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY).astype(np.float64)
                    )
        finally:
            cap.release()

        if len(frames_gray) < 2:
            return None

        # Compute entropic differences between consecutive frames
        ent_diffs = []
        for i in range(len(frames_gray) - 1):
            ed = _entropic_difference(frames_gray[i], frames_gray[i + 1])
            ent_diffs.append(np.mean(ed))

        ent_diffs = np.array(ent_diffs)

        # Low entropic difference = temporally consistent = high quality
        mean_ent_diff = np.mean(ent_diffs)
        var_ent_diff = np.var(ent_diffs)

        # Spatial quality from average frame entropy
        spatial_entropies = []
        for gray in frames_gray:
            ent = _local_entropy(gray)
            spatial_entropies.append(np.mean(ent))
        spatial_quality = min(np.mean(spatial_entropies) / 6.0, 1.0)

        # Temporal quality: lower/more consistent entropic diff = better
        temporal_quality = 1.0 / (1.0 + mean_ent_diff * 0.5)
        temporal_consistency = 1.0 / (1.0 + var_ent_diff * 2.0)

        score = (
            0.35 * spatial_quality
            + 0.40 * temporal_quality
            + 0.25 * temporal_consistency
        )

        return float(np.clip(score, 0.0, 1.0))
