import numpy as np
from typing import Dict, Tuple
import json


def create_color_gradient_mapping(val_sequence: np.ndarray) -> Dict[int, tuple]:
    """
    创建颜色渐变映射,颜色值范围为0-255,使用Python list类型,分界值四舍五入为整数

    Parameters:
        val_sequence: numpy数组,包含分界值序列

    Returns:
        Dict[int, tuple]: 键为四舍五入后的分界值(整数),值为对应的RGBA颜色元组(Python list)
    """
    # 确保输入是numpy数组
    val_sequence = np.array(val_sequence)

    # 定义基础颜色 (蓝色到红色),使用0-255格式和Python list类型
    blue = [0, 0, 255, 255]  # 起始蓝色 (R,G,B,A)
    red = [255, 0, 0, 255]  # 结束红色 (R,G,B,A)

    # 创建结果字典
    result = {}

    # 添加第一个元素（纯蓝色）
    result[round(float(val_sequence[0]))] = tuple(blue)

    # 为中间值生成渐变色
    for i in range(1, len(val_sequence) - 1):
        # 计算当前位置的比例
        ratio = i / (len(val_sequence) - 1)
        # 线性插值计算当前颜色
        current_color = [int(b * (1 - ratio) + r * ratio) for b, r in zip(blue, red)]

        result[round(float(val_sequence[i]))] = tuple(current_color)

    # 添加最后一个元素（纯红色）
    result[round(float(val_sequence[-1]))] = tuple(red)

    return result


def get_color_from_value(
    value: float, color_mapping: Dict[int, Tuple[int, int, int, int]]
) -> Tuple[int, int, int, int]:
    """
    根据输入的值,从颜色映射字典中查找对应的颜色,使用线性插值

    Parameters:
        value: 输入值
        color_mapping: 颜色映射字典

    Returns:
        Tuple[int, int, int, int]: RGBA颜色元组
    """
    if value <= 10:
        return (0, 0, 255, 255)
    if value >= 3500:
        return (255, 0, 0, 255)
    # 找到小于等于输入值的最大分界值
    max_key = max([k for k in color_mapping if k <= value])

    # 找到大于输入值的最小分界值
    min_key = min([k for k in color_mapping if k > value])

    # 计算当前值在两个分界值之间的比例
    max_val = color_mapping[max_key]
    min_val = color_mapping[min_key]
    ratio = (value - max_key) / (min_key - max_key)

    # 线性插值计算颜色
    current_color = [
        int(max_c * (1 - ratio) + min_c * ratio)
        for max_c, min_c in zip(max_val, min_val)
    ]
    return tuple(current_color)


# 使用示例:
if __name__ == "__main__":
    # 生成原始序列
    x_min = 10
    x_max = 3500
    a = 0.25
    n = 19
    val = (np.linspace((x_min**a), (x_max**a), n)) ** (1 / a)

    # 生成颜色映射
    color_mapping = create_color_gradient_mapping(val)

    print(json.dumps(color_mapping))

    # 使用get_color_from_value()方法查找颜色
    print(get_color_from_value(89.47, color_mapping))  # (15, 0, 240, 255)
    print(get_color_from_value(3500, color_mapping))  # (255, 0, 0, 255)
    print(get_color_from_value(50, color_mapping))  # (64, 0, 191, 255)