# saber 重力波

from io import BytesIO
import netCDF4 as nc
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from scipy.optimize import curve_fit
# from matplotlib.colors import LinearSegmentedColormap
# 设置字体为支持中文的字体
plt.rcParams['font.family'] = 'SimHei'  # 设置为黑体（需要你的环境中有该字体）
plt.rcParams['axes.unicode_minus'] = False  # 解决负号'-'显示为方块的问题

# ----------------------------------------------------------------------------------------------------------------------------
# 1---打开文件并读取不同变量数据                                                                                              -
# ----------------------------------------------------------------------------------------------------------------------------


def data_nc_load(file_path):

    dataset = nc.Dataset(file_path, 'r')

    # 纬度数据，二维数组形状为(42820,379)   42820为事件，379则为不同高度
    tplatitude = dataset.variables['tplatitude'][:, :]
    tplongitude = dataset.variables['tplongitude'][:, :]   # 经度数据
    tpaltitude = dataset.variables['tpaltitude'][:, :]  # 高度，二维数组形状为(42820,379)
    time = dataset.variables['time'][:, :]  # 二维数组形状为(42820,379)
    date = dataset.variables['date'][:]
    date_time = np.unique(date)  # 输出数据时间信息
    ktemp = dataset.variables['ktemp'][:, :]   # 温度数据，二维数组形状为(42820,379)

    return dataset, tplatitude, tplongitude, tpaltitude, ktemp, time, date, date_time
# ----------------------------------------------------------------------------------------------------------------------------

# ----------------------------------------------------------------------------------------------------------------------------
# 2---筛选某一天、某个纬度和高度范围15个不同cycle的温度数据                                                                    -
# ----------------------------------------------------------------------------------------------------------------------------
# 2-1 读取某一天的所有事件及其对应的纬度数据


def day_data_read(date, day_read, tplatitude):

    events = np.where(date == day_read)[0]   # 读取筛选天的事件编号  4294-5714位置，从0开始编号
    time_events = date[date == day_read]       # 读取筛选天的事件编号  4294-5714位置，从0开始编号
    latitudes = tplatitude[events, 189]         # 输出每个事件中间位置 即第189个经纬度
    df = pd.DataFrame([                       # 创建一个包含事件编号、纬度的列表，共1421个事件
        {'time': time, 'event': event, 'latitude': lat}
        for time, event, lat in zip(time_events, events, latitudes)])

    # print(df.head())  # 打印前几行数据以检查

    return df
# ----------------------------------------------------------------------------------------------------------------------------
# 2-2 将事件按照卫星轨迹周期进行输出处理，并输出落在纬度范围内的每个周期的事件的行号


def data_cycle_identify(df, latitude_min, latitude_max):

    cycles = []  # 存储每个周期的事件编号列表
    # 存储当前周期的事件编号
    current_cycle_events = []
    prev_latitude = None

    # 遍历DataFrame中的每一行以识别周期和筛选事件
    for index, row in df.iterrows():
        current_event = int(row['event'])
        current_latitude = row['latitude']

        if prev_latitude is not None and prev_latitude < 0 and current_latitude >= 0:  # 检查是否是新周期的开始（纬度从负变正，且首次变正）
            # 重置当前周期的事件编号列表
            current_cycle_events = []

        if latitude_min <= current_latitude <= latitude_max:                           # 如果事件的纬度在指定范围内，添加到当前周期的事件编号列表
            current_cycle_events.append(current_event)

        if prev_latitude is not None and prev_latitude >= 0 and current_latitude < 0:   # 检查是否是周期的结束（纬度从正变负）

            # 添加当前周期的事件编号列表到周期列表
            if current_cycle_events:  # 确保当前周期有事件编号
                cycles.append(current_cycle_events)
                current_cycle_events = []  # 重置当前周期的事件编号列表
        prev_latitude = current_latitude

    if current_cycle_events:                                                            # 处理最后一个周期，如果存在的话
        cycles.append(current_cycle_events)

    print(f"一天周期为 {len(cycles)}")
    for cycle_index, cycle in enumerate(cycles, start=0):
        # 屏幕显示每个循环周期的事件
        print(f"周期 {cycle_index} 包含事件个数: {len(cycle)}  具体事件为: {cycle}  ")

    return cycles
# ----------------------------------------------------------------------------------------------------------------------------
# 2-3---按照循环周期合并同周期数据，并输出处理后的温度数据、对应的高度数据


def data_cycle_generate(cycles, ktemp, tpaltitude, altitude_min, altitude_max):

    # 初始化列表存储每个周期的温度
    ktemp_cycles = []
    # 初始化每个循环周期的高度数据
    altitude_cycles = []
    print(f"周期数为 {len(cycles)}")
    if cycles is not None:                                                            # 检查周期是否有事件编号
        for event in cycles:
            # 获取每个周期各个事件的ktemp数据
            ktemp_cycles_events = np.array(ktemp[event, :])
            ktemp_cycles_events[np.logical_or(
                ktemp_cycles_events == -999, ktemp_cycles_events > 999)] = np.NaN  # 缺失值处理，避免影响结果
            ktemp_cycles_mean = np.nanmean(
                ktemp_cycles_events, axis=0)  # 对所有周期的 ktemp 数据取均值

            # altitude_cycles_events =np.array(tpaltitude[event,:])                    # 获取每个周期各个事件的tpaltitude数据
            # altitude_cycles_mean = np.nanmean(altitude_cycles_events, axis=0)        # 对所有周期的 altitude 数据取均值
            # 使用第一个的高度来表征所有的
            altitude_cycles_mean = tpaltitude[event[0], :]
            altitude_indices = np.where((altitude_cycles_mean >= altitude_min) & (
                altitude_cycles_mean <= altitude_max))[0]

            ktemp_cycles.append(np.array(ktemp_cycles_mean[altitude_indices]))
            altitude_cycles.append(
                np.array(altitude_cycles_mean[altitude_indices]))

    # 找到最短列表的长度
    min_length = min(len(arr) for arr in ktemp_cycles)
    # 创建新的列表，将每个子列表截断为最短长度
    ktemp_cycles = np.vstack([arr[:min_length] for arr in ktemp_cycles])
    altitude_cycles = np.vstack([arr[:min_length] for arr in altitude_cycles])

    return ktemp_cycles, altitude_cycles
