# 原始文件 TIDI行星波一年逐日图

# 此代码是对数据处理后的txt数据进行行星波参数提取绘图
# 2006_TIDI_V_Meridional_data.txt也可以做同样处理，一个是经向风提取的行星波，一个是纬向风的

import os
import pandas as pd
import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt

# 解决绘图中中文不能显示的问题
import matplotlib

from CONSTANT import DATA_BASEPATH
from modules.tidi.tidi_planetw_process import extract_tidi_planetw_data
# 设置中文显示和负号正常显示
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 显示中文
matplotlib.rcParams['axes.unicode_minus'] = False  # 正常显示负号
# 读取一年的数据文件
# 设置初始参数
# initial_guess = [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0]  # v0, a1, b1, a2, b2, a3, b3
initial_guess = [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5,
                 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]  # 9个 a 和 9个 b 参数

# 设置参数界限
bounds = (
    [0, -np.inf, 0, -np.inf, 0, -np.inf, 0, -np.inf, 0, -np.inf,
        0, -np.inf, 0, -np.inf, 0, -np.inf, 0, -np.inf],  # 下界
    [np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf,
     np.inf, np.inf, np.inf])  # 上界


# 定义拟合函数

def TidiPlotPlanetWDaily(
        mode,
        year,
        k,
        input_height,
        lat_range=(0, 5),
        T=10
):
    # 用于存储拟合参数结果的列表
    all_fit_results = []

    # 设置最小数据量的阈值
    min_data_points = 36
    if mode != 'V_Zonal' and mode != 'V_Meridional':
        raise ValueError('mode must be V_Zonal or V_Meridional')
    if k not in list(range(-4, 5)):
        raise ValueError('k must be in range(-4, 5)')
    data_cache_path = f'{DATA_BASEPATH.tidi}/{year}/TIDI_{mode}h{input_height}r{lat_range[0]}-{lat_range[1]}_data.parquet'
    # if cache exist, read from cache
    if os.path.exists(data_cache_path):
        df = pd.read_parquet(data_cache_path)
    else:
        df = extract_tidi_planetw_data(
            year, input_height, lat_range
        )[mode]
        df.to_parquet(data_cache_path)

    # 删除有 NaN 的行
    # V_Meridional
    df = df.dropna(subset=[mode])
    # 进行多个时间窗口的拟合
    # T应该取5、10、16
    # todo:对于5日波，滑动窗口选择15天

    def u_func(x, *params):
        a1, b1, a2, b2, a3, b3, a4, b4, a5, b5, a6, b6, a7, b7, a8, b8, a9, b9 = params
        return (
            a1 * np.sin((2 * np.pi / T) * t - 4 * x + b1)
            + a2 * np.sin((2 * np.pi / T) * t - 3 * x + b2)
            + a3 * np.sin((2 * np.pi / T) * t - 2 * x + b3)
            + a4 * np.sin((2 * np.pi / T) * t - x + b4)
            + a5 * np.sin((2 * np.pi / T) * t + b5)
            + a6 * np.sin((2 * np.pi / T) * t + x + b6)
            + a7 * np.sin((2 * np.pi / T) * t + 2 * x + b7)
            + a8 * np.sin((2 * np.pi / T) * t + 3 * x + b8)
            + a9 * np.sin((2 * np.pi / T) * t + 4 * x + b9)
        )
    for start_day in range(0, 365-3*T):  # 最后一个窗口为[351, 366]
        end_day = start_day + 3 * T  # 每个窗口的结束时间为 start_day + 3*T

        # 选择当前窗口的数据
        df_8 = df[(df['UTime'] >= start_day) & (df['UTime'] <= end_day)]

        # 检查当前窗口的数据量
        if len(df_8) < min_data_points:
            # 输出数据量不足的警告
            print(f"数据量不足，无法拟合：{start_day} 到 {end_day}，数据点数量：{len(df_8)}")
            # 将拟合参数设置为 NaN
            all_fit_results.append([np.nan] * 18)
            continue  # 跳过当前时间窗口，继续下一个窗口

        # 提取时间、经度、温度数据
        t = np.array(df_8['UTime'])  # 时间
        x = np.array(df_8['Longitude_Radians'])  # 经度弧度制
        temperature = np.array(df_8[mode])  # 经向风速，因变量

        # 用T进行拟合
        popt, pcov = curve_fit(u_func, x, temperature,
                               p0=initial_guess, bounds=bounds, maxfev=50000)

        # 将拟合结果添加到列表中
        all_fit_results.append(popt)

    # 将结果转换为DataFrame
    columns = ['a1', 'b1', 'a2', 'b2', 'a3', 'b3', 'a4', 'b4',
               'a5', 'b5', 'a6', 'b6', 'a7', 'b7', 'a8', 'b8', 'a9', 'b9']
    fit_df = pd.DataFrame(all_fit_results, columns=columns)  # fit_df即为拟合的参数汇总

    # -------------------------------画图----------------------------
    # a1-a9,对应波数-4、-3、-2、-1、0、1、2、3、4的行星波振幅
    a_columns = ['a1', 'a2', 'a3', 'a4', 'a5', 'a6', 'a7', 'a8', 'a9']
    k_values = list(range(-4, 5))  # 从 -4 到 4

    # 创建一个字典映射 k 值到 a_columns
    k_to_a = {f'k={k}': a for k, a in zip(k_values, a_columns)}

    # 获取索引并转换为 numpy 数组
    x_values = fit_df.index.to_numpy()

    # 对每一列生成独立的图
    col = k_to_a[f'k={k}']
    plt.figure(figsize=(8, 6))  # 创建新的图形
    plt.plot(x_values, fit_df[col].values)
    plt.title(f'k={k} 振幅图')
    plt.xlabel('Day')
    plt.ylabel('振幅')

    # 设置横坐标的动态调整
    adjusted_x_values = x_values + (3 * T + 1) / 2
    if len(adjusted_x_values) > 50:
        step = 30
        tick_positions = adjusted_x_values[::step]  # 选择每30个点
        tick_labels = [f'{int(val)}' for val in tick_positions]
    else:
        tick_positions = adjusted_x_values
        tick_labels = [f'{int(val)}' for val in tick_positions]

    plt.xticks(ticks=tick_positions, labels=tick_labels)

    # plt.show()  # 显示图形


