"""
功能 2: 连续数据类型智能识别
自动检测数据应该用 sequential 还是 diverging colormap
"""
import numpy as np
import warnings
from typing import Union


def detect_colormap_type(data: Union[np.ndarray, list], verbose: bool = True) -> str:
    """
    智能检测数据应该用什么类型的 colormap
    
    Args:
        data: 数值数据（numpy 数组或列表）
        verbose: 是否打印检测信息
    
    Returns:
        "sequential" 或 "diverging"
        
    Detection Logic:
        1. 如果数据跨越 0 且接近对称 → diverging
        2. 如果数据均值接近中点且有负值 → diverging
        3. 如果数据包含强烈的正负对比 → diverging
        4. 其他情况 → sequential
        
    Example:
        >>> data = np.array([[-1, 0, 1], [-0.5, 0, 0.5]])
        >>> detect_colormap_type(data)
        'diverging'
        
        >>> data = np.array([[0, 50, 100], [10, 60, 90]])
        >>> detect_colormap_type(data)
        'sequential'
    """
    # 转换为 numpy 数组
    if not isinstance(data, np.ndarray):
        data = np.array(data)
    
    # 展平数据
    data_flat = data.flatten()
    
    # 移除 NaN 和 inf
    data_clean = data_flat[np.isfinite(data_flat)]
    
    if len(data_clean) == 0:
        if verbose:
            warnings.warn("Data contains only NaN or inf values. Using sequential colormap.", UserWarning)
        return "sequential"
    
    data_min = data_clean.min()
    data_max = data_clean.max()
    data_mean = data_clean.mean()
    data_median = np.median(data_clean)
    
    # 规则 1：数据跨越 0 且接近对称
    if data_min < 0 < data_max:
        # 计算对称性
        abs_min = abs(data_min)
        abs_max = abs(data_max)
        symmetry_ratio = min(abs_min, abs_max) / max(abs_min, abs_max)
        
        if symmetry_ratio > 0.6:  # 相对对称（降低阈值，更容易检测到发散）
            if verbose:
                print(f"[Huez Intelligence] Detected diverging data (symmetric around 0)")
                print(f"  Range: [{data_min:.2f}, {data_max:.2f}], Symmetry: {symmetry_ratio:.2f}")
            return "diverging"
    
    # 规则 2：数据均值接近中点
    if data_max != data_min:
        data_range = data_max - data_min
        mean_position = (data_mean - data_min) / data_range
        median_position = (data_median - data_min) / data_range
        
        # 均值或中位数在中间 30%-70% 范围
        if 0.3 < mean_position < 0.7 or 0.3 < median_position < 0.7:
            # 且数据跨越较大范围（避免误判平坦数据）
            if data_range > abs(data_mean) * 0.5:
                if verbose:
                    print(f"[Huez Intelligence] Detected diverging data (centered distribution)")
                    print(f"  Mean position: {mean_position:.2%}, Median position: {median_position:.2%}")
                return "diverging"
    
    # 规则 3：相关矩阵类型（典型的 [-1, 1] 范围）
    if -1.5 <= data_min <= -0.3 and 0.3 <= data_max <= 1.5:
        if data_min < 0 < data_max:
            if verbose:
                print(f"[Huez Intelligence] Detected diverging data (correlation-like range)")
                print(f"  Range: [{data_min:.2f}, {data_max:.2f}]")
            return "diverging"
    
    # 默认：顺序数据
    if verbose:
        print(f"[Huez Intelligence] Detected sequential data")
        print(f"  Range: [{data_min:.2f}, {data_max:.2f}]")
    
    return "sequential"


def detect_colormap_for_heatmap(data: Union[np.ndarray, list], 
                                 center: float = None,
                                 vmin: float = None,
                                 vmax: float = None,
                                 verbose: bool = True) -> str:
    """
    专门为热力图检测 colormap 类型
    考虑用户指定的 center, vmin, vmax 参数
    
    Args:
        data: 数值数据
        center: 用户指定的中心值（如果有）
        vmin: 用户指定的最小值（如果有）
        vmax: 用户指定的最大值（如果有）
        verbose: 是否打印检测信息
    
    Returns:
        "sequential" 或 "diverging"
    """
    # 如果用户明确指定了 center，说明想用 diverging
    if center is not None:
        if verbose:
            print(f"[Huez Intelligence] User specified center={center}, using diverging colormap")
        return "diverging"
    
    # 转换为 numpy 数组
    if not isinstance(data, np.ndarray):
        data = np.array(data)
    
    # 使用通用检测逻辑
    cmap_type = detect_colormap_type(data, verbose=verbose)
    
    # 如果用户指定了 vmin/vmax，检查是否跨越 0
    if vmin is not None and vmax is not None:
        if vmin < 0 < vmax:
            if verbose:
                print(f"[Huez Intelligence] vmin={vmin}, vmax={vmax} crosses 0, using diverging")
            return "diverging"
    
    return cmap_type


if __name__ == "__main__":
    # 测试
    print("=" * 70)
    print("Colormap Type Detection Tests")
    print("=" * 70)
    
    # Test 1: 发散数据（相关矩阵）
    print("\n[Test 1] Correlation matrix:")
    corr_data = np.array([
        [1.0, 0.8, -0.6],
        [0.8, 1.0, -0.4],
        [-0.6, -0.4, 1.0]
    ])
    result = detect_colormap_type(corr_data)
    print(f"Result: {result}\n")
    
    # Test 2: 顺序数据（温度）
    print("[Test 2] Temperature data:")
    temp_data = np.random.rand(10, 10) * 30 + 10  # 10-40度
    result = detect_colormap_type(temp_data)
    print(f"Result: {result}\n")
    
    # Test 3: 发散数据（对称分布）
    print("[Test 3] Symmetric diverging data:")
    div_data = np.random.randn(10, 10)  # 均值≈0
    result = detect_colormap_type(div_data)
    print(f"Result: {result}\n")
    
    # Test 4: 顺序数据（只有正值）
    print("[Test 4] Positive-only data:")
    pos_data = np.random.rand(10, 10) * 100
    result = detect_colormap_type(pos_data)
    print(f"Result: {result}\n")

