import sys
import numpy as np
import pandas as pd


def generate_normal_distribution(rainfall_total, max_intensity, duration):
    """
    生成正态分布的降雨数据
    :param rainfall_total: 降雨总量
    :param max_intensity: 最大雨强
    :param duration: 降雨时间（分钟）
    :return: 降雨数据列表
    """
    mean = rainfall_total / duration
    std_dev = (max_intensity - mean) / 3
    rainfall = np.random.normal(mean, std_dev, duration)
    rainfall = np.clip(rainfall, 0, None)  # 确保降雨强度非负
    rainfall = rainfall * rainfall_total / rainfall.sum()  # 归一化
    return rainfall


def generate_increasing_distribution(rainfall_total, max_intensity, duration):
    """
    生成逐渐上升的降雨数据
    :param rainfall_total: 降雨总量
    :param max_intensity: 最大雨强
    :param duration: 降雨时间（分钟）
    :return: 降雨数据列表
    """
    rainfall = np.linspace(0, max_intensity, duration)
    rainfall = rainfall * rainfall_total / rainfall.sum()  # 归一化
    return rainfall


def generate_uniform_distribution(rainfall_total, max_intensity, duration):
    """
    生成平均分布的降雨数据
    :param rainfall_total: 降雨总量
    :param max_intensity: 最大雨强
    :param duration: 降雨时间（分钟）
    :return: 降雨数据列表
    """
    rainfall = np.full(duration, rainfall_total / duration)
    return rainfall


def generate_fluctuating_distribution(rainfall_total, max_intensity, duration):
    """
    生成波动分布的降雨数据
    :param rainfall_total: 降雨总量
    :param max_intensity: 最大雨强
    :param duration: 降雨时间（分钟）
    :return: 降雨数据列表
    """
    rainfall = np.random.uniform(0, max_intensity, duration)
    rainfall = rainfall * rainfall_total / rainfall.sum()  # 归一化
    return rainfall


def save_to_file(rainfall, filename):
    """
    将降雨数据保存到文件
    :param rainfall: 降雨数据列表
    :param filename: 文件名
    """
    df = pd.DataFrame({'rainfall': rainfall})
    df.to_csv(filename, index=False)


if __name__ == "__main__":
    # 检查参数数量（sys.argv[0] 是脚本名称）
    if len(sys.argv) != 6:
        print("错误：请提供5个参数！")
        print(
            "用法：python 脚本名.py <参数1-csv文件名> <参数2-降雨模式:正态分布|平均分布|波动平均分布|持续上升> <参数3-降雨总量> <参数4-最大雨强> <参数5-降雨时间(分钟)>")
        sys.exit(1)  # 非零退出码表示错误

    rainfall_total = float(sys.argv[3])
    max_intensity = float(sys.argv[4])
    duration = int(sys.argv[5])

    match sys.argv[2]:
        case "正态分布":
            normal_rainfall = generate_normal_distribution(rainfall_total, max_intensity, duration)
            save_to_file(normal_rainfall, sys.argv[1])
        case "持续上升":
            increasing_rainfall = generate_increasing_distribution(rainfall_total, max_intensity, duration)
            save_to_file(increasing_rainfall, sys.argv[1])
        case "平均分布":
            uniform_rainfall = generate_uniform_distribution(rainfall_total, max_intensity, duration)
            save_to_file(uniform_rainfall, sys.argv[1])
        case "波动平均分布":
            fluctuating_rainfall = generate_fluctuating_distribution(rainfall_total, max_intensity, duration)
            save_to_file(fluctuating_rainfall, sys.argv[1])
        case _:  # 正态分布，默认分支
            normal_rainfall = generate_normal_distribution(rainfall_total, max_intensity, duration)
            save_to_file(normal_rainfall, sys.argv[1])

    # rainfall_total = 100  # 降雨总量
    # max_intensity = 10  # 最大雨强
    # duration = 120  # 降雨时间（分钟）

    # 生成不同分布的降雨数据
    # normal_rainfall = generate_normal_distribution(rainfall_total, max_intensity, duration)
    # increasing_rainfall = generate_increasing_distribution(rainfall_total, max_intensity, duration)
    # uniform_rainfall = generate_uniform_distribution(rainfall_total, max_intensity, duration)
    # fluctuating_rainfall = generate_fluctuating_distribution(rainfall_total, max_intensity, duration)

    # 保存到文件
    # save_to_file(normal_rainfall, 'normal_rainfall.csv')
    # save_to_file(increasing_rainfall, 'increasing_rainfall.csv')
    # save_to_file(uniform_rainfall, 'uniform_rainfall.csv')
    # save_to_file(fluctuating_rainfall, 'fluctuating_rainfall.csv')