# 原始文件 TIDI行星波数据处理.py

import os
import numpy as np
from scipy.io import loadmat
import pandas as pd

from CONSTANT import DATA_BASEPATH

# 第156行代码，纬度可以自己选择，但是都是每5°做一个选择
# 70km-120km每2.5km一个，共21个，第六个就是82.5km高度的情况，高度也可以自己选择

HEIGHTS_CONSTANT = [
    70, 72.5, 75, 77.5, 80, 82.5, 85, 87.5, 90, 92.5, 95, 97.5, 100, 102.5, 105, 107.5, 110, 112.5, 115, 117.5, 120
]


def extract_tidi_planetw_data(
    year,
    input_height: float = 82.5,
    lat_range: tuple = (0, 5),
):
    # 初始化空列表来存储每天的数据
    if lat_range[1] - lat_range[0] != 5:
        raise ValueError("lat_range must be 5°")
    if input_height not in HEIGHTS_CONSTANT:
        raise ValueError(f"input_height must be in {HEIGHTS_CONSTANT}")
    fixed_height_index = HEIGHTS_CONSTANT.index(input_height)
    height_list = []
    lat_list = []
    lon_list = []
    vmeridional_list = []
    vzonal_list = []

    # 文件路径
    base_path = f"{DATA_BASEPATH.tidi}/{year}/"
    # if dir not exist, raise error
    if not os.path.exists(base_path):
        raise FileNotFoundError(f"{base_path} not found")

    # 循环读取从第1天到第365天的数据
    for day in range(1, 366):  # 365天数据，确保循环次数为365
        # 构建文件名（使用三位数格式）
        day_str = f"{day:03d}"  # 格式化数字为三位数，前面补0
        height_file = f"{day_str}_Height.mat"
        lat_file = f"{day_str}_Lat.mat"
        lon_file = f"{day_str}_Lon.mat"
        vmeridional_file = f"{day_str}_VMerdional.mat"
        vzonal_file = f"{day_str}_Vzonal.mat"

        # 检查文件是否存在
        try:
            if not os.path.exists(os.path.join(base_path, height_file)):
                raise FileNotFoundError(f"{height_file} not found")
            if not os.path.exists(os.path.join(base_path, lat_file)):
                raise FileNotFoundError(f"{lat_file} not found")
            if not os.path.exists(os.path.join(base_path, lon_file)):
                raise FileNotFoundError(f"{lon_file} not found")
            if not os.path.exists(os.path.join(base_path, vmeridional_file)):
                raise FileNotFoundError(f"{vmeridional_file} not found")
            if not os.path.exists(os.path.join(base_path, vzonal_file)):
                raise FileNotFoundError(f"{vzonal_file} not found")

            # 读取.mat文件
            height_data = loadmat(os.path.join(base_path, height_file))
            lat_data = loadmat(os.path.join(base_path, lat_file))
            lon_data = loadmat(os.path.join(base_path, lon_file))
            vmeridional_data = loadmat(
                os.path.join(base_path, vmeridional_file))
            vzonal_data = loadmat(os.path.join(base_path, vzonal_file))

            # 将数据转换为DataFrame
            height_df = pd.DataFrame(height_data['Height'])
            lat_df = pd.DataFrame(lat_data['Lat'])
            lon_df = pd.DataFrame(lon_data['Lon'])
            vmeridional_df = pd.DataFrame(vmeridional_data['VMerdional'])
            vzonal_df = pd.DataFrame(vzonal_data['Vzonal'])

            # 将每天的数据添加到列表中
            height_list.append(height_df)
            lat_list.append(lat_df)
            lon_list.append(lon_df)
            vmeridional_list.append(vmeridional_df)
            vzonal_list.append(vzonal_df)
        except FileNotFoundError as e:
            print(f"Error: {e}")
            continue  # 跳过当前文件，继续处理其他文件

    # 合并所有天的数据
    height_df = pd.concat(height_list, ignore_index=True)
    lat_df = pd.concat(lat_list, ignore_index=True)
    lon_df = pd.concat(lon_list, ignore_index=True)
    vmeridional_df = pd.concat(vmeridional_list, axis=1)  # 按列合并
    vzonal_df = pd.concat(vzonal_list, axis=1)  # 按列合并

    # 初始化空列表来存储每天的数据
    UTime_list = []
    # 初始化累加的时间偏移量
    time_offset = 0

    # 循环读取从第1天到第365天的数据
    for day in range(1, 366):  # 365天数据
        # 构建文件名（使用三位数格式）
        day_str = f"{day:03d}"  # 格式化数字为三位数，前面补0
        UTime_file = f"{day_str}_UTime.mat"

        # 检查UTime文件是否存在
        try:
            if not os.path.exists(os.path.join(base_path, UTime_file)):
                raise FileNotFoundError(f"{UTime_file} not found")

            # 读取.mat文件
            UTime_data = loadmat(os.path.join(base_path, UTime_file))

            # 将数据转换为DataFrame
            UTime_df = pd.DataFrame(UTime_data['UTime'])

            # 对UTime进行累加处理
            UTime_df += time_offset

            # 将每天的数据添加到列表中
            UTime_list.append(UTime_df)

            # 更新时间偏移量，为下一天增加24小时
            time_offset += 24
        except FileNotFoundError as e:
            print(f"Error: {e}")
            continue  # 跳过当前文件，继续处理其他文件

    # 合并所有天的数据
    UTime_df = pd.concat(UTime_list, ignore_index=True)
    # 将UTime_df的单位从小时转换为天（除以24）
    UTime_df = UTime_df / 24

    # 获取固定高度下的纬度数据
    fixed_height_lat = lat_df.iloc[:, 0].values
    # 获取固定高度下的经度数据
    fixed_height_lon = lon_df.iloc[:, 0].values
    # 获取固定高度下的经向风数据
    fixed_height_vmeridional = vmeridional_df.iloc[fixed_height_index, :].values
    # 获取固定高度下的纬向风数据
    fixed_height_vzonal = vzonal_df.iloc[fixed_height_index, :].values
    # 获取对应的世界时数据（这里直接使用整个UTime数据，你可以根据实际情况判断是否需要处理下格式等）
    fixed_height_utime = UTime_df.iloc[:, 0].values
    # 将这些数据整合到一个新的DataFrame（可选操作，方便后续查看等）
    combined_data = pd.DataFrame({
        'Latitude': fixed_height_lat,
        'Longitude': fixed_height_lon,
        'V_Meridional': fixed_height_vmeridional,
        'V_Zonal': fixed_height_vzonal,
        'UTime': fixed_height_utime
    })

    # 简单打印下整合后的数据前几行查看下内容
    print(combined_data)

    # 确定纬度和经度的范围
    lat_min, lat_max = np.min(fixed_height_lat), np.max(fixed_height_lat)
    lon_min, lon_max = np.min(fixed_height_lon), np.max(fixed_height_lon)

    # 计算网格的边界
    lat_bins = np.linspace(lat_min, lat_max, 10)  # 纬度分为6个网格，需要7个边界
    lon_bins = np.linspace(lon_min, lon_max, 13)  # 经度分为20个网格，需要21个边界

    # 将数据分配到网格中
    grid_data = []
    for i in range(len(fixed_height_lat)):
        lat_idx = np.digitize(fixed_height_lat[i], lat_bins) - 1  # 找到纬度所在的网格索引
        lon_idx = np.digitize(fixed_height_lon[i], lon_bins) - 1  # 找到经度所在的网格索引
        grid_idx = lat_idx * 12 + lon_idx  # 计算网格的唯一索引（9x12=108）
        grid_data.append([fixed_height_lat[i], fixed_height_lon[i], fixed_height_vmeridional[i],
                         fixed_height_vzonal[i], fixed_height_utime[i], grid_idx])
    # 将网格数据转换为DataFrame
    grid_df = pd.DataFrame(grid_data, columns=[
                           'Latitude', 'Longitude', 'V_Meridional', 'V_Zonal', 'UTime', 'Grid_Index'])
    # 打印网格数据的前几行查看
    print(grid_df.head())
    # 筛选纬度在 0 到 5° 范围内的数据
    lat_range_down, lat_range_upper = lat_range

    filtered_df = grid_df[(grid_df['Latitude'] >= lat_range_down) &
                          (grid_df['Latitude'] <= lat_range_upper)]
    # 将经度从度转换为弧度
    filtered_df['Longitude_Radians'] = np.radians(filtered_df['Longitude'])
    # 选择需要输出的列
    output_columns = ['Longitude_Radians', 'UTime', 'V_Zonal']
    # 将数据输出为txt文件
    output_file = base_path+'TIDI_V_Zonal_data.txt'  # 经向风速分量
    filtered_df[output_columns].to_csv(output_file, sep='\t', index=False)

    zonal_data = filtered_df[output_columns]
    # 选择需要输出的列
    output_columns1 = ['Longitude_Radians', 'UTime', 'V_Meridional']
    # 将数据输出为txt文件
    merdi_data = filtered_df[output_columns1]

    return {
        "V_Zonal": zonal_data,
        "V_Meridional": merdi_data
    }
    # output_file1 = base_path+'TIDI_V_Meridional_data.txt'  # 纬向风速分量
    # final_result.to_csv(output_file1, sep='\t', index=False)
    # # 打印结果确认
    # print(f"数据已保存到 {output_file1}")