# ----------------------------------------------------------------------------------------------------------------------------

# ----------------------------------------------------------------------------------------------------------------------------
# 4---高度相同下不同循环周期数据的波拟合和滤波处理                                                                             -
# ----------------------------------------------------------------------------------------------------------------------------


# 对输入数据按照行（即纬向）进行波数为k的滤波,数据为15*157
def fit_wave(ktemp_wn0, k):

    wave_fit = []

    def single_harmonic(x, A, phi, C, k):
        return A * np.sin(2 * np.pi / (15/k) * x + phi) + C

    # 数据转置并对每行进行操作，按照同高度数据进行处理
    for rtemp in ktemp_wn0.T:
        # 为当前高度层创建索引数组
        indices = np.arange(rtemp.size)
        def fit_temp(x, A, phi, C): return single_harmonic(
            x, A, phi, C, k)             # 定义只拟合 A, phi, C 的 lambda 函数，k 固定
        params, params_covariance = curve_fit(
            fit_temp, indices, rtemp)              # 使用 curve_fit 进行拟合
        # 提取拟合参数 A, phi, C
        A, phi, C = params
        # 使用拟合参数计算 wn3
        rtemp1 = single_harmonic(indices, A, phi, C, k)
        # 存储拟合参数和拟合曲线
        wave_fit.append(np.array(rtemp1))
    wave_fit = np.vstack(wave_fit)

    wave_fit = wave_fit.T
    wave_wn = ktemp_wn0-wave_fit

    return wave_fit, wave_wn
# ----------------------------------------------------------------------------------------------------------------------------


# ----------------------------------------------------------------------------------------------------------------------------
# 4---同周期下不同高度数据的波拟合和滤波处理                                                                             -
# ----------------------------------------------------------------------------------------------------------------------------
# 对输入数据按照列（即纬向）进行滤波，滤除波长2和15km以内的波, 数据为15*157
def fft_ifft_wave(ktemp_wn5, lamda_low, lamda_high, altitude_min, altitude_max, lvboin):

    ktemp_fft = []
    ktemp_ifft = []
    ktemp_fft_lvbo = []

    for rtemp in ktemp_wn5:
        # 采样点数或长度
        N = len(rtemp)
        # 采样时间间隔，其倒数等于采用频率，以1km为标准尺度等同于1s，假设波的速度为1km/s
        dt = (altitude_max-altitude_min)/(N-1)
        # 时间序列索引
        n = np.arange(N)
        # # t = altitude_min + n * dt  # 时间向量
        # t = np.round(np.linspace(altitude_min, altitude_max, N),2)
        # 频率索引向量
        f = n / (N * dt)
        # 对输入信号进行傅里叶变换
        y = np.fft.fft(rtemp)

        # f_low = 2*np.pi / lamda_high                                                                  # 定义波长滤波范围（以频率计算）   # 高频截止频率
        # f_high = 2*np.pi / lamda_low

        # 定义波长滤波范围（以频率计算）   # 高频截止频率
        f_low = 1 / lamda_high
        # 低频截止频率
        f_high = 1 / lamda_low

        # 创建滤波后的频域信号
        yy = y.copy()

        # 使用逻辑索引过滤特定频段（未确定）
        if lvboin:
            freq_filter = (f > f_low) & (f < f_high)  # 创建逻辑掩码
        else:
            freq_filter = (f < f_low) | (f > f_high)  # 创建逻辑掩码

        yy[freq_filter] = 0  # 过滤掉指定频段
        yy_ifft = np.real(np.fft.ifft(yy))

        ktemp_fft.append(y)
        # 存储拟合参数和拟合曲线
        ktemp_ifft.append(np.array(yy_ifft))
        ktemp_fft_lvbo.append(yy)

    ktemp_fft = np.vstack(ktemp_fft)
    ktemp_ifft = np.vstack(ktemp_ifft)
    ktemp_fft_lvbo = np.vstack(ktemp_fft_lvbo)

    return ktemp_fft, ktemp_fft_lvbo, ktemp_ifft
# ----------------------------------------------------------------------------------------------------------------------------
# ----------------------------------------------------------------------------------------------------------------------------
# 5---同周期下不同高度数据的BT_z背景位等指标计算                                                                             -
# ----------------------------------------------------------------------------------------------------------------------------
# ----------------------------------------------------------------------------------


