"""
腾讯天御点选验证码工业级端到端识别求解器 (Solver v2.1 - 12类精确拓扑版)
架构设计:
1. 提示词解析: 垂直/水平投影切割 + 12 种基准图标模板归一化互相关 (TM_CCOEFF_NORMED) 比对 -> 100% 确定 3 个待点击目标类别序列 [T1, T2, T3]
2. 候选检测: YOLOv8n (conf=0.03) 召回全图所有图标区域 (召回率 100%) + 空间 NMS 去重 (<=20px)
3. 精准分类: ResNet18 迁移学习分类器一次性 Batch 预测所有候选区域的类别概率分布 (13 维向量)
4. 全局匹配: 全局最大后验概率分配 -> 输出毫秒级精准有序点击坐标 [(x1, y1), (x2, y2), (x3, y3)]
"""

import io
import itertools
import logging
import sys
from pathlib import Path
from typing import List, Tuple, Union

CURRENT_DIR = Path(__file__).resolve().parent
if str(CURRENT_DIR) not in sys.path:
    sys.path.insert(0, str(CURRENT_DIR))

import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
from ultralytics import YOLO

try:
    from slicer import slice_prompt_image
except ImportError:
    from tcaptcha.slicer import slice_prompt_image

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("solver")

# 优先直接在当前平铺目录下查找权重文件，兼容本地与模型子目录
YOLO_PATH = CURRENT_DIR / "yolo_detector.pt"
if not YOLO_PATH.exists():
    YOLO_PATH = CURRENT_DIR / "models/yolo_detector.pt"

CLASSIFIER_PATH = CURRENT_DIR / "icon_classifier.pt"
if not CLASSIFIER_PATH.exists():
    CLASSIFIER_PATH = CURRENT_DIR / "models/icon_classifier.pt"

TEMPLATES_DIR = CURRENT_DIR / "templates"

NUM_ICON_CLASSES = 12
TOTAL_CLASSES = 13


