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

# 读取GeoJSON文件
gdf = gpd.read_file('input/1.geojson')

# 定义模拟径流数据的函数
def simulate_flow(river_id, hours):
    time = np.arange(hours)
    flow = 10 * np.sin(2 * np.pi * time / 24) + 20
    return flow

# 为每条河流生成径流数据
gdf['flow_data'] = gdf['Index'].apply(lambda x: simulate_flow(x, 336))

# 定义生成CZML的函数
def generate_czml(gdf, start_time):
    czml = [
        {
            "id": "document",
            "version": "1.0",
            "name": "River Flow Visualization",
            "clock": {
                "interval": f"{start_time.isoformat()}Z/{start_time + timedelta(hours=336)}Z",
                "currentTime": start_time.isoformat() + 'Z',
                "multiplier": 1,
                "range": "LOOP_STOP",
                "step": "SYSTEM_CLOCK_MULTIPLIER"
            }
        }
    ]

    for idx, row in gdf.iterrows():
        positions = []
        colors = []
        for i, flow in enumerate(row['flow_data']):
            time = start_time + timedelta(hours=i)
            positions.append({
                "interval": time.isoformat() + 'Z',
                "cartographicDegrees": row['geometry'].coords[0]
            })
            colors.append({
                "interval": time.isoformat() + 'Z',
                "rgba": [0, int(255 * (flow - 20) / 10), 255, 255]
            })

        czml.append({
            "id": f"river_{row['Index']}",
            "name": f"River {row['Index']}",
            "polyline": {
                "positions": {
                    "epoch": start_time.isoformat() + 'Z',
                    "cartographicDegrees": positions
                },
                "material": {
                    "polylineOutline": {
                        "color": {
                            "epoch": start_time.isoformat() + 'Z',
                            "rgba": colors
                        },
                        "outlineColor": {
                            "rgba": [255, 255, 255, 255]
                        },
                        "outlineWidth": 2
                    }
                },
                "width": 5
            }
        })

    return czml

# 设置开始时间
start_time = datetime.utcnow()

# 生成CZML数据
czml_data = generate_czml(gdf, start_time)

# 保存为CZML文件
with open('output/river_flow.czml', 'w') as f:
    json.dump(czml_data, f, indent=2)