#!/usr/bin/env python3
"""
🔥 Icon Challenge Solver - B&W Exact Pipeline
"""

import os
os.environ["GRADIO_ANALYTICS_ENABLED"] = "0"
os.environ["PYTHONASYNCIODEBUG"] = "0"
os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"

import asyncio
_original_del = asyncio.BaseEventLoop.__del__
def _safe_del(self):
    try:
        if not self.is_closed():
            self.close()
    except (ValueError, OSError, RuntimeError, Exception):
        pass
asyncio.BaseEventLoop.__del__ = _safe_del

import warnings
warnings.filterwarnings("ignore")

import gc
import json
import logging
from pathlib import Path
from typing import List, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T

import numpy as np
from PIL import Image
import cv2

import gradio as gr
import spaces

from ultralytics import YOLO
import open_clip
from peft import PeftModel

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# ===================== CONFIG =====================
BASE_DIR = Path(__file__).resolve().parent

class Config:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    yolo_model_path = BASE_DIR / "yolo" / "best.pt"
    lora_model_dir = BASE_DIR / "icon_model_hybrid"
    adapter_subdir = "lora_adapter"
    model_name = "ViT-L-16-SigLIP-256"
    yolo_conf = 0.25
    # FIX: imgsz [128, 224] exact dari referensi
    yolo_img_size = [128, 224]
    yolo_min_crop_size = 16
    yolo_min_conf = 0.25

CONFIG = Config()

# ===================== PREPROCESS (EXACT MATCH REFERENSI) =====================
preprocess = T.Compose([
    T.Resize((256, 256), interpolation=T.InterpolationMode.NEAREST),
    T.ToTensor(),
    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])

def preprocess_image(img: Image.Image):
    """
    Untuk model SigLIP yang expect 3 channel:
    - Kalau image mode 'L' (B&W), duplicate ke 3 channel
    - Kalau sudah 'RGB', langsung preprocess
    Ini BUKAN convert back ke warna, cuma replicate channel untuk model.
    """
    if img.mode != 'RGB':
        img = img.convert('RGB')
    return preprocess(img)

# ===================== PROJECTION HEAD (EXACT MATCH) =====================
class ProjectionHead(nn.Module):
    def __init__(self, in_dim=1024, hidden_dim=512, out_dim=256, dropout=0.1):
        super().__init__()
        self.layer1 = nn.Linear(in_dim, hidden_dim)
        self.bn1 = nn.BatchNorm1d(hidden_dim)
        self.act = nn.GELU()
        self.dropout = nn.Dropout(dropout)
        self.layer2 = nn.Linear(hidden_dim, out_dim)
        self.bn2 = nn.BatchNorm1d(out_dim)
        self.skip = nn.Linear(in_dim, out_dim)

    def forward(self, x):
        if x.shape[0] <= 1:
            self.eval()
        residual = self.skip(x)
        x = self.layer1(x)
        x = self.bn1(x)
        x = self.act(x)
        x = self.dropout(x)
        x = self.layer2(x)
        x = self.bn2(x)
        return F.normalize(x + residual, p=2, dim=1)

# ===================== ICON MATCHER (EXACT MATCH REFERENSI) =====================
class IconMatcher(nn.Module):
    def __init__(self, lora_path, head_path, device='cuda'):
        super().__init__()
        backbone, _, _ = open_clip.create_model_and_transforms(
            'ViT-L-16-SigLIP-256', pretrained='webli'
        )
        self.backbone = PeftModel.from_pretrained(backbone.visual, lora_path)
        self.backbone = self.backbone.to(device)

        head_state = torch.load(head_path, map_location=device)
        self.projection = ProjectionHead(1024, 512, 256)
        self.projection.load_state_dict(head_state['projection'])
        self.projection = self.projection.to(device)

    def forward(self, x):
        with torch.no_grad():
            features = self.backbone(x)
            return self.projection(features)

