app/models/line_art.py

"""
Multi-model line art extraction with intelligent fusion.
"""

import cv2
import numpy as np
import onnxruntime as ort

from app.models.registry import ModelRegistry
from app.config import LineArtStyle
from app.processing.utils import (
    multi_scale_close,
    directional_close,
    bridge_endpoints,
    hysteresis_threshold,
    declutter_lines,
)


class MultiLineArtExtractor:
    """Extracts line art using one or more AI models."""

    def __init__(self, registry: ModelRegistry):
        self.registry = registry

    def _run_model(
        self,
        session: ort.InferenceSession,
        image_bgr: np.ndarray,
        target_size: int = 512,
    ) -> np.ndarray:
        """Run a single ONNX line art model."""
        original_h, original_w = image_bgr.shape[:2]

        rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(rgb, (target_size, target_size))
        blob = resized.astype(np.float32) / 255.0
        blob = np.transpose(blob, (2, 0, 1))
        blob = np.expand_dims(blob, axis=0)

        input_name = session.get_inputs()[0].name
        outputs = session.run(None, {input_name: blob})
        raw = np.squeeze(outputs[0])
        raw = np.clip(raw, 0.0, 1.0)

        return cv2.resize(raw, (original_w, original_h))

    def _preprocess(self, image: np.ndarray) -> np.ndarray:
        """Bilateral filter + CLAHE for cleaner model input."""
        filtered = cv2.bilateralFilter(image, d=9, sigmaColor=75, sigmaSpace=75)
        lab = cv2.cvtColor(filtered, cv2.COLOR_BGR2LAB)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        lab[:, :, 0] = clahe.apply(lab[:, :, 0])
        return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    def extract(
        self,
        image: np.ndarray,
        style: LineArtStyle = LineArtStyle.FUSED,
        density: str = "normal",
        bridge_gaps: bool = True,
    ) -> np.ndarray:
        """
        Extract line art using specified style.
        """
        preprocessed = self._preprocess(image)
        outputs: list[np.ndarray] = []

        # ── Decide which models to run ──
        if style == LineArtStyle.AUTO:
            style = self._detect_best_style(image)

        if style == LineArtStyle.FUSED:
            models_to_run = ["informative_drawings", "anyline"]
        elif style == LineArtStyle.MANGA:
            models_to_run = ["manga_line"]
        elif style == LineArtStyle.ANYLINE:
            models_to_run = ["anyline"]
        else:
            models_to_run = ["informative_drawings"]

        print(f"  [LineArt] Styles to run: {models_to_run}")

        # ── Run each available model ──
        for model_name in models_to_run:
            session = self.registry.get(model_name)
            if session is None:
                continue

            sizes = self._get_inference_sizes(model_name, density)
            scale_outputs = []
            for size in sizes:
                out = self._run_model(session, preprocessed, target_size=size)
                scale_outputs.append(out)

            merged = scale_outputs[0]
            for o in scale_outputs[1:]:
                merged = np.maximum(merged, o)

            outputs.append(merged)

        # ── Fallback: Canny if no models available ──
        if not outputs:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, 50, 150)
            outputs.append(edges.astype(np.float32) / 255.0)

        # ── Fuse multiple model outputs ──
        if len(outputs) == 1:
            fused = outputs[0]
        else:
            print(f"  [LineArt] Fusing {len(outputs)} model outputs...")
            fused = self._fuse_outputs(outputs)

        # ── To uint8 ──
        gray = (fused * 255).astype(np.uint8)

        # ── Hysteresis threshold ──
        thresholds = {
            "sparse": (100, 200),
            "normal": (90, 180),
            "dense":  (80, 160),
        }
        lo, hi = thresholds.get(density, (90, 180))
        binary = hysteresis_threshold(gray, low=lo, high=hi)
        binary = cv2.bitwise_not(binary)  # lines = 0, bg = 255

        # ── Aggressive closing pipeline ──
        print("  [LineArt] Running closing pipeline...")
        binary = self._close_pipeline(binary, density, bridge_gaps=bridge_gaps)

        return binary

    def _fuse_outputs(self, outputs: list[np.ndarray]) -> np.ndarray:
        """Intelligent fusion of multiple model outputs."""
        if len(outputs) == 2:
            a, b = outputs[0], outputs[1]
            agreement = np.minimum(a, b)
            union = np.maximum(a, b)
            fused = 0.7 * agreement + 0.3 * union
            both_strong = (a > 0.3) & (b > 0.3)
            fused[both_strong] = np.maximum(fused[both_strong], 0.8)
            return np.clip(fused, 0.0, 1.0)
        else:
            stacked = np.stack(outputs)
            avg = np.mean(stacked, axis=0)
            agree_count = np.sum(stacked > 0.3, axis=0).astype(np.float32)
            boost = agree_count / len(outputs)
            fused = avg * (0.5 + 0.5 * boost)
            return np.clip(fused, 0.0, 1.0)

    def _close_pipeline(self, binary: np.ndarray, density: str, bridge_gaps: bool = True) -> np.ndarray:
        """Full morphological closing pipeline."""
        h, w = binary.shape
        scale_factor = min(w, h) / 1024 # Normalize based on 1024px baseline
        
        configs = {
            "sparse": {"kernels": [3], "dir_len": 5, "iters": 1, "bridge": 5},
            "normal": {"kernels": [3, 5], "dir_len": 7, "iters": 1, "bridge": 10},
            "dense":  {"kernels": [3, 5, 7], "dir_len": 9, "iters": 2, "bridge": 15},
        }
        cfg = configs.get(density, configs["normal"])
        
        # Scale bridge radius
        bridge_dist = max(1, int(cfg["bridge"] * scale_factor))
        if scale_factor < 0.5: # Small image
            bridge_dist = min(bridge_dist, 8)

        print("            - Multi-scale close...")
        result = multi_scale_close(binary, cfg["kernels"], cfg["iters"])
        print("            - Directional close...")
        result = directional_close(result, cfg["dir_len"], cfg["iters"])

        smooth = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
        result = cv2.morphologyEx(result, cv2.MORPH_CLOSE, smooth, iterations=1)

        if bridge_gaps:
            print(f"            - Bridging endpoints (radius {bridge_dist})...")
            result = bridge_endpoints(result, search_radius=bridge_dist)

        # ── Declutter: Merge close lines and thin out ──
        print("            - Decluttering lines (merge & thin)...")
        result = declutter_lines(result, merge_radius=2, thin_out=True)

        return result

    def _detect_best_style(self, image: np.ndarray) -> LineArtStyle:
        """Auto-detect content type to pick the best model."""
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        saturation = hsv[:, :, 1].mean()
        if saturation > 120:
            if self.registry.available("manga_line"):
                return LineArtStyle.MANGA
        return LineArtStyle.FUSED

    def _get_inference_sizes(self, model_name: str, density: str) -> list[int]:
        """Get multi-scale inference sizes per model and density."""
        base = {
            "informative_drawings": 512,
            "anyline": 768,
            "manga_line": 512,
        }.get(model_name, 512)

        if density == "sparse":
            return [base]
        elif density == "dense":
            return [max(256, base - 128), base, base + 256]
        else:
            return [base, base + 256]