import geopandas as gpd
import numpy as np
from datetime import datetime, timedelta
import json
import random


def simulate_river_flow(hours=336):
    """模拟河流径流数据
    使用简单的正弦函数加随机波动来模拟径流变化
    """
    # 基础流量(立方米/秒)
    base_flow = random.uniform(50, 200)

    # 生成时间序列数据
    time_steps = np.linspace(0, 2 * np.pi, hours)
    # 使用正弦函数生成周期性变化
    flow = base_flow + base_flow * 0.3 * np.sin(time_steps)
    # 添加随机波动
    flow += np.random.normal(0, base_flow * 0.1, hours)
    return flow.clip(min=0)  # 确保流量非负


def flow_to_color(flow, min_flow, max_flow):
    """根据径流量生成颜色
    低流量 (0-30%): 深蓝 [0, 0, 180, 255]
    中等流量 (30-65%): 绿色过渡 [0, 180, 0, 255]
    高流量 (65-100%): 深红 [180, 0, 0, 255]
    """
    # 将流量归一化到0-1区间
    normalized_flow = (flow - min_flow) / (max_flow - min_flow)

    transition_point = 0.7  # 70%处开始第一次过渡

    if normalized_flow < transition_point:
        # 从蓝色过渡到绿色 (0-70%)
        ratio = normalized_flow / transition_point  # 将0-0.3映射到0-1
        red = 0
        green = int(180 * ratio)
        blue = int(180 * (1 - ratio))
    else:
        # 从绿色过渡到红色 (70%-100%)
        ratio = (normalized_flow - transition_point) / (
            1 - transition_point
        )  # 将0.3-1映射到0-1
        red = int(180 * ratio)
        green = int(180 * (1 - ratio))
        blue = 0

    return [red, green, blue, 255]  # RGBA格式


def create_czml(geojson_path, output_path, year=2024, month=6, day=1, hours=168):
    # 读取GeoJSON数据
    gdf = gpd.read_file(geojson_path)

    # 设置时间范围
    start_time = datetime(year, month, day, 0, 0, 0)

    # 创建CZML文档
    czml = [
        {
            "id": "document",
            "name": "River Network Visualization",
            "version": "1.0",
            "clock": {
                "interval": f"{start_time.isoformat()}Z/{(start_time + timedelta(hours=hours)).isoformat()}Z",
                "currentTime": f"{start_time.isoformat()}Z",
                "multiplier": 3600,  # 1小时/秒
            },
        }
    ]

    # 为每条河流生成数据
    for idx, row in gdf.iterrows():
        # 模拟径流数据
        flow_data = simulate_river_flow(hours)
        min_flow = flow_data.min()
        max_flow = flow_data.max()

        # 提取坐标
        coords = []
        if row.geometry.type == "MultiLineString":
            for line in row.geometry.geoms:
                coords.extend([[x, y, 0] for x, y in line.coords])

        # 创建每个时间点的颜色数据
        color_property = []
        for hour in range(hours):
            current_time = start_time + timedelta(hours=hour)
            color = flow_to_color(flow_data[hour], min_flow, max_flow)

            color_property.append(
                {
                    "interval": f"{current_time.isoformat()}Z/{(current_time + timedelta(hours=1)).isoformat()}Z",
                    "rgba": color,
                }
            )

        # 创建河流对象
        river = {
            "id": f"river_{row['Index']}",
            "name": f"River {row['Index']}",
            "polyline": {
                "positions": {
                    "cartographicDegrees": [
                        coord for point in coords for coord in point
                    ]
                },
                "material": {"solidColor": {"color": color_property}},
                "width": 3,
                "clampToGround": True,  # 贴地显示
                "classificationType": "BOTH",  # 同时在地形和3D切片上显示
                "zIndex": 1,  # 控制叠加顺序
                "distanceDisplayCondition": {
                    "nearFar": [0, 1000000]  # 显示距离范围（米）
                },
                "heightReference": "CLAMP_TO_GROUND",  # 确保贴地
            },
        }

        czml.append(river)

    # 保存CZML文件
    with open(output_path, "w") as f:
        json.dump(czml, f, indent=2)


if __name__ == "__main__":
    # 使用示例
    geojson_path = "input/1.geojson"
    czml_output = "output/rivers.czml"
    create_czml(geojson_path, czml_output, hours=168)