def power_indices(ktemp_cycles, ktemp_wn5, ktemp_ifft, altitude_min, altitude_max):

    # 定义Brunt-Väisälä频率计算函数
    def brunt_vaisala_frequency(g, BT_z, c_p, height):
        # 计算位温随高度的变化率
        dBT_z_dz = np.gradient(BT_z, height)
        # 计算 Brunt-Väisälä 频率
        return np.sqrt((g / BT_z) * ((g / c_p) + dBT_z_dz))

    # 定义势能计算函数
    def calculate_gravitational_potential_energy(g, BT_z, N_z, PT_z):
        return 0.5 * ((g / N_z) ** 2) * ((PT_z / BT_z) ** 2)

    # 重力加速度
    g = 9.81
    # 空气比热容
    c_p = 1004.5

    height = np.linspace(altitude_min, altitude_max,
                         ktemp_cycles.shape[1]) * 1000  # 高度
    background = ktemp_cycles - ktemp_wn5

    # 初始化结果矩阵
    N_z_matrix = []
    # 初始化结果矩阵
    PT_z_matrix = []

    # 循环处理background和filtered_perturbation所有行
    for i in range(background.shape[0]):
        BT_z = np.array(background[i])
        # 滤波后的扰动
        PT_z = np.array(ktemp_ifft[i])

        # 调用Brunt-Väisälä频率函数
        N_z = brunt_vaisala_frequency(g, BT_z, c_p, height)
        PT_z = calculate_gravitational_potential_energy(
            g, BT_z, N_z, PT_z)              # 调用势能函数
        N_z_matrix.append(N_z)
        PT_z_matrix.append(PT_z)

    ktemp_Nz = np.vstack(N_z_matrix)
    ktemp_Ptz = np.vstack(PT_z_matrix)
    return ktemp_Nz, ktemp_Ptz
# ----------------------------------------------------------------------------------------------------------------------------
# 5---main程序环节                                                                             -
# ----------------------------------------------------------------------------------------------------------------------------
# 按日处理资料


def day_process_main(file_path, day_read, latitude_min, latitude_max, altitude_min, altitude_max, lamda_low, lamda_high, lvboin):

    dataset, tplatitude, tplongitude, tpaltitude, ktemp, time, date, date_time = data_nc_load(
        file_path)
    # 2018年的第94天 # 4月4日的日期，以年份和年内的第几天表示
    df = day_data_read(date, day_read, tplatitude)
    cycles = data_cycle_identify(df, latitude_min, latitude_max)
    ktemp_cycles, altitude_cycles = data_cycle_generate(
        cycles, ktemp, tpaltitude, altitude_min, altitude_max)
    # 按照纬向计算平均温度 wn0_temp
    ktemp_wn0 = ktemp_cycles - np.mean(ktemp_cycles, axis=0)
    ktemp_fit_wn1, ktemp_wn1 = fit_wave(ktemp_wn0, 1)
    ktemp_fit_wn2, ktemp_wn2 = fit_wave(ktemp_wn1, 2)
    ktemp_fit_wn3, ktemp_wn3 = fit_wave(ktemp_wn2, 3)
    ktemp_fit_wn4, ktemp_wn4 = fit_wave(ktemp_wn3, 4)
    ktemp_fit_wn5, ktemp_wn5 = fit_wave(ktemp_wn4, 5)
    ktemp_fft, ktemp_fft_lvbo, ktemp_ifft = fft_ifft_wave(
        ktemp_wn5, lamda_low, lamda_high, altitude_min, altitude_max, lvboin)
    ktemp_Nz, ktemp_Ptz = power_indices(
        ktemp_cycles, ktemp_wn5, ktemp_ifft, altitude_min, altitude_max)

    return ktemp_cycles, altitude_cycles, ktemp_wn0, ktemp_fit_wn1, ktemp_wn1, ktemp_fit_wn2, ktemp_wn2, ktemp_fit_wn3, ktemp_wn3, ktemp_fit_wn4, ktemp_wn4, ktemp_fit_wn5, ktemp_wn5, ktemp_fft, ktemp_fft_lvbo, ktemp_ifft, ktemp_Nz, ktemp_Ptz
# ----------------------------------------------------------------------------------------------------------------------------
# 按文件中单日处理资料


def day_process_maing(dataset, tplatitude, tplongitude, tpaltitude, ktemp, time, date, date_time, day_read, latitude_min, latitude_max, altitude_min, altitude_max, lamda_low, lamda_high, lvboin):

    # dataset, tplatitude, tplongitude, tpaltitude, ktemp, time, date, date_time = data_nc_load (file_path)
    # 2018年的第94天 # 4月4日的日期，以年份和年内的第几天表示
    df = day_data_read(date, day_read, tplatitude)
    cycles = data_cycle_identify(df, latitude_min, latitude_max)
    ktemp_cycles, altitude_cycles = data_cycle_generate(
        cycles, ktemp, tpaltitude, altitude_min, altitude_max)
    # 按照纬向计算平均温度 wn0_temp
    ktemp_wn0 = ktemp_cycles - np.mean(ktemp_cycles, axis=0)
    ktemp_fit_wn1, ktemp_wn1 = fit_wave(ktemp_wn0, 1)
    ktemp_fit_wn2, ktemp_wn2 = fit_wave(ktemp_wn1, 2)
    ktemp_fit_wn3, ktemp_wn3 = fit_wave(ktemp_wn2, 3)
    ktemp_fit_wn4, ktemp_wn4 = fit_wave(ktemp_wn3, 4)
    ktemp_fit_wn5, ktemp_wn5 = fit_wave(ktemp_wn4, 5)
    ktemp_fft, ktemp_fft_lvbo, ktemp_ifft = fft_ifft_wave(
        ktemp_wn5, lamda_low, lamda_high, altitude_min, altitude_max, lvboin)
    ktemp_Nz, ktemp_Ptz = power_indices(
        ktemp_cycles, ktemp_wn5, ktemp_ifft, altitude_min, altitude_max)

    return ktemp_cycles, altitude_cycles, ktemp_wn0, ktemp_fit_wn1, ktemp_wn1, ktemp_fit_wn2, ktemp_wn2, ktemp_fit_wn3, ktemp_wn3, ktemp_fit_wn4, ktemp_wn4, ktemp_fit_wn5, ktemp_wn5,  ktemp_fft, ktemp_fft_lvbo, ktemp_ifft,  ktemp_Nz, ktemp_Ptz