# ===================== BLACK & WHITE PIPELINE (EXACT REFERENSI) =====================
def to_black_white(img_np: np.ndarray) -> np.ndarray:
    """
    Exact pipeline dari referensi:
    1. Median Blur (3x3) - hapus noise
    2. Grayscale
    3. OTSU Threshold → Binary Black & White (0 atau 255)
    
    Input: RGB numpy array (H, W, 3)
    Output: Binary B&W numpy array (H, W) uint8
    """
    # 1. Median Blur untuk hapus noise
    blurred = cv2.medianBlur(img_np, 3)
    
    # 2. Ubah ke Grayscale
    gray = cv2.cvtColor(blurred, cv2.COLOR_RGB2GRAY)
    
    # 3. Thresholding OTSU → Hitam + Putih total
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    
    return binary

# ===================== MODEL MANAGER =====================
class ModelManager:
    _instance = None
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance.initialized = False
        return cls._instance

    def init(self):
        if self.initialized:
            return

        logger.info(f"Base dir: {BASE_DIR}")
        logger.info("Loading models...")

        # YOLO
        yolo_path = Path(CONFIG.yolo_model_path)
        if not yolo_path.exists():
            yolo_path = BASE_DIR / "best.pt"
        if not yolo_path.exists():
            raise FileNotFoundError(f"YOLO not found: {yolo_path}")

        self.yolo = YOLO(str(yolo_path)).to(CONFIG.device)
        logger.info(f"✅ YOLO loaded: {yolo_path}")

        # LoRA
        lora_base = Path(CONFIG.lora_model_dir)
        if not lora_base.exists():
            lora_base = BASE_DIR / "icon_model_hybrid"
        if not lora_base.exists():
            raise FileNotFoundError(f"LoRA dir not found")

        adapter_dir = lora_base / CONFIG.adapter_subdir
        lora_path = str(adapter_dir) if adapter_dir.exists() else str(lora_base)
        logger.info(f"LoRA path: {lora_path}")

        # Head Weights
        head_candidates = [
            lora_base / "head_weights.pt",
            lora_base / "weights.pt",
        ]
        head_path = None
        for h in head_candidates:
            if h.exists():
                head_path = str(h)
                logger.info(f"✅ Head weights: {h}")
                break
        if head_path is None:
            raise FileNotFoundError(f"Head weights not found in {lora_base}")

        # Load IconMatcher
        self.matcher = IconMatcher(lora_path, head_path, device=CONFIG.device).eval()
        logger.info("✅ IconMatcher loaded")

        # Config class names
        config_path = lora_base / "config.json"
        if config_path.exists():
            with open(config_path) as f:
                cfg = json.load(f)
            self.idx_to_class = {
                int(k): v for k, v in cfg.get('idx_to_class', {}).items()
            }
        else:
            self.idx_to_class = {}

        self.initialized = True
        logger.info("✅ All models ready!")

    # ===================== YOLO: Detect + Crop + B&W =====================
    def yolo_detect_and_crop(self, image: Image.Image) -> Tuple[Image.Image, Image.Image, dict]:
        """
        Returns:
        - raw_crop: crop asli RGB (untuk display)
        - bw_crop: hasil B&W OTSU (mode 'L', untuk model + display)
        - debug_info
        """
        if image.mode != 'RGB':
            image = image.convert('RGB')

        img_np = np.array(image)  # RGB
        img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

        # FIX: YOLO dengan imgsz [128, 224] exact dari referensi
        results = self.yolo(img_bgr, conf=CONFIG.yolo_conf, imgsz=CONFIG.yolo_img_size, verbose=False)

        best_box = None
        best_conf = 0.0
        for result in results:
            if result.boxes is None:
                continue
            for box in result.boxes:
                conf = float(box.conf[0])
                if conf > best_conf:
                    best_conf = conf
                    best_box = box.xyxy[0].cpu().numpy().astype(int)

        h, w = img_np.shape[:2]

        if best_box is not None and best_conf >= CONFIG.yolo_min_conf:
            x1, y1, x2, y2 = best_box
            x1, y1, x2, y2 = max(0, x1), max(0, y1), min(w, x2), min(h, y2)
            cw, ch = x2 - x1, y2 - y1

            if cw >= CONFIG.yolo_min_crop_size and ch >= CONFIG.yolo_min_crop_size:
                # Crop asli (RGB)
                raw_crop = img_np[y1:y2, x1:x2]
                raw_pil = Image.fromarray(raw_crop)

                # B&W Pipeline exact dari referensi user
                bw_np = to_black_white(raw_crop)
                # Simpan sebagai PIL 'L' (grayscale/B&W) - TIDAK convert back ke RGB
                bw_pil = Image.fromarray(bw_np, mode='L')

                debug = {
                    "detected": True,
                    "conf": best_conf,
                    "box": [int(x1), int(y1), int(x2), int(y2)],
                    "crop_size": [cw, ch]
                }
                return raw_pil, bw_pil, debug

        # Fallback: center crop 80% + B&W
        margin_x = int(w * 0.1)
        margin_y = int(h * 0.1)
        x1, y1 = margin_x, margin_y
        x2, y2 = w - margin_x, h - margin_y

        raw_crop = img_np[y1:y2, x1:x2]
        raw_pil = Image.fromarray(raw_crop)
        bw_np = to_black_white(raw_crop)
        bw_pil = Image.fromarray(bw_np, mode='L')

        debug = {
            "detected": False,
            "conf": 0.0,
            "box": [x1, y1, x2, y2],
            "crop_size": [x2-x1, y2-y1],
            "fallback": "Center crop 80% + B&W (YOLO tidak detect)"
        }
        return raw_pil, bw_pil, debug

    # ===================== INFERENCE: EXACT MATCH REFERENSI =====================
    def match_options(self, challenge_crop: Image.Image, options: List[Image.Image]):
        """
        challenge_crop: PIL Image mode 'L' (B&W hasil OTSU)
        options: List PIL Image (user upload, biasanya mode 'L' B&W juga)
        """
        # Gabungkan challenge + 8 options
        imgs = [challenge_crop] + options
        
        # Preprocess: convert ke 3-channel untuk model (duplicate L channel)
        # Ini BUKAN convert back ke warna, cuma replicate untuk SigLIP
        tensors = [preprocess_image(img) for img in imgs]
        batch = torch.stack(tensors).to(CONFIG.device)

        with torch.no_grad():
            embs = self.matcher(batch)

        # Exact match referensi: torch.mm(embs[0:1], embs[1:].t()).squeeze(0)
        scores = torch.mm(embs[0:1], embs[1:].t()).squeeze(0)
        best_idx = int(torch.argmax(scores).item())
        max_score = float(scores.max().item())

        probs = torch.softmax(scores * 10, dim=0).cpu().numpy()

        return {
            "predicted_index": best_idx,
            "confidence": float(probs[best_idx]),
            "max_score": max_score,
            "all_scores": {
                f"option_{i}": {
                    "score": float(scores[i]),
                    "probability": float(probs[i])
                }
                for i in range(len(options))
            },
            "predicted_class": self.idx_to_class.get(best_idx, "unknown")
        }