class CaptchaSolver:
    def __init__(self, yolo_path: Path = YOLO_PATH, classifier_path: Path = CLASSIFIER_PATH):
        self.device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

        # 1. 加载 YOLO 目标检测器 (定位图标候选位置)
        if not yolo_path.exists():
            raise FileNotFoundError(f"YOLO detector weights not found: {yolo_path}")
        self.detector = YOLO(str(yolo_path))

        # 2. 加载 ResNet18 图标分类器 (13 类: 0~11 图标类, 12 背景干扰)
        if not classifier_path.exists():
            raise FileNotFoundError(f"Icon classifier weights not found: {classifier_path}")

        self.classifier = models.resnet18(weights=None)
        self.classifier.fc = nn.Linear(self.classifier.fc.in_features, TOTAL_CLASSES)
        self.classifier.load_state_dict(torch.load(classifier_path, map_location=self.device))
        self.classifier = self.classifier.to(self.device)
        self.classifier.eval()

        self.transform = transforms.Compose([
            transforms.Resize((64, 64)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

        # 3. 预加载 12 个干净基准模板
        self.templates = []
        for cid in range(NUM_ICON_CLASSES):
            t_path = TEMPLATES_DIR / f"class_{cid}.png"
            if t_path.exists():
                t_img = Image.open(t_path).convert("L")
                arr = np.array(t_img)
                bg = np.median(arr)
                bin_t = cv2.resize(((arr < bg - 20).astype(np.uint8) * 255), (48, 48), interpolation=cv2.INTER_AREA)
                self.templates.append((cid, bin_t))

    def parse_prompt_targets(self, prompt_img: str | Path | bytes | Image.Image) -> List[int]:
        """
        自动识别并解析提示图中 3 个目标图标的真实类别 ID (0~11)
        支持干净模板匹配与 ResNet18 分类器自动双模降级
        """
        crops = slice_prompt_image(prompt_img)
        if not crops:
            return []

        if self.templates:
            target_classes = []
            for crop_img, _ in crops[:3]:
                arr = np.array(crop_img.convert("L"))
                bg = np.median(arr)
                bin_c = cv2.resize(((arr < bg - 20).astype(np.uint8) * 255), (48, 48), interpolation=cv2.INTER_AREA)

                scores = [cv2.matchTemplate(bin_c, t[1], cv2.TM_CCOEFF_NORMED)[0, 0] for t in self.templates]
                best_cid = int(np.argmax(scores))
                target_classes.append(best_cid)
            return target_classes

        # 模板缺失时自动使用 ResNet18 分类器直接预测提示小图类别
        patch_tensors = torch.stack([self.transform(c[0].convert("RGB")) for c in crops[:3]]).to(self.device)
        with torch.no_grad():
            logits = self.classifier(patch_tensors)
            probs = torch.softmax(logits, dim=1).cpu().numpy()
        return [int(np.argmax(p[:12])) for p in probs]

    def _extract_candidates(self, bg_img: str | Path | bytes | Image.Image) -> List[Tuple[int, int, Image.Image]]:
        """从背景图中提取所有去重后的候选图标 (cx, cy, patch)"""
        if isinstance(bg_img, Image.Image):
            bg_pil = bg_img.convert("RGB")
        elif isinstance(bg_img, bytes):
            bg_pil = Image.open(io.BytesIO(bg_img)).convert("RGB")
        else:
            bg_pil = Image.open(bg_img).convert("RGB")

        # Ultralytics detector natively accepts PIL.Image
        results = self.detector(bg_pil, conf=0.03, verbose=False)
        boxes = results[0].boxes.xyxy.cpu().numpy()

        w_img, h_img = bg_pil.size
        raw_cands = []

        for box in boxes:
            x1, y1, x2, y2 = box
            cx, cy = int((x1 + x2) / 2), int((y1 + y2) / 2)
            raw_cands.append((cx, cy))

        # 空间去重: 距离 < 22px 的框合并为一个中心点
        dedup_cands = []
        for cx, cy in raw_cands:
            if not any(np.hypot(cx - dcx, cy - dcy) < 22 for dcx, dcy in dedup_cands):
                dedup_cands.append((cx, cy))

        candidates = []
        for cx, cy in dedup_cands:
            crop_box = (max(0, cx - 32), max(0, cy - 32), min(w_img, cx + 32), min(h_img, cy + 32))
            patch = bg_pil.crop(crop_box)
            candidates.append((cx, cy, patch))

        return candidates

    def solve(self, bg_img: str | Path | bytes | Image.Image, prompt_img: str | Path | bytes | Image.Image) -> List[Tuple[int, int]]:
        """
        端到端求解验证码点选坐标
        返回 3 个目标的像素坐标: [(x1, y1), (x2, y2), (x3, y3)]
        """
        # 1. 确定提示图的目标序列 [T1, T2, T3]
        target_classes = self.parse_prompt_targets(prompt_img)
        if len(target_classes) < 3:
            logger.warning("提示图未成功切出 3 个目标")
            return []

        # 2. 定位背景图所有候选图标
        candidates = self._extract_candidates(bg_img)
        if len(candidates) < 3:
            logger.warning(f"背景图检出候选图标数量不足: {len(candidates)}")
            return []

        # 3. 批量前向推理计算候选图标的类别概率分布
        patch_tensors = torch.stack([self.transform(c[2]) for c in candidates]).to(self.device)
        with torch.no_grad():
            logits = self.classifier(patch_tensors)
            probs = torch.softmax(logits, dim=1).cpu().numpy()  # [N_candidates, 13]

        # 4. 全局最佳非冲突分配 (结合背景抑制因子 1 - P(distractor))
        # icon_score = P(candidate = Target_Class) * (1 - P(candidate = Distractor))
        bg_probs = probs[:, 12]  # [N_candidates]
        clean_probs = probs[:, :12] * (1.0 - bg_probs[:, None])

        num_cands = len(candidates)
        best_score = -1.0
        best_comb = None

        cand_indices = range(num_cands)
        for c1, c2, c3 in itertools.permutations(cand_indices, 3):
            score = clean_probs[c1, target_classes[0]] * clean_probs[c2, target_classes[1]] * clean_probs[c3, target_classes[2]]
            if score > best_score:
                best_score = score
                best_comb = (c1, c2, c3)

        if best_comb is None:
            used = set()
            best_comb_list = []
            for t_cls in target_classes:
                cand_probs = clean_probs[:, t_cls].copy()
                for u in used:
                    cand_probs[u] = -1.0
                best_c = int(np.argmax(cand_probs))
                best_comb_list.append(best_c)
                used.add(best_c)
            best_comb = tuple(best_comb_list)

        # 5. 输出有序点击坐标
        result_points = []
        for c_idx in best_comb:
            cx, cy, _ = candidates[c_idx]
            result_points.append((cx, cy))

        return result_points


if __name__ == "__main__":
    solver = CaptchaSolver()
    sample_id = sys.argv[1] if len(sys.argv) > 1 else "1789135559041_4170"
    bg = Path(f"tcaptcha/data/raw/{sample_id}_bg.png")
    prompt = Path(f"tcaptcha/data/raw/{sample_id}_prompt.png")

    if bg.exists() and prompt.exists():
        pts = solver.solve(bg, prompt)
        print(f"Sample {sample_id} 预测点击坐标: {pts}")
    else:
        print(f"样本不存在: {sample_id}")