# ----------------------------------------------------------------------------------------------------------------------------
# 按月处理资料


def mon_process_main(file_path, latitude_min, latitude_max, altitude_min, altitude_max, lamda_low, lamda_high, lvboin):

    dataset, tplatitude, tplongitude, tpaltitude, ktemp, time, date, date_time = data_nc_load(
        file_path)

    ktemp_cycles_mon = []
    altitude_cycles_mon = []
    ktemp_wn0_mon = []
    ktemp_fit_wn1_mon = []
    ktemp_wn1_mon = []
    ktemp_fit_wn2_mon = []
    ktemp_wn2_mon = []
    ktemp_fit_wn3_mon = []
    ktemp_wn3_mon = []
    ktemp_fit_wn4_mon = []
    ktemp_wn4_mon = []
    ktemp_fit_wn5_mon = []
    ktemp_wn5_mon = []
    ktemp_fft_mon = []
    ktemp_fft_lvbo_mon = []
    ktemp_ifft_mon = []
    ktemp_Nz_mon = []
    ktemp_Ptz_mon = []

    for day_read in date_time:
        print(f"读取日期 {day_read}")
        (ktemp_cycles0, altitude_cycles0,
         ktemp_wn00,
         ktemp_fit_wn10, ktemp_wn10,
         ktemp_fit_wn20, ktemp_wn20,
         ktemp_fit_wn30, ktemp_wn30,
         ktemp_fit_wn40, ktemp_wn40,
         ktemp_fit_wn50, ktemp_wn50,
         ktemp_fft0, ktemp_fft_lvbo0, ktemp_ifft0,
         ktemp_Nz0, ktemp_Ptz0) = day_process_maing(
            dataset, tplatitude, tplongitude, tpaltitude, ktemp, time, date, date_time,
            day_read, latitude_min, latitude_max, altitude_min, altitude_max, lamda_low, lamda_high, lvboin)

        ktemp_cycles_mon.append(ktemp_cycles0)
        altitude_cycles_mon.append(altitude_cycles0)
        ktemp_wn0_mon.append(ktemp_wn00)
        ktemp_fit_wn1_mon.append(ktemp_fit_wn50)
        ktemp_wn1_mon.append(ktemp_wn50)
        ktemp_fit_wn2_mon.append(ktemp_fit_wn50)
        ktemp_wn2_mon.append(ktemp_wn50)
        ktemp_fit_wn3_mon.append(ktemp_fit_wn50)
        ktemp_wn3_mon.append(ktemp_wn50)
        ktemp_fit_wn4_mon.append(ktemp_fit_wn50)
        ktemp_wn4_mon.append(ktemp_wn50)
        ktemp_fit_wn5_mon.append(ktemp_fit_wn50)
        ktemp_wn5_mon.append(ktemp_wn50)
        ktemp_fft_mon.append(ktemp_fft0)
        ktemp_fft_lvbo_mon.append(ktemp_fft_lvbo0)
        ktemp_ifft_mon.append(ktemp_ifft0)
        ktemp_Nz_mon.append(ktemp_Nz0)
        ktemp_Ptz_mon.append(ktemp_Ptz0)

    return [date_time, ktemp_cycles_mon, altitude_cycles_mon,                                 # 某月份 逐日时间 每日不同循环周期温度  不同循环周期高度数据
            # 距平扰动温度
            ktemp_wn0_mon,
            # 波数为1的拟合温度  滤波后的扰动温度
            ktemp_fit_wn1_mon, ktemp_wn1_mon,
            # 波数为2的拟合温度  滤波后的扰动温度
            ktemp_fit_wn2_mon, ktemp_wn2_mon,
            # 波数为3的拟合温度  滤波后的扰动温度
            ktemp_fit_wn3_mon, ktemp_wn3_mon,
            # 波数为4的拟合温度  滤波后的扰动温度
            ktemp_fit_wn4_mon, ktemp_wn4_mon,
            # 波数为5的拟合温度  滤波后的扰动温度
            ktemp_fit_wn5_mon, ktemp_wn5_mon,
            # 滤波0-5后扰动温度傅里叶频谱分析  滤除波长为2km-15km内后频谱分析 滤波后的扰动温度
            ktemp_fft_mon, ktemp_fft_lvbo_mon, ktemp_ifft_mon,
            ktemp_Nz_mon, ktemp_Ptz_mon]                                                      # 滤波后NZ、PTZ重力波势能指标计算

# -----------------------------------------------------------------------------------------------------------------------------
# 6-主要统计分析图绘制                                                                                                         -
# -----------------------------------------------------------------------------------------------------------------------------
# 6-1 示例的逐层滤波效果图---不同波数  曲线图


day_fit_wave_modes = ["k=1", "k=2", "k=3", "k=4", "k=5", "滤波"]


