# 原始文件 saber行星波处理.py
# 高度、纬度可以指定，得到指定高度纬度下的数据，高度在101行，一般输入70、90、110，纬度在115行，纬度一般输入-20、30、60

import os
import netCDF4 as nc
import numpy as np
import netCDF4 as nc
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import pandas as pd
from scipy.optimize import curve_fit
from scipy.optimize import least_squares

from CONSTANT import DATA_BASEPATH
# 定义月份的英文名称
months = [
    "January", "February", "March", "April", "May", "June",
    "July", "August", "September", "October", "November", "December"
]


def saber_planetw_process(year):
    # 定义文件路径模板和年份

    base_path = DATA_BASEPATH.saber + \
        "/{year}/SABER_Temp_O3_{month}{year}_v2.0.nc"

    cache_path = DATA_BASEPATH.saber + \
        f"/{year}/gravity_result.parquet"

    if os.path.exists(cache_path):
        # load from cache
        return pd.read_parquet(cache_path)

    # 初始化空数组来存储拼接后的数据
    tplongitude = None
    tplatitude = None
    ktemp = None
    tpaltitude = None
    date = None
    tpSolarLT = None

    # 循环遍历每个月
    for month_name in months:
        # 构造文件路径
        file_path = base_path.format(month=month_name, year=year)

        # 打开文件并读取数据
        dataset = nc.Dataset(file_path, 'r')
        tplongitude_month = dataset.variables['tplongitude'][:]
        tplatitude_month = dataset.variables['tplatitude'][:]
        ktemp_month = dataset.variables['ktemp'][:]
        tpaltitude_month = dataset.variables['tpaltitude'][:]
        date_month = dataset.variables['date'][:]
        tpSolarLT_month = dataset.variables['tpSolarLT'][:]
        dataset.close()

        # 动态获取两个数组的较小维度，并截取较大的数组以匹配较小的维度
        if tplatitude is None:
            # 如果是第一个月，初始化数组
            tplongitude = tplongitude_month
            tplatitude = tplatitude_month
            ktemp = ktemp_month
            tpaltitude = tpaltitude_month
            date = date_month
            tpSolarLT = tpSolarLT_month
        else:
            min_dim = min(tplatitude.shape[1], tplatitude_month.shape[1])
            if tplatitude.shape[1] < tplatitude_month.shape[1]:
                tplatitude_month = tplatitude_month[:, :min_dim]
                ktemp_month = ktemp_month[:, :min_dim]
                tpaltitude_month = tpaltitude_month[:, :min_dim]
                tpSolarLT_month = tpSolarLT_month[:, :min_dim]
                tplongitude_month = tplongitude_month[:, :min_dim]
            else:
                tplatitude = tplatitude[:, :min_dim]
                ktemp = ktemp[:, :min_dim]
                tpaltitude = tpaltitude[:, :min_dim]
                tpSolarLT = tpSolarLT[:, :min_dim]
                tplongitude = tplongitude[:, :min_dim]

            # 拼接数据
            tplongitude = np.concatenate(
                (tplongitude, tplongitude_month), axis=0)
            tplatitude = np.concatenate((tplatitude, tplatitude_month), axis=0)
            ktemp = np.concatenate((ktemp, ktemp_month), axis=0)
            tpaltitude = np.concatenate((tpaltitude, tpaltitude_month), axis=0)
            date = np.concatenate((date, date_month), axis=0)
            tpSolarLT = np.concatenate((tpSolarLT, tpSolarLT_month), axis=0)

    # ----------数据处理------------------------------------------------------
    # 时间处理
    time_hour = tpSolarLT / 1000 / 3600
    # 假设date的格式为YYYYDDD，其中DDD是一年中的第几天
    dates = [datetime.strptime(f'{d:07d}', '%Y%j') for d in date]
    # 创建一个辅助函数来调整time_hour

    def adjust_time_hour(time_hour, dates):
        # 初始化一个和time_hour形状相同的数组，用于存储调整后的时间
        adjusted_time_hour = np.zeros_like(time_hour)
        # 获取第一个日期作为参考点
        first_date = dates[0]
        # 遍历dates和time_hour，根据日期调整时间
        for i, (dt, th) in enumerate(zip(dates, time_hour)):
            # 计算当前日期和第一个日期之间的天数差
            delta_days = (dt - first_date).days
            # 根据天数差调整时间
            adjusted_time_hour[i] = th + delta_days * 24

        return adjusted_time_hour

    # 调整time_hour
    adjusted_time_hour = adjust_time_hour(time_hour, dates)
    # 将小时转换为天
    adjusted_time_day = adjusted_time_hour / 24.0  # 转换为天
    # 找到第一个事件中与70最接近的高度的索引
    # 输入高度
    target_h = 70
    first_event_altitudes = tpaltitude[0]
    diff = np.abs(first_event_altitudes - target_h)
    closest_height_index = np.argmin(diff)
    # 读取每个事件对应closest_height_index下的纬度、tplongitude和adjusted_time_day数据
    selected_latitude = tplatitude[:, closest_height_index]
    selected_tplongitude = tplongitude[:, closest_height_index]
    selected_adjusted_time_day = adjusted_time_day[:, closest_height_index]
    selected_ktemp = ktemp[:, closest_height_index]
    # 将时间数据转换为整数天数，方便按天分组
    int_days = np.array([int(day) for day in selected_adjusted_time_day])
    # 1. 创建经度分组（0到360°，每20°一个区间）
    longitude_bins = np.arange(0, 361, 20)
    # 2. 创建纬度目标值
    target_lat = 60  # 假设目标纬度为60°
    # 3. 用于存储结果
    closest_event_indices = []

    # 4. 遍历每天的所有数据（按天分组），并查找与每个经度区间最接近的经度事件索引
    for day in np.unique(int_days):  # 遍历每个独特的天数
        # 获取当前天数下的所有经度和纬度
        current_day_indices = np.where(int_days == day)[0]
        current_longitudes = selected_tplongitude[current_day_indices]
        current_latitudes = selected_latitude[current_day_indices]

        # 初始化列表用于存储该天的最接近经度和纬度的索引
        closest_indices_for_day = []

        # 遍历每个经度区间（20°为一个区间）
        for lon_bin_start in longitude_bins[:-1]:
            lon_bin_end = lon_bin_start + target_lat
            # 找到当前经度区间内的经度
            indices_in_bin = np.where((current_longitudes >= lon_bin_start) & (
                current_longitudes < lon_bin_end))[0]
            # 如果有经度落入该区间，找到与目标纬度20°最接近的纬度
            if len(indices_in_bin) > 0:
                # 计算纬度与目标纬度的差值
                latitudes_in_bin = current_latitudes[indices_in_bin]
                diff_latitude = np.abs(latitudes_in_bin - target_lat)
                # 找到与目标纬度最接近的索引
                closest_latitude_index = indices_in_bin[np.argmin(
                    diff_latitude)]
                # 将该经度区间和最接近纬度的索引添加到结果中
                closest_indices_for_day.append(
                    current_day_indices[closest_latitude_index])
        # 将该天的所有最接近的经度和纬度索引添加到总结果中
        closest_event_indices.extend(closest_indices_for_day)

    # 4. 获取这些事件的相关信息
    closest_latitudes = selected_latitude[closest_event_indices]
    closest_longitudes = selected_tplongitude[closest_event_indices]
    closest_times = selected_adjusted_time_day[closest_event_indices]
    closest_temperatures = selected_ktemp[closest_event_indices]
    # 5. 将经度从度数转换为弧度制
    closest_longitudes_radians = np.radians(closest_longitudes)
    # 将三列数据组合成一个二维数组
    combined_data = np.column_stack(
        (closest_longitudes_radians, closest_times, closest_temperatures))

    df = pd.DataFrame(combined_data, columns=[
                      "Longitude_Radians", "Time", "Temperature"])

    # dump to cache
    df.to_parquet(cache_path)
    return df


# 将数据保存为txt文件
# np.savetxt('combined_data.txt', combined_data, fmt='%f', delimiter='\t',
#            header="Longitude_Radians\tTime\tTemperature", comments='')

# print("数据已保存为 'combined_data.txt'")
