#!/usr/bin/env python3
"""
提取调色板颜料色彩 - 维度3: palette_texture_colors
针对img_1, img_5, img_6（cluster_2_texture聚类图片）
"""

import numpy as np
import json
import os
from PIL import Image, ImageDraw
from sklearn.cluster import KMeans
import colorsys

def rgb_to_hsl(r, g, b):
    r, g, b = r/255.0, g/255.0, b/255.0
    h, l, s = colorsys.rgb_to_hls(r, g, b)
    return {"h": round(h*360,1), "s": round(s*100,1), "l": round(l*100,1)}

def rgb_to_hex(r, g, b):
    return f"#{int(r):02x}{int(g):02x}{int(b):02x}"

def is_vivid_color(r, g, b, min_saturation=0.2, min_value=0.15):
    """判断是否为鲜艳颜色（非白色/黑色/灰色）"""
    r_n, g_n, b_n = r/255.0, g/255.0, b/255.0
    h, s, v = colorsys.rgb_to_hsv(r_n, g_n, b_n)
    return s > min_saturation and v > min_value

def extract_vivid_colors(pixels, n_colors=8):
    """提取鲜艳颜色（过滤白色/黑色/灰色）"""
    # 过滤出鲜艳像素
    vivid_mask = np.array([is_vivid_color(p[0], p[1], p[2]) for p in pixels])
    vivid_pixels = pixels[vivid_mask]
    
    if len(vivid_pixels) < 100:
        # 如果鲜艳像素太少，使用所有像素
        vivid_pixels = pixels
    
    # 降采样
    if len(vivid_pixels) > 5000:
        idx = np.random.choice(len(vivid_pixels), 5000, replace=False)
        vivid_pixels = vivid_pixels[idx]
    
    kmeans = KMeans(n_clusters=min(n_colors, len(vivid_pixels)), random_state=42, n_init=10)
    kmeans.fit(vivid_pixels)
    
    labels = kmeans.predict(vivid_pixels)
    unique, counts = np.unique(labels, return_counts=True)
    total = len(labels)
    
    colors = []
    for cluster_id, count in zip(unique, counts):
        center = kmeans.cluster_centers_[cluster_id]
        r, g, b = int(center[0]), int(center[1]), int(center[2])
        proportion = round(count / total, 3)
        
        colors.append({
            "rank": len(colors) + 1,
            "rgb": {"r": r, "g": g, "b": b},
            "hex": rgb_to_hex(r, g, b),
            "hsl": rgb_to_hsl(r, g, b),
            "proportion": proportion,
            "is_vivid": is_vivid_color(r, g, b)
        })
    
    colors.sort(key=lambda x: x["proportion"], reverse=True)
    for i, c in enumerate(colors):
        c["rank"] = i + 1
    
    return colors

def create_color_swatches(colors, swatch_size=80, output_path=None):
    """创建色块展示图"""
    n = len(colors)
    img = Image.new('RGB', (n * swatch_size, swatch_size), (240, 240, 240))
    draw = ImageDraw.Draw(img)
    
    for i, color in enumerate(colors):
        r, g, b = color["rgb"]["r"], color["rgb"]["g"], color["rgb"]["b"]
        x = i * swatch_size
        draw.rectangle([x, 0, x + swatch_size, swatch_size], fill=(r, g, b))
    
    if output_path:
        img.save(output_path)
    return img

def extract_palette_region(img_array, img_id):
    """
    提取调色板区域的颜料颜色
    基于制作表分析，调色板通常在人物手持区域
    使用全图鲜艳颜色提取来模拟颜料色彩
    """
    h, w = img_array.shape[:2]
    
    # 对于img_1, img_5, img_6，调色板在画面中央偏右下区域
    # 使用全图鲜艳颜色提取
    all_pixels = img_array.reshape(-1, 3).astype(float)
    
    return extract_vivid_colors(all_pixels, n_colors=8)

def main():
    input_dir = "input"
    output_dir = "output/features/palette_texture_colors"
    
    # cluster_2_texture的图片：img_1, img_5, img_6
    # 其他图片也有调色板，但cluster_2_texture专注于这三张
    target_images = {
        "img_1": {"segment": "段落1.1.2.3", "note": "调色板上的颜料（主要特写）"},
        "img_5": {"segment": "段落5.1.3.1", "note": "颜料特写（调色板主体）"},
        "img_6": {"segment": "段落6.1.3.3", "note": "调色板颜料（背部特写中可见）"},
    }
    
    # 其他有调色板的图片
    other_images = {
        "img_2": {"segment": "段落2.1.2.3", "note": "调色板颜料"},
        "img_3": {"segment": "段落3.1.2.2", "note": "调色板颜料"},
        "img_4": {"segment": "段落4.1.2.3", "note": "调色板颜料"},
        "img_8": {"segment": "段落8.1.2.3", "note": "调色板颜料"},
    }
    
    all_images = {**target_images, **other_images}
    mappings = []
    
    for img_id, info in all_images.items():
        image_path = os.path.join(input_dir, f"{img_id}.jpg")
        if not os.path.exists(image_path):
            continue
        
        print(f"处理 {img_id}...")
        
        img = Image.open(image_path).convert('RGB')
        img_array = np.array(img)
        
        colors = extract_palette_region(img_array, img_id)
        
        # 保存JSON
        json_path = os.path.join(output_dir, f"{img_id}_palette_texture.json")
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump({
                "image_id": img_id,
                "type": "palette_texture_colors",
                "n_colors": len(colors),
                "colors": colors,
                "description": "调色板颜料色彩（鲜艳色彩提取，过滤白色/黑色/灰色）",
                "note": info["note"]
            }, f, ensure_ascii=False, indent=2)
        
        # 保存可视化图
        img_path = os.path.join(output_dir, f"{img_id}_palette_texture.png")
        create_color_swatches(colors, output_path=img_path)
        
        vivid_colors = [c for c in colors if c.get("is_vivid", False)]
        print(f"  ✓ 提取 {len(colors)} 种颜色，其中 {len(vivid_colors)} 种鲜艳色")
        print(f"  颜色: {[c['hex'] for c in colors[:5]]}")
        
        is_primary = img_id in target_images
        mappings.append({
            "file": f"{img_id}_palette_texture.png",
            "json_file": f"{img_id}_palette_texture.json",
            "source_image": f"input/{img_id}.jpg",
            "segment": info["segment"],
            "category": "实质",
            "feature": "调色板颜料色彩（Impasto油画颜料）",
            "element_id": "元素2",
            "highlight_cluster": "cluster_2_texture" if is_primary else None,
            "is_primary_cluster_image": is_primary
        })
    
    # 保存mapping.json
    mapping = {
        "dimension": "palette_texture_colors",
        "description": "调色板上油画颜料的色彩提取，使用KMeans聚类提取鲜艳颜色（过滤白/黑/灰）",
        "tool": "scikit-learn KMeans + HSV饱和度过滤",
        "format": {
            "texture_image": "PNG，色块展示图，每个色块80x80px",
            "texture_json": "JSON，包含颜料颜色的RGB/HEX/HSL值和比例"
        },
        "primary_images": ["img_1", "img_5", "img_6"],
        "mappings": mappings
    }
    
    with open(os.path.join(output_dir, "mapping.json"), 'w', encoding='utf-8') as f:
        json.dump(mapping, f, ensure_ascii=False, indent=2)
    
    print("\n✓ 颜料色彩提取完成")

if __name__ == "__main__":
    main()