def day_fit_wave_plot(height_no, ktemp_wn0, ktemp_fit_wn1, ktemp_wn1, ktemp_fit_wn2, ktemp_wn2, ktemp_fit_wn3, ktemp_wn3, ktemp_fit_wn4, ktemp_wn4, ktemp_fit_wn5, ktemp_wn5, mode):

    N = len(ktemp_wn0[:, height_no])
    # 循环周期索引
    x = np.arange(N)

    y1_1 = ktemp_wn0[:, height_no]
    y1_2 = ktemp_fit_wn1[:, height_no]

    y2_1 = ktemp_wn1[:, height_no]
    y2_2 = ktemp_fit_wn2[:, height_no]

    y3_1 = ktemp_wn2[:, height_no]
    y3_2 = ktemp_fit_wn3[:, height_no]

    y4_1 = ktemp_wn3[:, height_no]
    y4_2 = ktemp_fit_wn4[:, height_no]

    y5_1 = ktemp_wn4[:, height_no]
    y5_2 = ktemp_fit_wn5[:, height_no]

    y6 = ktemp_wn5[:, height_no]

    if mode not in ["k=1", "k=2", "k=3", "k=4", "k=5", "滤波"]:
        raise ValueError(
            "mode 参数应为 'k=1', 'k=2', 'k=3', 'k=4', 'k=5', 'lvbo' 中的一个")

    plt.figure(figsize=(16, 10))  # 调整图形大小
    # 原始信号的时间序列
    if mode == "k=1":
        plt.plot(x, y1_1, label='原始信号')
        plt.plot(x, y1_2, label='拟合信号', linestyle='--')
        plt.title('（a）波数k=1')
        plt.xlabel('Cycles', labelpad=10)  # 增加标签间距
        plt.ylabel('温度 (°C)', labelpad=10)  # 增加标签间距
        plt.legend()
        plt.tight_layout()  # 调整子图参数以适应图形区域

    if mode == "k=2":
        plt.plot(x, y2_1, label='原始信号')
        plt.plot(x, y2_2, label='拟合信号', linestyle='--')
        plt.title('（b）波数k=2')
        plt.xlabel('Cycles', labelpad=10)  # 增加标签间距
        plt.ylabel('温度 (°C)', labelpad=10)  # 增加标签间距
        plt.legend()
        plt.tight_layout()  # 调整子图参数以适应图形区域

    if mode == "k=3":
        plt.plot(x, y3_1, label='原始信号')
        plt.plot(x, y3_2, label='拟合信号', linestyle='--')
        plt.title('（c）波数k=3')
        plt.xlabel('Cycles', labelpad=10)  # 增加标签间距
        plt.ylabel('温度 (°C)', labelpad=10)  # 增加标签间距
        plt.legend()
        plt.tight_layout()  # 调整子图参数以适应图形区域

    if mode == "k=4":
        plt.plot(x, y4_1, label='原始信号')
        plt.plot(x, y4_2, label='拟合信号', linestyle='--')
        plt.title('（d）波数k=4')
        plt.xlabel('Cycles', labelpad=10)  # 增加标签间距
        plt.ylabel('温度 (°C)', labelpad=10)  # 增加标签间距
        plt.legend()
        plt.tight_layout()  # 调整子图参数以适应图形区域

    if mode == "k=5":
        plt.plot(x, y5_1, label='原始信号')
        plt.plot(x, y5_2, label='拟合信号', linestyle='--')
        plt.title('（e）波数k=5')
        plt.xlabel('Cycles', labelpad=10)  # 增加标签间距
        plt.ylabel('温度 (°C)', labelpad=10)  # 增加标签间距
        plt.legend()
        plt.tight_layout()  # 调整子图参数以适应图形区域

    if mode == "滤波":
        plt.plot(x, y6, label='滤波信号')
        plt.title('（f）滤波1-5后信号')
        plt.xlabel('Cycles', labelpad=10)  # 增加标签间距
        plt.ylabel('温度 (°C)', labelpad=10)  # 增加标签间距
        plt.tight_layout()  # 调整子图参数以适应图形区域


def day_fit_wave_plotg(height_no, ktemp_wn0, ktemp_fit_wn1, ktemp_wn1, ktemp_fit_wn2,
                       ktemp_wn2, ktemp_fit_wn3, ktemp_wn3, ktemp_fit_wn4,
                       ktemp_wn4, ktemp_fit_wn5, ktemp_wn5):

    N = len(ktemp_wn0[:, height_no])
    x = np.arange(N)

    y1_1, y1_2 = ktemp_wn0[:, height_no], ktemp_fit_wn1[:, height_no]
    y2_1, y2_2 = ktemp_wn1[:, height_no], ktemp_fit_wn2[:, height_no]
    y3_1, y3_2 = ktemp_wn2[:, height_no], ktemp_fit_wn3[:, height_no]
    y4_1, y4_2 = ktemp_wn3[:, height_no], ktemp_fit_wn4[:, height_no]
    y5_1, y5_2 = ktemp_wn4[:, height_no], ktemp_fit_wn5[:, height_no]
    y6 = ktemp_wn5[:, height_no]

    plt.figure(figsize=(16, 10))

    y_limits = (min(min(y1_1), min(y2_1), min(y3_1), min(y4_1), min(y5_1), min(y6)),
                max(max(y1_1), max(y2_1), max(y3_1), max(y4_1), max(y5_1), max(y6)))

    for i, (y1, y2) in enumerate([(y1_1, y1_2), (y2_1, y2_2), (y3_1, y3_2),
                                  (y4_1, y4_2), (y5_1, y5_2), (y6, None)]):
        plt.subplot(2, 3, i + 1)
        plt.plot(x, y1, label='原始信号')
        if y2 is not None:
            plt.plot(x, y2, label='拟合信号', linestyle='--')
        plt.title(f'（{"abcdef"[i]}）波数k={i + 1 if i < 5 else "滤波"}')
        plt.xlabel('Cycles', labelpad=10)
        plt.ylabel('温度 (°C)', labelpad=10)
        plt.legend()
        plt.xticks(x)  # 设置横坐标为整数
        plt.ylim(y_limits)  # 设置统一纵坐标范围
        plt.tight_layout()

    plt.subplots_adjust(top=0.8, bottom=0.2, left=0.1,
                        right=0.8, hspace=0.3, wspace=0.2)
    plt.show()


