# 重力波 一天cosmic重力波

import os
import time
import numpy as np
import pandas as pd
from scipy.interpolate import interp1d
from scipy.optimize import curve_fit
import netCDF4 as nc
import matplotlib.pyplot as plt

from CONSTANT import DATA_BASEPATH

# 模块1: 数据读取与预处理


def read_and_preprocess_data(base_folder_path, day_num):
    """
    读取指定路径下指定天数的相关数据文件，进行插值等预处理操作

    参数:
    base_folder_path (str): 基础文件夹路径
    day_num (int): 表示天数的序号（范围按代码中原本逻辑处理）

    返回:
    final_df (pd.DataFrame): 预处理后的最终数据框
    """
    dfs = []
    # 构建当前文件夹的路径
    folder_name = f"atmPrf_repro2021_2008_00{day_num}" if day_num < 10 else f"atmPrf_repro2021_2008_0{day_num}"
    folder_path = os.path.join(base_folder_path, folder_name)
    # 检查文件夹是否存在
    if os.path.exists(folder_path):
        # 遍历文件夹中的文件
        for file_name in os.listdir(folder_path):
            if file_name.endswith('.0390_nc'):
                finfo = os.path.join(folder_path, file_name)
                print(f"正在处理文件: {finfo}")
                try:
                    dataset = nc.Dataset(finfo, 'r')
                    # 提取变量数据
                    temp = dataset.variables['Temp'][:]
                    altitude = dataset.variables['MSL_alt'][:]
                    lat = dataset.variables['Lat'][:]
                    lon = dataset.variables['Lon'][:]
                    # 创建DataFrame
                    df = pd.DataFrame({
                        'Longitude': lon,
                        'Latitude': lat,
                        'Altitude': altitude,
                        'Temperature': temp
                    })
                    dataset.close()
                    # 剔除高度大于60的行
                    df = df[df['Altitude'] <= 60]

                    # 对每个文件的数据进行插值
                    alt_interp = np.linspace(
                        df['Altitude'].min(), df['Altitude'].max(), 3000)
                    f_alt = interp1d(
                        df['Altitude'], df['Altitude'], kind='linear', fill_value="extrapolate")
                    f_lon = interp1d(
                        df['Altitude'], df['Longitude'], kind='linear', fill_value="extrapolate")
                    f_lat = interp1d(
                        df['Altitude'], df['Latitude'], kind='linear', fill_value="extrapolate")
                    f_temp = interp1d(
                        df['Altitude'], df['Temperature'], kind='linear', fill_value="extrapolate")

                    # 计算插值结果
                    interpolated_alt = f_alt(alt_interp)
                    interpolated_lon = f_lon(alt_interp)
                    interpolated_lat = f_lat(alt_interp)
                    interpolated_temp = f_temp(alt_interp)

                    # 创建插值后的DataFrame
                    interpolated_df = pd.DataFrame({
                        'Altitude': interpolated_alt,
                        'Longitude': interpolated_lon,
                        'Latitude': interpolated_lat,
                        'Temperature': interpolated_temp
                    })

                    # 将插值后的DataFrame添加到列表中
                    dfs.append(interpolated_df)
                except Exception as e:
                    print(f"处理文件 {finfo} 时出错: {e}")
    else:
        print(f"文件夹 {folder_path} 不存在。")
    # 按行拼接所有插值后的DataFrame
    final_df = pd.concat(dfs, axis=0, ignore_index=True)
    # 获取 DataFrame 的长度
    num_rows = len(final_df)
    # 生成一个每3000个数从0到2999的序列并重复
    altitude_values = np.tile(np.arange(3000), num_rows // 3000 + 1)[:num_rows]
    # 将生成的值赋给 DataFrame 的 'Altitude' 列
    final_df['Altitude'] = altitude_values
    # 摄氏度换算开尔文
    final_df['Temperature'] = final_df['Temperature'] + 273.15
    return final_df


# 模块2: 数据筛选与网格划分
def filter_and_grid_data(final_df):
    """
    对输入的数据框进行纬度筛选以及经度网格划分等操作

    参数:
    final_df (pd.DataFrame): 预处理后的完整数据框

    返回:
    latitude_filtered_df (pd.DataFrame): 经过纬度筛选和经度网格划分后的数据框
    altitude_temperature_mean (pd.DataFrame): 按高度分组求平均温度后的数据框
    """
    # 筛选出纬度在30到40度之间的数据
    latitude_filtered_df = final_df[(
        final_df['Latitude'] >= 30) & (final_df['Latitude'] <= 40)]
    # 划分经度网格，20°的网格
    lon_min, lon_max = latitude_filtered_df['Longitude'].min(
    ), latitude_filtered_df['Longitude'].max()
    lon_bins = np.arange(lon_min, lon_max + 20, 20)  # 创建经度网格边界
    # 将数据分配到网格中
    latitude_filtered_df['Longitude_Grid'] = np.digitize(
        latitude_filtered_df['Longitude'], lon_bins) - 1
    # 对相同高度的温度取均值，忽略NaN
    altitude_temperature_mean = latitude_filtered_df.groupby(
        'Altitude')['Temperature'].mean().reset_index()
    # 重命名列，使其更具可读性
    altitude_temperature_mean.columns = ['Altitude', 'Mean_Temperature']
    return latitude_filtered_df, altitude_temperature_mean


# 模块3: 计算wn0和初始扰动相关操作
def calculate_wn0_and_perturbation(latitude_filtered_df, altitude_temperature_mean):
    """
    计算wn0以及相关的温度差值等操作

    参数:
    latitude_filtered_df (pd.DataFrame): 经过纬度筛选和经度网格划分后的数据框
    altitude_temperature_mean (pd.DataFrame): 按高度分组求平均温度后的数据框

    返回:
    merged_df (pd.DataFrame): 合并相关数据后的结果数据框，包含温度差值等列
    """
    altitude_range = range(0, 3000)
    all_heights_mean_temperature = []  # 用于存储所有高度下的温度均值结果
    for altitude in altitude_range:
        # 筛选出当前高度的所有数据
        altitude_df = latitude_filtered_df[latitude_filtered_df['Altitude'] == altitude]
        # 对Longitude_Grid同一区间的温度取均值
        temperature_mean_by_grid = altitude_df.groupby(
            'Longitude_Grid')['Temperature'].mean().reset_index()
        # 重命名列，使其更具可读性
        temperature_mean_by_grid.columns = [
            'Longitude_Grid', 'Mean_Temperature']
        # 添加高度信息列，方便后续区分不同高度的结果
        temperature_mean_by_grid['Altitude'] = altitude
        # 将当前高度的结果添加到列表中
        all_heights_mean_temperature.append(temperature_mean_by_grid)
    # 将所有高度的结果合并为一个DataFrame
    combined_mean_temperature_df = pd.concat(
        all_heights_mean_temperature, ignore_index=True)
    # 基于Altitude列合并两个DataFrame，只保留能匹配上的行
    merged_df = pd.merge(combined_mean_temperature_df,
                         altitude_temperature_mean, on='Altitude', how='inner')
    # 计算差值（减去wn0的扰动）
    merged_df['Temperature_Difference'] = merged_df['Mean_Temperature_x'] - \
        merged_df['Mean_Temperature_y']
    return merged_df


# 模块4: 高度相同下不同区间网格数据的波拟合和滤波处理
def wave_fitting_and_filtering(merged_df):
    """
    对合并后的数据框按高度进行波拟合以及滤波处理

    参数:
    merged_df (pd.DataFrame): 合并相关数据后的结果数据框，包含温度差值等列

    返回:
    fit_results (dict): 存储每个高度的最佳拟合参数的字典
    residuals (dict): 存储每个高度的残差值的字典
    background_temperatures (dict): 存储每个高度的背景温度的字典
    """
    def single_harmonic(x, A, phi):
        return A * np.sin(2 * np.pi / (18 / k) * x + phi)

    fit_results = {}
    fitted_curves = {}
    residuals = {}
    background_temperatures = {}
    grouped = merged_df.groupby('Altitude')
    for altitude, group in grouped:
        y_data = group['Temperature_Difference'].values
        x_data = np.arange(len(y_data))
        wn0_data = group['Mean_Temperature_y'].values  # 获取同一高度下的wn0数据
        # 检查Temperature_Difference列是否全部为NaN
        if np.all(np.isnan(y_data)):
            fit_results[altitude] = {'A': [np.nan] * 5, 'phi': [np.nan] * 5}
            fitted_curves[altitude] = [np.nan * x_data] * 5
            residuals[altitude] = np.nan * x_data
            background_temperatures[altitude] = np.nan * \
                x_data  # 对应高度全部为NaN时，背景温度也设为NaN
        else:
            # 替换NaN值为非NaN值的均值
            y_data = np.where(np.isnan(y_data), np.nanmean(y_data), y_data)
            # 初始化存储WN参数和曲线的列表
            wn_params = []
            wn_curves = []
            # 计算wn0（使用Mean_Temperature_y列数据）
            wn0 = wn0_data
            # 对WN1至WN5进行拟合
            for k in range(1, 6):
                # 更新单谐波函数中的k值
                def harmonic_func(x, A, phi): return single_harmonic(x, A, phi)
                # 使用curve_fit进行拟合
                popt, pcov = curve_fit(harmonic_func, x_data, y_data, p0=[
                                       np.nanmax(y_data) - np.nanmin(y_data), 0])
                A_fit, phi_fit = popt
                # 存储拟合结果
                wn_params.append({'A': A_fit, 'phi': phi_fit})
                # 使用拟合参数生成拟合曲线
                WN = harmonic_func(x_data, A_fit, phi_fit)
                wn_curves.append(WN)
                # 计算残差值
                y_data = y_data - WN  # 使用残差值作为下一次拟合的y_data
            # 存储结果
            fit_results[altitude] = wn_params
            fitted_curves[altitude] = wn_curves
            residuals[altitude] = y_data  # 最终残差值
            # 计算同一高度下的背景温度（wn0 + wn1 + wn2 + wn3 + wn4 + wn5）
            wn_sum = np.sum([wn0] + wn_curves, axis=0)
            background_temperatures[altitude] = wn_sum
    return fit_results, residuals, background_temperatures


# 模块5: 带通滤波处理
def bandpass_filtering(residuals):
    """
    对残差数据进行带通滤波处理

    参数:
    residuals (dict): 存储每个高度的残差值的字典

    返回:
    result (dict): 存储滤波后每个高度结果的字典
    """
    df = pd.DataFrame(residuals)
    # 使用前向填充（用上一个有效值填充 NaN）
    df.ffill(axis=1, inplace=True)
    result = {}
    lambda_low = 2  # 2 km
    lambda_high = 15  # 15 km
    f_low = 2 * np.pi / lambda_high
    f_high = 2 * np.pi / lambda_low
    for idx, residuals_array in df.iterrows():
        # 提取有效值
        valid_values = np.ma.masked_array(
            residuals_array, np.isnan(residuals_array))
        compressed_values = valid_values.compressed()  # 去除NaN值后的数组
        N = len(compressed_values)  # 有效值的数量
        # 如果有效值为空（即所有值都是NaN），则将结果设置为NaN
        if N == 0:
            result[idx] = np.full_like(residuals_array, np.nan)
        else:
            # 时间序列和频率
            dt = 0.02  # 假设的时间间隔
            n = np.arange(N)
            f = n / (N * dt)
            # 傅里叶变换
            y = np.fft.fft(compressed_values)  # 使用去除NaN后的数组进行FFT
            # 频率滤波
            yy = y.copy()
            freq_filter = (f >= f_low) & (f <= f_high)  # 保留指定频段
            yy[~freq_filter] = 0  # 过滤掉指定频段外的值
            # 逆傅里叶变换
            perturbation_after = np.real(np.fft.ifft(yy))
            # 将处理结果插回到result字典中
            result[idx] = perturbation_after
    return result


# 模块6: 计算势能相关操作
def calculate_potential_energy(background_temperatures, result, heights):
    """
    基于背景温度、滤波后的扰动以及高度数据计算浮力频率和势能，并求平均

    参数:
    background_temperatures (dict): 存储每个高度的背景温度的字典
    result (dict): 存储滤波后每个高度结果的字典
    heights (np.ndarray): 高度数据数组

    返回:
    mean_ktemp_Nz (np.ndarray): 平均浮力频率数据
    mean_ktemp_Ptz (np.ndarray): 平均势能数据
    """
    # 处理背景温度和扰动温度数据格式
    heights1 = list(background_temperatures.keys())
    data_length1 = len(next(iter(background_temperatures.values())))
    background_matrix = np.zeros((data_length1, len(heights1)))
    for idx, height in enumerate(heights1):
        background_matrix[:, idx] = background_temperatures[height]
    heights2 = list(result.keys())
    data_length2 = len(next(iter(result.values())))
    perturbation_matrix = np.zeros((data_length2, len(heights2)))
    for idx, height in enumerate(heights2):
        perturbation_matrix[:, idx] = result[height]
    perturbation_matrix = perturbation_matrix.T

    def brunt_vaisala_frequency(g, BT_z, c_p, heights):
        # 计算位温随高度的变化率
        dBT_z_dz = np.gradient(BT_z, heights)
        # 计算 Brunt-Väisälä 频率，根号内取绝对值
        frequency_squared = (g / BT_z) * ((g / c_p) + dBT_z_dz)
        frequency = np.sqrt(np.abs(frequency_squared))
        return frequency

    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  # 比热容
    N_z_matrix = []
    PT_z_matrix = []
    for i in range(background_matrix.shape[0]):
        BT_z = np.array(background_matrix[i])
        PT_z = np.array(perturbation_matrix[i])  # 滤波后的扰动

        # 调用Brunt-Väisälä频率函数
        N_z = brunt_vaisala_frequency(g, BT_z, c_p, heights)

        # 调用势能函数
        PW = calculate_gravitational_potential_energy(g, BT_z, N_z, PT_z)

        # 将结果添加到矩阵中
        N_z_matrix.append(N_z)
        PT_z_matrix.append(PW)
    ktemp_Nz = np.vstack(N_z_matrix)
    ktemp_Ptz = np.vstack(PT_z_matrix)
    mean_ktemp_Nz = np.mean(ktemp_Nz, axis=0)
    mean_ktemp_Ptz = np.mean(ktemp_Ptz, axis=0)
    # 对mean_ktemp_Ptz中的值做处理，如果超过40，除以10
    mean_ktemp_Ptz[mean_ktemp_Ptz > 40] /= 10
    return mean_ktemp_Nz, mean_ktemp_Ptz
# 模块7 绘制重力波势能和浮力频率日图


def plot_results_mean_ktemp_Nz(mean_ktemp_Nz, heights):
    """
    绘制平均浮力频率和平均势能随高度变化的图像

    参数:
    mean_ktemp_Nz (np.ndarray): 平均浮力频率数据
    mean_ktemp_Ptz (np.ndarray): 平均势能数据
    heights (np.ndarray): 高度数据数组
    """
    # 绘制平均浮力频率（mean_ktemp_Nz）随高度变化的图像
    plt.clf()
    plt.figure(figsize=(10, 6))
    plt.plot(mean_ktemp_Nz, heights / 1000)  # 高度单位换算为km，方便展示
    plt.xlabel('Average (N_z)')
    plt.ylabel('H(km)')
    # plt.gca().invert_yaxis()  # 使高度坐标轴从上到下递增，符合常规习惯
    # plt.show()


def plot_results_mean_ktemp_Ptz(mean_ktemp_Ptz, heights):
    """

    # 绘制平均势能（mean_ktemp_Ptz）随高度变化的图像
    """
    # clean the plot
    plt.clf()
    plt.figure(figsize=(10, 6))
    plt.plot(mean_ktemp_Ptz, heights / 1000)  # 高度单位换算为km，方便展示
    plt.xlabel('Average (PT_z)')
    plt.ylabel('H (km)')
    # plt.gca().invert_yaxis()  # 使高度坐标轴从上到下递增，符合常规习惯
    # plt.show()


class CosmicGravitywPlot:

    def __init__(self, year, day):
        BASE_PATH_COSMIC = DATA_BASEPATH.cosmic

        base_folder_path = f"./{BASE_PATH_COSMIC}/{year}"
        day_num = day
        # 模块1调用
        final_df = read_and_preprocess_data(base_folder_path, day_num)
        # 模块2调用
        latitude_filtered_df, altitude_temperature_mean = filter_and_grid_data(
            final_df)
        # 模块3调用
        merged_df = calculate_wn0_and_perturbation(
            latitude_filtered_df, altitude_temperature_mean)
        # 模块4调用
        fit_results, residuals, background_temperatures = wave_fitting_and_filtering(
            merged_df)
        # 模块5调用
        result = bandpass_filtering(residuals)
        # 创建从1到80km的高度列表（按照代码中原本逻辑转换为对应单位的数组）
        heights = np.linspace(0, 60000, 3000)
        # 模块6调用，计算势能相关数据并获取平均浮力频率和平均势能
        mean_ktemp_Nz, mean_ktemp_Ptz = calculate_potential_energy(
            background_temperatures, result, heights)
        # print("平均浮力频率（mean_ktemp_Nz）:", mean_ktemp_Nz)
        # print("平均势能（mean_ktemp_Ptz）:", mean_ktemp_Ptz)
        self.mean_ktemp_Nz = mean_ktemp_Nz
        self.mean_ktemp_Ptz = mean_ktemp_Ptz
        self.heights = heights
    # 调用绘图模块函数进行绘图

    def plot_results_mean_ktemp_Nz(self):
        plot_results_mean_ktemp_Nz(self.mean_ktemp_Nz, self.heights)

    def plot_results_mean_ktemp_Ptz(self):
        plot_results_mean_ktemp_Ptz(self.mean_ktemp_Ptz, self.heights)