manager = ModelManager()

# ===================== GRADIO FUNCTION =====================
@spaces.GPU(request_timeout=180, enable=True)
def solve_gradio(challenge_img, *options):
    if not manager.initialized:
        try:
            manager.init()
        except Exception as e:
            logger.error(f"Init failed: {e}", exc_info=True)
            return "❌ Init failed: " + str(e), None, None, None

    try:
        if challenge_img is None or any(o is None for o in options):
            return "❌ Semua gambar harus diisi!", None, None, None

        if isinstance(challenge_img, np.ndarray):
            challenge_img = Image.fromarray(challenge_img)
        elif isinstance(challenge_img, str):
            challenge_img = Image.open(challenge_img)

        options_pil = []
        for opt in options:
            if isinstance(opt, np.ndarray):
                options_pil.append(Image.fromarray(opt))
            elif isinstance(opt, str):
                options_pil.append(Image.open(opt))
            elif isinstance(opt, Image.Image):
                options_pil.append(opt)
            else:
                options_pil.append(None)

        if any(o is None for o in options_pil):
            return "❌ Format option tidak valid!", None, None, None

        # YOLO detect + crop + B&W pipeline
        raw_crop, bw_crop, debug = manager.yolo_detect_and_crop(challenge_img)

        # Inference pakai B&W crop (mode 'L')
        result = manager.match_options(bw_crop, options_pil)

        best_idx = result["predicted_index"]

        # Build result text
        lines = [
            "🎯 HASIL\n",
            f"✅ Best Option: {best_idx + 1}",
            f"Confidence: {result['confidence']:.4f} ({result['confidence']*100:.2f}%)",
            f"Raw Score: {result['max_score']:.4f}",
            f"Class: {result['predicted_class']}\n",
            "🔧 YOLO DEBUG:"
        ]

        if debug["detected"]:
            lines.append(f"   YOLO Detect: ✅ YES (conf={debug['conf']:.3f})")
            lines.append(f"   Box: {debug['box']}")
            lines.append(f"   Crop Size: {debug['crop_size']}")
        else:
            lines.append(f"   YOLO Detect: ❌ NO")
            lines.append(f"   Fallback: {debug['fallback']}")
            lines.append(f"   Crop Size: {debug['crop_size']}")

        lines.append("\n📊 Skor Semua Opsi:")
        for i in range(len(options_pil)):
            s = result["all_scores"][f"option_{i}"]
            marker = "✅" if i == best_idx else "   "
            lines.append(f"{marker} Option {i+1}: prob={s['probability']:.4f} score={s['score']:.4f}")

        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        # Return: text, best match, raw crop RGB, B&W crop
        return "\n".join(lines), options_pil[best_idx], raw_crop, bw_crop

    except Exception as e:
        logger.error(str(e), exc_info=True)
        return "❌ Error: " + str(e), None, None, None