# -----------------------------------------------------------------------------------------------------------------------------
# 6-2 示例的高度滤波处理--不同循环周期  曲线图
day_fft_ifft_modes = ["原始信号", "fft", "ifft", "滤波", "比较"]


def day_fft_ifft_plot(cycle_no, ktemp_wn5, ktemp_fft, ktemp_ifft, altitude_min, altitude_max, lamda_low, lamda_high, ktemp_fft_lvbo, mode):

    N = len(ktemp_wn5[cycle_no, :])
    # 采样时间间隔，其倒数等于采用频率，以1km为标准尺度等同于1s，假设波的速度为1km/s
    dt = (altitude_max-altitude_min)/(N-1)
    # 时间序列索引
    n = np.arange(N)
    f = n / (N * dt)
    t = np.round(np.linspace(altitude_min, altitude_max, N), 2)

    # 原始扰动温度
    x = ktemp_wn5[cycle_no, :]
    # 傅里叶变换频谱分析
    y = ktemp_fft[cycle_no, :]
    # 滤波后的傅里叶变换频谱分析
    yy = ktemp_fft_lvbo[cycle_no, :]
    # 傅里叶逆变换后的扰动温度
    yyy = ktemp_ifft[cycle_no, :]

    if mode not in ["原始信号", "fft", "ifft", "滤波", "比较"]:
        raise ValueError(
            "mode 参数应为 '原始信号', 'fft', 'ifft', 'lv bo' 中的一个")

        plt.figure(figsize=(15, 10))  # 调整图形大小
        # 原始信号的时间序列
    if mode == "原始信号":
        plt.plot(t, x)
        plt.title('（a）原始信号')
        plt.xlabel('高度 (km)', labelpad=10)  # 增加标签间距
        plt.ylabel('温度 (°C)', labelpad=10)  # 增加标签间距
        # 原始振幅谱
    if mode == "fft":
        plt.plot(f, np.abs(y) * 2 / N)
        plt.title('（b））原始振幅谱')
        plt.xlabel('频率/Hz', labelpad=10)  # 增加标签间距
        plt.ylabel('振幅', labelpad=10)  # 增加标签间距

        # 通过IFFT回到时间域
    if mode == "ifft":
        plt.plot(t, yyy)
        plt.title('（c））傅里叶逆变换')
        plt.xlabel('高度 (km)', labelpad=10)  # 增加标签间距
        plt.ylabel('温度 (°C)', labelpad=10)  # 增加标签间距

    if mode == "滤波":
        plt.plot(f, np.abs(yy) * 2 / N)
        plt.title(f'（d）滤除波长 < {lamda_low} km, > {lamda_high} km的波')
        plt.xlabel('频率/Hz', labelpad=10)  # 增加标签间距
        plt.ylabel('振幅', labelpad=10)  # 增加标签间距

    # 调整子图之间的边距

    # 绘制原始信号与滤波后信号
    if mode == '比较':
        plt.plot(t, x, label='原始信号')
        plt.plot(t, yyy, label='滤波后信号', linestyle='--')
        plt.title('信号比较')
        plt.xlabel('高度 (km)', labelpad=10)  # 增加标签间距
        plt.ylabel('温度 (°C)', labelpad=10)  # 增加标签间距
        plt.legend()
        plt.tight_layout()  # 调整子图参数以适应图形区域


day_fft_ifft_modes = ["原始信号", "原始振幅谱", "傅里叶逆变换", "滤波后振幅谱", "信号比较"]


