#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Color Extractor
颜色提取器 - 从图像中提取主要颜色和生成配色方案
"""

import logging
from typing import List, Tuple, Dict, Any, Optional
import colorsys
import math
from collections import Counter
from PIL import Image
import numpy as np

logger = logging.getLogger(__name__)


class ColorExtractor:
    """颜色提取工具类"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
    
    def extract_dominant_colors(
        self,
        image: Image.Image,
        num_colors: int = 5,
        quality: int = 10
    ) -> List[Tuple[int, int, int]]:
        """
        提取图像的主要颜色

        Args:
            image: 源图像
            num_colors: 提取的颜色数量
            quality: 采样质量（数值越小质量越高，速度越慢）

        Returns:
            颜色列表 [(R, G, B), ...]
        """
        try:
            # 转换为RGB模式
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # 缩小图像以提高处理速度
            small_image = image.resize((150, 150))
            
            # 获取像素数据
            pixels = list(small_image.getdata())
            
            # 按质量参数采样
            pixels = pixels[::quality]
            
            # 统计颜色频率
            color_counts = Counter(pixels)
            
            # 获取最常见的颜色
            dominant_colors = [color for color, count in color_counts.most_common(num_colors)]
            
            return dominant_colors
            
        except Exception as e:
            self.logger.error(f"提取主要颜色失败: {e}")
            return [(128, 128, 128)]  # 返回灰色作为默认值
    
    def extract_palette_kmeans(
        self,
        image: Image.Image,
        num_colors: int = 5
    ) -> List[Tuple[int, int, int]]:
        """
        使用K-means聚类提取调色板

        Args:
            image: 源图像
            num_colors: 提取的颜色数量

        Returns:
            颜色列表
        """
        try:
            # 转换为RGB并缩小
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            small_image = image.resize((100, 100))
            
            # 转换为numpy数组
            data = np.array(small_image)
            data = data.reshape((-1, 3))
            
            # 简单的K-means实现
            colors = self._simple_kmeans(data, num_colors)
            
            return [tuple(map(int, color)) for color in colors]
            
        except Exception as e:
            self.logger.error(f"K-means颜色提取失败: {e}")
            return self.extract_dominant_colors(image, num_colors)
    
    def _simple_kmeans(self, data: np.ndarray, k: int, max_iterations: int = 20) -> List[np.ndarray]:
        """简单的K-means聚类实现"""
        # 随机初始化中心点
        centroids = data[np.random.choice(data.shape[0], k, replace=False)]
        
        for _ in range(max_iterations):
            # 计算每个点到中心点的距离
            distances = np.sqrt(((data - centroids[:, np.newaxis])**2).sum(axis=2))
            
            # 分配到最近的中心点
            closest_cluster = np.argmin(distances, axis=0)
            
            # 更新中心点
            new_centroids = np.array([data[closest_cluster == i].mean(axis=0) for i in range(k)])
            
            # 检查收敛
            if np.allclose(centroids, new_centroids):
                break
                
            centroids = new_centroids
        
        return centroids.tolist()
    
    def get_color_harmony(
        self,
        base_color: Tuple[int, int, int],
        harmony_type: str = "complementary"
    ) -> List[Tuple[int, int, int]]:
        """
        根据基础颜色生成和谐配色

        Args:
            base_color: 基础颜色 (R, G, B)
            harmony_type: 和谐类型 (complementary, analogous, triadic, split_complementary)

        Returns:
            和谐颜色列表
        """
        r, g, b = base_color
        h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
        
        colors = [base_color]  # 包含基础颜色
        
        if harmony_type == "complementary":
            # 互补色（色相相差180度）
            comp_h = (h + 0.5) % 1.0
            comp_r, comp_g, comp_b = colorsys.hsv_to_rgb(comp_h, s, v)
            colors.append((int(comp_r * 255), int(comp_g * 255), int(comp_b * 255)))
            
        elif harmony_type == "analogous":
            # 类似色（色相相差±30度）
            for offset in [-0.083, 0.083]:  # ±30度 = ±0.083
                analog_h = (h + offset) % 1.0
                analog_r, analog_g, analog_b = colorsys.hsv_to_rgb(analog_h, s, v)
                colors.append((int(analog_r * 255), int(analog_g * 255), int(analog_b * 255)))
                
        elif harmony_type == "triadic":
            # 三色调（色相相差120度）
            for offset in [0.333, 0.667]:  # 120度 = 0.333
                triadic_h = (h + offset) % 1.0
                triadic_r, triadic_g, triadic_b = colorsys.hsv_to_rgb(triadic_h, s, v)
                colors.append((int(triadic_r * 255), int(triadic_g * 255), int(triadic_b * 255)))
                
        elif harmony_type == "split_complementary":
            # 分裂互补色
            for offset in [0.417, 0.583]:  # 150度和210度
                split_h = (h + offset) % 1.0
                split_r, split_g, split_b = colorsys.hsv_to_rgb(split_h, s, v)
                colors.append((int(split_r * 255), int(split_g * 255), int(split_b * 255)))
        
        return colors
    
    def adjust_color_brightness(
        self,
        color: Tuple[int, int, int],
        factor: float
    ) -> Tuple[int, int, int]:
        """
        调整颜色亮度

        Args:
            color: 原始颜色 (R, G, B)
            factor: 亮度因子 (>1变亮, <1变暗)

        Returns:
            调整后的颜色
        """
        r, g, b = color
        h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
        
        # 调整明度
        v = min(1.0, max(0.0, v * factor))
        
        new_r, new_g, new_b = colorsys.hsv_to_rgb(h, s, v)
        return (int(new_r * 255), int(new_g * 255), int(new_b * 255))
    
    def adjust_color_saturation(
        self,
        color: Tuple[int, int, int],
        factor: float
    ) -> Tuple[int, int, int]:
        """
        调整颜色饱和度

        Args:
            color: 原始颜色 (R, G, B)
            factor: 饱和度因子 (>1更鲜艳, <1更灰)

        Returns:
            调整后的颜色
        """
        r, g, b = color
        h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
        
        # 调整饱和度
        s = min(1.0, max(0.0, s * factor))
        
        new_r, new_g, new_b = colorsys.hsv_to_rgb(h, s, v)
        return (int(new_r * 255), int(new_g * 255), int(new_b * 255))
    
    def get_color_temperature(self, color: Tuple[int, int, int]) -> str:
        """
        判断颜色的冷暖调

        Args:
            color: 颜色 (R, G, B)

        Returns:
            温度描述 ("warm", "cool", "neutral")
        """
        r, g, b = color
        h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
        
        # 将色相转换为度数
        hue_degrees = h * 360
        
        # 暖色调: 0-60度 (红-黄) 和 300-360度 (紫-红)
        if (0 <= hue_degrees <= 60) or (300 <= hue_degrees <= 360):
            return "warm"
        # 冷色调: 120-240度 (绿-蓝)
        elif 120 <= hue_degrees <= 240:
            return "cool"
        # 中性: 60-120度 和 240-300度
        else:
            return "neutral"
    
    def create_gradient_colors(
        self,
        start_color: Tuple[int, int, int],
        end_color: Tuple[int, int, int],
        steps: int = 5
    ) -> List[Tuple[int, int, int]]:
        """
        创建两个颜色之间的渐变

        Args:
            start_color: 起始颜色
            end_color: 结束颜色
            steps: 渐变步数

        Returns:
            渐变颜色列表
        """
        colors = []
        r1, g1, b1 = start_color
        r2, g2, b2 = end_color
        
        for i in range(steps):
            ratio = i / (steps - 1) if steps > 1 else 0
            
            r = int(r1 * (1 - ratio) + r2 * ratio)
            g = int(g1 * (1 - ratio) + g2 * ratio)
            b = int(b1 * (1 - ratio) + b2 * ratio)
            
            colors.append((r, g, b))
        
        return colors
    
    def find_best_contrast_color(
        self,
        background_color: Tuple[int, int, int],
        candidates: List[Tuple[int, int, int]] = None
    ) -> Tuple[int, int, int]:
        """
        找到与背景色对比度最好的颜色

        Args:
            background_color: 背景颜色
            candidates: 候选颜色列表，如果为None则使用黑白

        Returns:
            对比度最好的颜色
        """
        if candidates is None:
            candidates = [(0, 0, 0), (255, 255, 255)]  # 黑色和白色
        
        best_color = candidates[0]
        best_contrast = 0
        
        for candidate in candidates:
            contrast = self._calculate_contrast_ratio(background_color, candidate)
            if contrast > best_contrast:
                best_contrast = contrast
                best_color = candidate
        
        return best_color
    
    def _calculate_contrast_ratio(
        self,
        color1: Tuple[int, int, int],
        color2: Tuple[int, int, int]
    ) -> float:
        """计算两个颜色的对比度比值"""
        def get_luminance(color):
            r, g, b = [c / 255.0 for c in color]
            
            # 应用gamma校正
            r = r / 12.92 if r <= 0.03928 else pow((r + 0.055) / 1.055, 2.4)
            g = g / 12.92 if g <= 0.03928 else pow((g + 0.055) / 1.055, 2.4)
            b = b / 12.92 if b <= 0.03928 else pow((b + 0.055) / 1.055, 2.4)
            
            return 0.2126 * r + 0.7152 * g + 0.0722 * b
        
        lum1 = get_luminance(color1)
        lum2 = get_luminance(color2)
        
        lighter = max(lum1, lum2)
        darker = min(lum1, lum2)
        
        return (lighter + 0.05) / (darker + 0.05)
    
    def suggest_color_palette(
        self,
        image: Image.Image,
        palette_type: str = "balanced"
    ) -> Dict[str, Tuple[int, int, int]]:
        """
        为图像建议配色方案

        Args:
            image: 源图像
            palette_type: 配色类型 (balanced, warm, cool, high_contrast)

        Returns:
            配色方案字典
        """
        # 提取主要颜色
        dominant_colors = self.extract_dominant_colors(image, 3)
        
        if not dominant_colors:
            return self._get_default_palette()
        
        primary_color = dominant_colors[0]
        
        if palette_type == "balanced":
            colors = self.get_color_harmony(primary_color, "analogous")
        elif palette_type == "warm":
            colors = self._adjust_to_warm_palette(dominant_colors)
        elif palette_type == "cool":
            colors = self._adjust_to_cool_palette(dominant_colors)
        elif palette_type == "high_contrast":
            colors = self.get_color_harmony(primary_color, "complementary")
        else:
            colors = dominant_colors
        
        return {
            "primary": colors[0] if colors else primary_color,
            "secondary": colors[1] if len(colors) > 1 else self.adjust_color_brightness(primary_color, 0.7),
            "accent": colors[2] if len(colors) > 2 else self.get_color_harmony(primary_color, "complementary")[1],
            "text": self.find_best_contrast_color(primary_color),
            "background": self.adjust_color_brightness(primary_color, 1.8)
        }
    
    def _adjust_to_warm_palette(self, colors: List[Tuple[int, int, int]]) -> List[Tuple[int, int, int]]:
        """调整为暖色调配色"""
        warm_colors = []
        for color in colors:
            r, g, b = color
            h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
            
            # 调整色相向暖色方向
            if 60 < h * 360 < 300:  # 如果不是暖色
                h = 0.083 if h > 0.5 else 0.917  # 调整为30度或330度
            
            new_r, new_g, new_b = colorsys.hsv_to_rgb(h, s, v)
            warm_colors.append((int(new_r * 255), int(new_g * 255), int(new_b * 255)))
        
        return warm_colors
    
    def _adjust_to_cool_palette(self, colors: List[Tuple[int, int, int]]) -> List[Tuple[int, int, int]]:
        """调整为冷色调配色"""
        cool_colors = []
        for color in colors:
            r, g, b = color
            h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
            
            # 调整色相向冷色方向
            if not (120 <= h * 360 <= 240):  # 如果不是冷色
                h = 0.5  # 调整为180度（青色）
            
            new_r, new_g, new_b = colorsys.hsv_to_rgb(h, s, v)
            cool_colors.append((int(new_r * 255), int(new_g * 255), int(new_b * 255)))
        
        return cool_colors
    
    def _get_default_palette(self) -> Dict[str, Tuple[int, int, int]]:
        """获取默认配色方案"""
        return {
            "primary": (70, 130, 180),    # 钢蓝色
            "secondary": (100, 149, 237),  # 矢车菊蓝
            "accent": (255, 69, 0),       # 橙红色
            "text": (255, 255, 255),      # 白色
            "background": (240, 248, 255)  # 爱丽丝蓝
        } 