# ===================== GRADIO UI =====================
with gr.Blocks(title="🔥 Icon Challenge Solver", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔥 Icon Challenge Solver\n**YOLO(128×224) → Crop → MedianBlur → Grayscale → OTSU → SigLIP LoRA**")

    with gr.Row():
        with gr.Column(scale=1):
            challenge_input = gr.Image(label="🖼️ Challenge Image (Original)", type="numpy", height=240)
        with gr.Column(scale=2):
            gr.Markdown("### 🎲 8 Option Images (B&W Icons)")
            with gr.Row():
                opt_inputs = [gr.Image(label=f"Option {i+1}", type="numpy", height=130) for i in range(8)]

    solve_btn = gr.Button("🔍 Solve Challenge", variant="primary")

    with gr.Row():
        with gr.Column():
            result_text = gr.Textbox(label="📋 Result + Debug", lines=16, interactive=False)
        with gr.Column():
            best_image = gr.Image(label="🏆 Best Match")

    with gr.Row():
        with gr.Column():
            raw_crop_img = gr.Image(label="🔴 YOLO Raw Crop (RGB)", height=240)
        with gr.Column():
            bw_crop_img = gr.Image(label="⚫ YOLO B&W Crop (OTSU Binary)", height=240)

    solve_btn.click(
        fn=solve_gradio,
        inputs=[challenge_input] + opt_inputs,
        outputs=[result_text, best_image, raw_crop_img, bw_crop_img]
    )

    gr.Markdown("---\nMade with ❤️ for Icon Challenge")

if __name__ == "__main__":
    demo.launch()