def day_fft_ifft_plotg(cycle_no, ktemp_wn5, ktemp_fft, ktemp_ifft, altitude_min, altitude_max, lamda_low, lamda_high, mode):

    N = len(ktemp_wn5[cycle_no, :])
    dt = (altitude_max - altitude_min) / (N - 1)  # 采样时间间隔
    n = np.arange(N)  # 时间序列索引
    f = n / (N * dt)
    t = np.round(np.linspace(altitude_min, altitude_max, N), 2)

    x = ktemp_wn5[cycle_no, :]  # 原始扰动温度
    y = ktemp_fft[cycle_no, :]  # 傅里叶变换频谱分析
    yy = ktemp_fft_lvbo[cycle_no, :]  # 滤波后的傅里叶变换频谱分析
    yyy = ktemp_ifft[cycle_no, :]  # 傅里叶逆变换后的扰动温度

    plt.figure(figsize=(15, 10))  # 调整图形大小

    # 计算纵坐标范围
    temp_limits = (min(min(x), min(yyy)), max(max(x), max(yyy)))
    amp_limits = (0, max(np.max(np.abs(y) * 2 / N),
                  np.max(np.abs(yy) * 2 / N)))

    if mode not in ["原始信号", "原始振幅谱", "傅里叶逆变换", "滤波后振幅谱", "信号比较"]:
        raise ValueError(
            "mode 参数应为 '原始信号', '原始振幅谱', '傅里叶逆变换', '滤波后振幅谱', '信号比较' 中的一个")

    # 原始信号的时间序列
    if mode == "原始信号":
        plt.plot(t, x)
        plt.title('（a）原始信号')
        plt.xlabel('高度 (km)', labelpad=10)
        plt.ylabel('温度 (°C)', labelpad=10)
        plt.ylim(temp_limits)  # 设置统一纵坐标范围

        # 原始振幅谱
    if mode == "原始振幅谱":
        plt.plot(f, np.abs(y) * 2 / N)
        plt.title('（b）原始振幅谱')
        plt.xlabel('频率/Hz', labelpad=10)
        plt.ylabel('振幅', labelpad=10)
        plt.ylim(amp_limits)  # 设置统一纵坐标范围

        # 通过IFFT回到时间域
    if mode == "傅里叶逆变换":
        plt.plot(t, yyy)
        plt.title('（c）傅里叶逆变换')
        plt.xlabel('高度 (km)', labelpad=10)
        plt.ylabel('温度 (°C)', labelpad=10)
        plt.ylim(temp_limits)  # 设置统一纵坐标范围

        # 滤波后的振幅谱
    if mode == "滤波后振幅谱":
        plt.plot(f, np.abs(yy) * 2 / N)
        plt.title(f'（d）滤除波长 < {lamda_low} km, > {lamda_high} km的波')
        plt.xlabel('频率/Hz', labelpad=10)
        plt.ylabel('振幅', labelpad=10)
        plt.ylim(amp_limits)  # 设置统一纵坐标范围

    if mode == "信号比较":
        # 绘制原始信号与滤波后信号
        plt.figure(figsize=(6, 8))  # 调整图形大小
        plt.plot(x, t, label='原始信号')
        plt.plot(yyy, t, label='滤波后信号', linestyle='--')
        plt.title('信号比较')
        plt.ylabel('高度 (km)', labelpad=10)  # 增加标签间距
        plt.xlabel('温度 (°C)', labelpad=10)  # 增加标签间距
        plt.legend()
        plt.xlim(temp_limits)  # 设置统一纵坐标范围
        plt.tight_layout()
# -----------------------------------------------------------------------------------------------------------------------------
# 6-3 示例的按高度的重力波势能变化曲线图


def day_cycle_power_wave_plot(cycle_no, altitude_min, altitude_max, ktemp_Nz, ktemp_Ptz):

    N = len(ktemp_Nz[cycle_no, :])
    y = np.round(np.linspace(altitude_min, altitude_max, N), 2)
    x1 = ktemp_Nz[cycle_no, :]
    x2 = ktemp_Ptz[cycle_no, :]

    plt.figure(figsize=(12, 10))  # 调整图形大小
    # 原始信号的时间序列
    plt.subplot(1, 2, 1)
    plt.plot(x1, y, label='原始信号')
    plt.title('（a）Nz')
    plt.xlabel('Nz', labelpad=10)  # 增加标签间距
    plt.ylabel('高度 (km)', labelpad=10)  # 增加标签间距

    # 原始信号的时间序列
    plt.subplot(1, 2, 2)
    plt.plot(x2, y, label='原始信号')
    plt.title('（b）Ptz')
    plt.xlabel('Ptz', labelpad=10)  # 增加标签间距
    plt.ylabel('高度 (km)', labelpad=10)  # 增加标签间距

    # 调整子图之间的边距
    plt.subplots_adjust(top=0.8, bottom=0.2, left=0.1,
                        right=0.8, hspace=0.3, wspace=0.2)
    plt.tight_layout()  # 调整子图参数以适应图形区域

    plt.show()
# -----------------------------------------------------------------------------------------------------------------------------
# 6-4 按日统计的按周期计算的不同高度的重力波势能 平面图


def day_power_wave_plot(altitude_min, altitude_max, ktemp_Nz, ktemp_Ptz):
    # 假设 ktemp_Nz 和 ktemp_Ptz 以及 altitude_min, altitude_max 已经定义好
    x = np.arange(ktemp_Nz.shape[0])
    y = np.round(np.linspace(altitude_min, altitude_max, ktemp_Nz.shape[1]), 2)

    # 创建一个图形，并指定两个子图
    fig, axs = plt.subplots(1, 2, figsize=(15, 10))

    # 第一幅图 (a) NZ
    cax1 = axs[0].imshow(ktemp_Nz.T, aspect='auto', cmap='viridis',
                         origin='lower', extent=[x[0], x[-1], y[0], y[-1]])
    fig.colorbar(cax1, ax=axs[0])  # 为第一幅图添加颜色条
    axs[0].set_title('(a) NZ')
    axs[0].set_ylabel('Height (km)')
    axs[0].set_xlabel('Cycles')
    axs[0].set_yticks(np.linspace(30, 90, 7))
    axs[0].set_yticklabels(np.round(np.linspace(30, 90, 7), 1))
    axs[0].set_xticks(np.arange(15))
    axs[0].set_xticklabels(np.arange(15))

    # 第二幅图 (b) PTZ
    cax2 = axs[1].imshow(ktemp_Ptz.T, aspect='auto', cmap='viridis',
                         origin='lower', extent=[x[0], x[-1], y[0], y[-1]])
    fig.colorbar(cax2, ax=axs[1])  # 为第二幅图添加颜色条
    axs[1].set_title('(b) PTZ')
    axs[1].set_ylabel('Height (km)')
    axs[1].set_xlabel('Cycles')
    axs[1].set_yticks(np.linspace(30, 90, 7))
    axs[1].set_yticklabels(np.round(np.linspace(30, 90, 7), 1))
    axs[1].set_xticks(np.arange(15))
    axs[1].set_xticklabels(np.arange(15))

    # 调整子图之间的边距
    plt.subplots_adjust(top=0.9, bottom=0.1, left=0.05,
                        right=0.95, hspace=0.3, wspace=0.3)
    plt.tight_layout()  # 调整布局以避免重叠
    plt.show()