# # 对网格数据按照Grid_Index进行分组，并计算每个网格的统计数据
# grid_stats = grid_df.groupby('Grid_Index').agg({
#     'Latitude': 'mean',  # 计算纬度的平均值
#     'Longitude': 'mean',  # 计算经度的平均值
#     'V_Meridional': 'mean',  # 计算经向风的平均值、最大值和最小值
#     'V_Zonal': 'mean',  # 计算纬向风的平均值、最大值和最小值
#     'UTime': 'mean'  # 计算UTime的平均值
# }).reset_index()
#
# # 重命名列，使其更具可读性
# grid_stats.columns = ['Grid_Index', 'Mean_Latitude', 'Mean_Longitude', 'Mean_VMeridional',  'Mean_VZonal',  'Mean_UTime']
# # 将经度转换为2π范围
# grid_stats['Mean_Longitude_2pi'] = (grid_stats['Mean_Longitude'] / 360) * (2 * np.pi)
# # 定义纬度的最小值和最大值（起码要40°区间，否则数据太少无法拟合）
# latmin =0
# latmax=40
# # 筛选出Mean_Latitude在lat_min到lat_max之间的数据
# filtered_grid_stats = grid_stats[(grid_stats['Mean_Latitude'] >= latmin) & (grid_stats['Mean_Latitude'] <= latmax)]
#
# # 检查筛选后的数据行数是否小于18
# if filtered_grid_stats.shape[0] < 18:
#     print("无法拟合行星波")
# else:
#     # 打印筛选后的网格数据的前几行查看
#     print(filtered_grid_stats.head())
