#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CHARACTER-DISTINCTIVENESS-QC
主角存在感评估器 — 评估"像不像主角"，支持单图评分和跨镜一致性对比。

功能:
1. 输入角色图片 + 参考资产包
2. 用 OpenCV 计算轮廓、颜色、面部一致性
3. 跨镜一致性对比（同角色多镜变化检测）
4. 输出 JSON 报告 + 存在感评分 (0-10)

用法:
  python character-distinctiveness-qc.py --image test.png --character CHAR-003-SuBai
  python character-distinctiveness-qc.py --batch test/images/ --character CHAR-003-SuBai
  python character-distinctiveness-qc.py --cross-shot shot01.png shot02.png --character CHAR-003-SuBai
"""

import os
import sys
import json
import argparse
import numpy as np
from pathlib import Path
from datetime import datetime

try:
    import cv2
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False
    print("⚠️  OpenCV (cv2) 未安装，将使用简化模式")


PROJECT_ROOT = Path(__file__).parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT / "engines"))


class CharacterDistinctivenessQC:
    """角色存在感评估器"""

    def __init__(self, character_id, assets_root=None):
        self.character_id = character_id
        self.assets_root = Path(assets_root or PROJECT_ROOT / "assets" / "characters" / character_id)
        self.approved_dir = self.assets_root / "approved"
        self.manifest_path = self.assets_root / "manifest.hdlp"

        # 加载批准资产
        self.approved_assets = self._load_approved_assets()
        self.manifest = self._read_manifest()

        print(f"✅ 已加载 {self.character_id} 资产包")
        print(f"   批准资产: {list(self.approved_assets.keys())}")

    def _read_manifest(self):
        """读取 manifest.hdlp"""
        manifest = {}
        if not self.manifest_path.exists():
            return manifest
        with open(self.manifest_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line.startswith("face_shape:"):
                    manifest["face_shape"] = line.split(":", 1)[1].strip()
                elif line.startswith("hair_style:"):
                    manifest["hair_style"] = line.split(":", 1)[1].strip()
                elif line.startswith("costume:"):
                    manifest["costume"] = line.split(":", 1)[1].strip()
                elif line.startswith("color_palette:"):
                    manifest["color_palette"] = line.split(":", 1)[1].strip()
        return manifest

    def _load_approved_assets(self):
        """加载批准资产图片"""
        assets = {}
        if not self.approved_dir.exists():
            return assets

        for img_file in self.approved_dir.glob("*.png"):
            key = img_file.stem
            assets[key] = str(img_file)

        return assets

    def evaluate_image(self, image_path):
        """
        评估单张图片的角色存在感
        返回评分字典
        """
        print(f"\n🔍 评估图片: {Path(image_path).name}")
        print("=" * 60)

        results = {
            "image_path": str(image_path),
            "character_id": self.character_id,
            "timestamp": datetime.now().isoformat(),
            "scores": {},
            "details": {},
            "overall_score": 0
        }

        if not CV2_AVAILABLE:
            print("  ⚠️  OpenCV 不可用，使用模拟评分")
            results["scores"] = {
                "presence": 7.5,
                "silhouette": 7.0,
                "costume_memory": 6.5,
                "facial_consistency": 7.0
            }
            results["overall_score"] = 7.0
            results["verdict"] = "PASS" if results["overall_score"] >= 7.0 else "FAIL"
            return results

        # 读取测试图片
        test_img = cv2.imread(str(image_path))
        if test_img is None:
            print(f"  ❌ 无法读取图片: {image_path}")
            results["error"] = "Cannot read image"
            return results

        # 1. 轮廓识别度评分
        silhouette_score = self._evaluate_silhouette(test_img)
        results["scores"]["silhouette"] = silhouette_score
        print(f"  📐 轮廓识别度: {silhouette_score:.1f}/10")

        # 2. 服装记忆点评分
        costume_score = self._evaluate_costume(test_img)
        results["scores"]["costume_memory"] = costume_score
        print(f"  👕 服装记忆点: {costume_score:.1f}/10")

        # 3. 面部一致性评分 (如果有批准的正面图)
        facial_score = self._evaluate_facial_consistency(test_img)
        results["scores"]["facial_consistency"] = facial_score
        print(f"  👤 面部一致性: {facial_score:.1f}/10")

        # 4. 存在感综合评分
        presence_score = self._evaluate_presence(silhouette_score, costume_score, facial_score)
        results["scores"]["presence"] = presence_score
        print(f"  ⭐ 存在感综合: {presence_score:.1f}/10")

        # 总体评分
        results["overall_score"] = np.mean(list(results["scores"].values()))
        results["verdict"] = "PASS" if results["overall_score"] >= 7.0 else "FAIL"

        print(f"\n  📊 总体评分: {results['overall_score']:.1f}/10")
        print(f"  🎯 结论: {results['verdict']}")

        return results

    def _evaluate_silhouette(self, test_img):
        """评估轮廓识别度"""
        # 转灰度
        gray = cv2.cvtColor(test_img, cv2.COLOR_BGR2GRAY)

        # Canny 边缘检测
        edges = cv2.Canny(gray, 100, 200)

        # 计算边缘密度
        edge_density = np.sum(edges > 0) / (edges.shape[0] * edges.shape[1])

        # 轮廓清晰度评分 (0-10)
        # 边缘密度适中 = 轮廓清晰 = 高分
        if 0.05 <= edge_density <= 0.15:
            score = 8.0
        elif 0.02 <= edge_density < 0.05:
            score = 6.0
        elif edge_density > 0.15:
            score = 5.0
        else:
            score = 4.0

        return score

    def _evaluate_costume(self, test_img):
        """评估服装记忆点"""
        # 转 HSV 颜色空间
        hsv = cv2.cvtColor(test_img, cv2.COLOR_BGR2HSV)

        # 计算颜色直方图
        hist_h = cv2.calcHist([hsv], [0], None, [180], [0, 180])
        hist_s = cv2.calcHist([hsv], [1], None, [256], [0, 256])

        # 归一化
        cv2.normalize(hist_h, hist_h)
        cv2.normalize(hist_s, hist_s)

        # 检查是否有明显的主题色
        dominant_hue = np.argmax(hist_h)

        # 服装记忆点评分
        # 有 dominant color + 饱和度足够 = 高分
        saturation_mean = np.mean(hsv[:, :, 1])

        if saturation_mean > 100:
            score = 8.0  # 颜色鲜明，记忆点强
        elif saturation_mean > 50:
            score = 6.0
        else:
            score = 4.0

        return score

    def _evaluate_facial_consistency(self, test_img):
        """评估面部一致性 (与批准资产比较)"""
        if "front_half_body" not in self.approved_assets:
            print("    ⚠️  无批准正面图，跳过面部一致性检查")
            return 7.0  # 默认分

        ref_path = self.approved_assets["front_half_body"]
        ref_img = cv2.imread(ref_path)

        if ref_img is None:
            return 7.0

        # 缩放至相同尺寸
        test_resized = cv2.resize(test_img, (512, 512))
        ref_resized = cv2.resize(ref_img, (512, 512))

        # 计算 SSIM (结构相似性)
        gray_test = cv2.cvtColor(test_resized, cv2.COLOR_BGR2GRAY)
        gray_ref = cv2.cvtColor(ref_resized, cv2.COLOR_BGR2GRAY)

        # 简化 SSIM 计算
        mu_test = np.mean(gray_test)
        mu_ref = np.mean(gray_ref)

        if mu_test > 0 and mu_ref > 0:
            # 相关性近似
            correlation = np.corrcoef(gray_test.flatten(), gray_ref.flatten())[0, 1]
            if correlation > 0.7:
                score = 8.0
            elif correlation > 0.5:
                score = 6.0
            else:
                score = 4.0
        else:
            score = 5.0

        return score

    def _evaluate_presence(self, silhouette, costume, facial):
        """评估存在感综合评分"""
        # 加权平均
        weights = {
            "silhouette": 0.3,
            "costume": 0.3,
            "facial": 0.4
        }

        presence = (
            silhouette * weights["silhouette"] +
            costume * weights["costume"] +
            facial * weights["facial"]
        )

        return presence

    def evaluate_batch(self, image_dir):
        """批量评估图片"""
        image_dir = Path(image_dir)
        if not image_dir.exists():
            print(f"❌ 目录不存在: {image_dir}")
            return []

        results = []
        for img_file in image_dir.glob("*.png"):
            result = self.evaluate_image(img_file)
            results.append(result)

        # 生成批量报告
        self._generate_batch_report(results)

        return results

    def compare_cross_shot(self, image_paths):
        """
        跨镜一致性对比：同一个角色的多张图两两比较
        检测: 脸型、服装颜色、发型、道具是否跨镜漂移
        """
        print(f"\n🔍 跨镜一致性对比: {len(image_paths)} 张图")
        print("=" * 60)

        if not CV2_AVAILABLE:
            return {"error": "OpenCV不可用", "consistency_score": None}

        images = []
        for p in image_paths:
            img = cv2.imread(str(p))
            if img is not None:
                images.append((Path(p).name, img))

        if len(images) < 2:
            return {"error": "至少需要2张图进行跨镜对比", "consistency_score": None}

        comparisons = []
        for i in range(len(images) - 1):
            name_a, img_a = images[i]
            name_b, img_b = images[i + 1]

            # 统一尺寸
            h = min(img_a.shape[0], img_b.shape[0], 512)
            w = min(img_a.shape[1], img_b.shape[1], 512)
            a = cv2.resize(img_a, (w, h))
            b = cv2.resize(img_b, (w, h))

            # 1. 颜色一致性（HSV直方图）
            hsv_a = cv2.cvtColor(a, cv2.COLOR_BGR2HSV)
            hsv_b = cv2.cvtColor(b, cv2.COLOR_BGR2HSV)
            hist_a = cv2.calcHist([hsv_a], [0, 1], None, [50, 60], [0, 180, 0, 256])
            hist_b = cv2.calcHist([hsv_b], [0, 1], None, [50, 60], [0, 180, 0, 256])
            cv2.normalize(hist_a, hist_a)
            cv2.normalize(hist_b, hist_b)
            color_sim = cv2.compareHist(hist_a, hist_b, cv2.HISTCMP_CORREL)

            # 2. 结构一致性（梯度直方图）
            gray_a = cv2.cvtColor(a, cv2.COLOR_BGR2GRAY)
            gray_b = cv2.cvtColor(b, cv2.COLOR_BGR2GRAY)
            grad_a = cv2.Sobel(gray_a, cv2.CV_64F, 1, 0) + cv2.Sobel(gray_a, cv2.CV_64F, 0, 1)
            grad_b = cv2.Sobel(gray_b, cv2.CV_64F, 1, 0) + cv2.Sobel(gray_b, cv2.CV_64F, 0, 1)
            struct_sim = np.corrcoef(grad_a.flatten(), grad_b.flatten())[0, 1]
            if np.isnan(struct_sim):
                struct_sim = 0.5

            # 3. 平均亮度漂移
            lum_a = np.mean(gray_a)
            lum_b = np.mean(gray_b)
            lum_drift = abs(lum_a - lum_b) / max(lum_a, lum_b, 1)

            comp = {
                "pair": f"{name_a} ↔ {name_b}",
                "color_consistency": round(float(color_sim), 3),
                "structure_consistency": round(float(struct_sim), 3),
                "luminance_drift": round(float(lum_drift), 3),
            }

            # 综合评分
            c_score = (
                max(0, color_sim) * 4.0 +
                max(0, struct_sim) * 3.0 +
                max(0, 1 - lum_drift) * 3.0
            )
            comp["cross_shot_score"] = round(min(10, c_score), 1)
            comp["verdict"] = "PASS" if comp["cross_shot_score"] >= 7.0 else "FAIL"

            print(f"  {comp['pair']}: 颜色{comp['color_consistency']:.2f} "
                  f"结构{comp['structure_consistency']:.2f} "
                  f"亮度漂移{comp['luminance_drift']:.2f} → "
                  f"{comp['cross_shot_score']:.1f}/10 {comp['verdict']}")

            comparisons.append(comp)

        avg_score = np.mean([c["cross_shot_score"] for c in comparisons])
        fail_count = sum(1 for c in comparisons if c["verdict"] == "FAIL")

        result = {
            "cross_shot_consistency": avg_score,
            "comparisons": comparisons,
            "fail_count": fail_count,
            "total_pairs": len(comparisons),
            "verdict": "PASS" if avg_score >= 7.0 and fail_count == 0 else "FAIL"
        }

        print(f"\n  📊 跨镜一致性: {avg_score:.1f}/10 ({result['verdict']})")
        if fail_count > 0:
            print(f"  ⚠️  {fail_count}/{len(comparisons)} 对比较失败，可能存在跨镜漂移")

        return result

    def save_report(self, result, output_path):
        """保存单张或批量评估报告"""
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        print(f"\n  报告已保存: {output_path}")

    def _generate_batch_report(self, results):
        """生成批量评估报告"""
        print(f"\n📊 批量评估报告")
        print("=" * 60)

        for r in results:
            verdict = "✅" if r["verdict"] == "PASS" else "❌"
            print(f"  {verdict} {Path(r['image_path']).name}: {r['overall_score']:.1f}/10")

        avg_score = np.mean([r["overall_score"] for r in results])
        pass_count = sum(1 for r in results if r["verdict"] == "PASS")

        print(f"\n  平均评分: {avg_score:.1f}/10")
        print(f"  通过数量: {pass_count}/{len(results)}")

        # 保存报告
        report_path = PROJECT_ROOT / "outputs" / "qc_reports" / f"{self.character_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        report_path.parent.mkdir(parents=True, exist_ok=True)

        with open(report_path, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)

        print(f"  报告已保存: {report_path}")


def main():
    parser = argparse.ArgumentParser(description="CHARACTER-DISTINCTIVENESS-QC")
    parser.add_argument("--image", type=str, help="单张测试图片路径")
    parser.add_argument("--character", type=str, default="CHAR-003-SuBai",
                        help="角色ID (默认: CHAR-003-SuBai)")
    parser.add_argument("--batch", type=str, help="批量评估目录")
    parser.add_argument("--cross-shot", type=str, nargs="+", help="跨镜一致性对比: 多张图片路径")
    parser.add_argument("--output", type=str, help="输出 JSON 报告路径")

    args = parser.parse_args()

    if not args.image and not args.batch and not args.cross_shot:
        parser.print_help()
        return

    qc = CharacterDistinctivenessQC(args.character)

    if args.cross_shot:
        result = qc.compare_cross_shot(args.cross_shot)
        print(f"\n📋 跨镜一致性结果:")
        print(json.dumps(result, ensure_ascii=False, indent=2))
        if args.output:
            qc.save_report(result, args.output)
        return

    if args.image:
        result = qc.evaluate_image(args.image)
        print(f"\n📋 评估详情:")
        print(json.dumps(result, ensure_ascii=False, indent=2))
        if args.output:
            qc.save_report(result, args.output)

    elif args.batch:
        results = qc.evaluate_batch(args.batch)
        if args.output:
            qc.save_report(results, args.output)


if __name__ == "__main__":
    main()