# -----------------------------------------------------------------------------------------------------------------------------
# 6-5 按月统计的每日重力波势能随天变化的图


# -----------------------------------------------------------------------------------------------------------------------------
# 6-6 按年统计的每月重力波势能随月变化的图


# -----------------------------------------------------------------------------------------------------------------------------
# 7 主程序，运行数据,并输出主要统计分析图                                                                                          -
# -----------------------------------------------------------------------------------------------------------------------------
# 7-1 挑选某一天进行运行主要程序
def render_saber(file_path, mode, mode_args=None):
    # day_read=2018113,
    day_read = 2018094
    # 初始化某一天、某个纬度、高度范围等参数
    latitude_min = 30.0
    latitude_max = 50.0
    altitude_min = 30.0
    altitude_max = 90.0
    lamda_low = 2
    lamda_high = 15
    lvboin = True

    (ktemp_cycles, altitude_cycles,
     ktemp_wn0,
     ktemp_fit_wn1, ktemp_wn1,
     ktemp_fit_wn2, ktemp_wn2,
     ktemp_fit_wn3, ktemp_wn3,
     ktemp_fit_wn4, ktemp_wn4,
     ktemp_fit_wn5, ktemp_wn5,
     ktemp_fft, ktemp_fft_lvbo,  ktemp_ifft,
     ktemp_Nz, ktemp_Ptz) = day_process_main(
        file_path,
        # day_read=2018113,
        day_read,
        # 初始化某一天、某个纬度、高度范围等参数
        latitude_min,
        latitude_max,
        altitude_min,
        altitude_max,
        lamda_low, lamda_high, lvboin)
    # -----------------------------------------------------------------------------------------------------------------------------
    # 7-2 挑选某一个月进行运行主要程序
    (date_time, ktemp_cycles_mon, altitude_cycles_mon,                                 # 某月份 逐日时间 每日不同循环周期温度  不同循环周期高度数据
     # 距平扰动温度
     ktemp_wn0_mon,
     # 波数为1的拟合温度  滤波后的扰动温度
     ktemp_fit_wn1_mon, ktemp_wn1_mon,
     # 波数为2的拟合温度  滤波后的扰动温度
     ktemp_fit_wn2_mon, ktemp_wn2_mon,
     # 波数为3的拟合温度  滤波后的扰动温度
     ktemp_fit_wn3_mon, ktemp_wn3_mon,
     # 波数为4的拟合温度  滤波后的扰动温度
     ktemp_fit_wn4_mon, ktemp_wn4_mon,
     # 波数为5的拟合温度  滤波后的扰动温度
     ktemp_fit_wn5_mon, ktemp_wn5_mon,
     # 滤波0-5后扰动温度傅里叶频谱分析  滤除波长为2km-15km内后频谱分析 滤波后的扰动温度
     ktemp_fft_mon, ktemp_fft_lvbo_mon, ktemp_ifft_mon,
     # 滤波后NZ、PTZ重力波势能指标计算
     ktemp_Nz_mon, ktemp_Ptz_mon) = mon_process_main(
        file_path,
        # 初始化某一天、某个纬度、高度范围等参数
        latitude_min,
        latitude_max,
        altitude_min,
        altitude_max,
        lamda_low, lamda_high, lvboin)

    # -----------------------------------------------------------------------------------------------------------------------------
    # 7-3 绘制不同截过图

    if mode == "day_fit_wave_plot":

        height_no = 1
        # day_fit_wave_plotg(height_no, ktemp_wn0, ktemp_fit_wn1, ktemp_wn1, ktemp_fit_wn2, ktemp_wn2,
        #                    ktemp_fit_wn3, ktemp_wn3, ktemp_fit_wn4, ktemp_wn4, ktemp_fit_wn5, ktemp_wn5)
        day_fit_wave_plot(height_no, ktemp_wn0, ktemp_fit_wn1, ktemp_wn1, ktemp_fit_wn2, ktemp_wn2,
                          ktemp_fit_wn3, ktemp_wn3, ktemp_fit_wn4, ktemp_wn4, ktemp_fit_wn5, ktemp_wn5, mode_args)
    if mode == "day_fft_ifft_plot":
        cycle_no = 1
        day_fft_ifft_plotg(cycle_no, ktemp_wn5, ktemp_fft, ktemp_ifft,
                           altitude_min, altitude_max, lamda_low, lamda_high, mode_args)
    if mode == "day_cycle_power_wave_plot":

        day_cycle_power_wave_plot(cycle_no, altitude_min,
                                  altitude_max, ktemp_Nz, ktemp_Ptz)
    if mode == "day_power_wave_plot":
        day_power_wave_plot(altitude_min, altitude_max, ktemp_Nz, ktemp_Ptz)

    buffer = BytesIO()
    plt.savefig(buffer, format='png')
    buffer.seek(0)
    return buffer


ALL_RENDER_MODES = ["day_fit_wave_plot", "day_fft_ifft_plot",
                    "day_cycle_power_wave_plot", "day_power_wave_plot"]