if __name__ == '__main__':
    TidiPlotPlanetWDaily('V_Zonal', 2015, 1)
    TidiPlotPlanetWDaily('V_Meridional', 2015, 1)

# # 用于存储拟合参数结果的列表
# all_fit_results = []
#
# # 设置最小数据量的阈值
# min_data_points = 36
#
# # 进行多个时间窗口的拟合
# # todo:对于5日波，滑动窗口选择15天
# for start_day in range(0, 351):  # 最后一个窗口为[351, 366]
#     end_day = start_day + 15  # 设置每个窗口的结束时间
#
#     # 选择当前窗口的数据
#     df_8 = df[(df['UTime'] >= start_day) & (df['UTime'] <= end_day)]
#
#     # 检查当前窗口的数据量
#     if len(df_8) < min_data_points:
#         # 输出数据量不足的警告
#         print(f"数据量不足，无法拟合：{start_day} 到 {end_day}，数据点数量：{len(df_8)}")
#         # 将拟合参数设置为 NaN
#         all_fit_results.append([np.nan] * 18)
#         continue  # 跳过当前时间窗口，继续下一个窗口
#
#
#     # 提取时间、经度、温度数据
#     t = np.array(df_8['UTime'])  # 时间
#     x = np.array(df_8['Longitude_Radians'])  # 经度弧度制
#     temperature = np.array(df_8['V_Zonal'])  # 经向风速，因变量
#
#     # 用T=5进行拟合
#     for T in [5]:
#         popt, pcov = curve_fit(u_func, x, temperature, p0=initial_guess, bounds=bounds, maxfev=50000)
#
#         # 将拟合结果添加到列表中
#         all_fit_results.append(popt)
#
# # 将结果转换为DataFrame
# columns = ['a1', 'b1', 'a2', 'b2', 'a3', 'b3', 'a4', 'b4', 'a5', 'b5', 'a6', 'b6', 'a7', 'b7', 'a8', 'b8', 'a9', 'b9']
# fit_df = pd.DataFrame(all_fit_results, columns=columns) # fit_df即为拟合的参数汇总
#
# # -------------------------------画图----------------------------
# # 定义 a1 到 a9 的列名
# a_columns = ['a1', 'a2', 'a3', 'a4', 'a5', 'a6', 'a7', 'a8', 'a9']
# a_values = fit_df[a_columns].values
#
# # 获取索引并转换为 numpy 数组
# x_values = fit_df.index.to_numpy()
#
# # 对每一列生成独立的图
# for i, col in enumerate(a_columns):
#     plt.figure(figsize=(8, 6))  # 创建新的图形
#     plt.plot(x_values, a_values[:, i])
#     plt.title(f'{col} 振幅图')
#     plt.xlabel('Day')
#     plt.ylabel('振幅')
#     plt.show()  # 显示图形

# # 提取 a1 到 a9 列的数据
# a_columns = ['a1', 'a2', 'a3', 'a4', 'a5', 'a6', 'a7', 'a8', 'a9']
# a_values = fit_df[a_columns].values
# # 获取索引并转换为 numpy 数组
# x_values = fit_df.index.to_numpy()
#
# # 创建一个图形和子图
# plt.figure(figsize=(10, 6))
#
# # 绘制每一条线
# for i, col in enumerate(a_columns):
#     plt.plot(x_values, a_values[:, i], label=col)
#
# # 添加标题、标签和图例
# plt.title('a1 to a9 振幅图')
# plt.xlabel('Day')
# plt.ylabel('振幅')
# plt.legend()
# # 显示图形
# plt.tight_layout()
# plt.show()
