import numpy as np
import pandas as pd
from scipy.optimize import curve_fit
from scipy.signal import lombscargle, hilbert
from sklearn.decomposition import PCA
import sympy as sp


def calculate_n(Av, Au, Bu, Bv):
    # 计算 Fuv
    Fuv = Av * Au * np.cos(Bu - Bv)
    # 根据条件计算 n
    if Au > Av:
        n = 1
    elif Au < Av:
        if Fuv > 0:
            n = 0
        else:  # Fuv <= 0
            n = 2
    else:
        raise ValueError("Au 和 Av 不应相等")
    return n


def calculate_seta(Av, Au, Bu, Bv):
    # 计算 n
    n = calculate_n(Av, Au, Bu, Bv)
    # 计算 Fuv
    Fuv = Av * Au * np.cos(Bu - Bv)
    # 计算 arctan 部分
    arctan_value = np.atan2(2 * Fuv, Av**2 - Au**2)  # 更通用，返回的角度范围是 ([-π, π])
    # arctan_value = np.atan(2 * Fuv / (Av ** 2 - Au ** 2)) # 返回的角度范围是 ([-π/2, π/2])
    # 计算 seta
    seta = round((1 / 2) * (n * np.pi + arctan_value), 2)
    return seta  # 返回结果为弧度


# 判断椭圆旋转方向
def analyze_direction(df1, df2, seta_degrees):
    # 初始值
    # a = 0
    # b = 0

    # 计算元素
    element1 = df1.iloc[0, 1] * df1.iloc[1, 2] - \
        df1.iloc[0, 2] * df1.iloc[1, 1]
    element2 = df2.iloc[0, 1] * df2.iloc[1, 2] - \
        df2.iloc[0, 2] * df2.iloc[1, 1]

    # 判断 element1
    if element1 < 0:
        a = 1  # 表示顺时针，重力波向上传播
    else:
        a = -1

    # 判断 element2
    if element2 < 0:
        b = seta_degrees  # 表示顺时针，水平传播方向与假设方向同，seta
    else:
        b = seta_degrees + 180  # 逆时针，seta+180
    return a, b


# 计算u_v椭圆长短轴之比（先标准化椭圆，再计算最远距离最近距离之比）
def calculate_axial_ratio(u_fit, v_fit):
    # 将 ucmp 和 vcmp 组成一个数据矩阵
    data_matrix = np.vstack((u_fit, v_fit)).T

    # 使用 PCA 进行主成分分析
    pca = PCA(n_components=2)
    transformed_data = pca.fit_transform(data_matrix)

    # 提取主成分
    principal_components = pca.components_

    # 计算每个数据点在标准椭圆上的坐标
    scaling_factors = np.linalg.norm(principal_components, axis=1)
    standardized_data = transformed_data / scaling_factors

    # 创建 DataFrame
    standardized_df = pd.DataFrame(standardized_data, columns=["PC1", "PC2"])
    # 获取第一列和第二列的最大值和最小值
    col1_max = standardized_df.iloc[:, 0].max()
    col1_min = standardized_df.iloc[:, 0].min()
    col2_max = standardized_df.iloc[:, 1].max()
    col2_min = standardized_df.iloc[:, 1].min()

    # 计算长短轴比值
    axial_ratio = (col1_max - col1_min) / (col2_max - col2_min)

    return axial_ratio


def calculate_wave(data: pd.DataFrame, lat: float, g: float):
    height = data["alt"].values
    temp = data["temp"].values
    ucmp = data["u"].values
    vcmp = data["v"].values
    wcmp = data["wspeed"].values
    press = (data["press"].values) * 100  # 单位Pa

    coef_temp = np.polyfit(height, temp, 2)
    poly_temp = np.poly1d(coef_temp)
    coef_ucmp = np.polyfit(height, ucmp, 2)
    poly_ucmp = np.poly1d(coef_ucmp)
    coef_vcmp = np.polyfit(height, vcmp, 2)
    poly_vcmp = np.poly1d(coef_vcmp)
    coef_wcmp = np.polyfit(height, wcmp, 2)
    poly_wcmp = np.poly1d(coef_wcmp)

    residual_temp = temp - poly_temp(height)
    residual_ucmp = ucmp - poly_ucmp(height)
    residual_vcmp = vcmp - poly_vcmp(height)
    residual_wcmp = wcmp - poly_wcmp(height)

    ff = np.linspace(0.01, 25, 1000)  # 横坐标范围
    pgram_temp = lombscargle(height, residual_temp, ff)
    pgram_ucmp = lombscargle(height, residual_ucmp, ff)
    pgram_vcmp = lombscargle(height, residual_vcmp, ff)

    # todo:我的频率等于垂直波数（单位10-3rad/m）？
    main_frequency_temp = round(ff[np.argmax(pgram_temp)], 2)
    main_frequency_ucmp = round(ff[np.argmax(pgram_ucmp)], 2)
    main_frequency_vcmp = round(ff[np.argmax(pgram_vcmp)], 2)

    # 由波数计算波长
    λ1 = round(1 / ((main_frequency_temp / 2.5) * 0.4), 2)
    λ2 = round(1 / ((main_frequency_ucmp / 2.5) * 0.4), 2)
    λ3 = round(1 / ((main_frequency_vcmp / 2.5) * 0.4), 2)

    mean = round(sum([λ1, λ2, λ3]) / 3, 2)
    rsd_temp = abs(round(((λ1 - mean) / λ1) * 100, 2))
    rsd_ucmp = abs(round(((λ2 - mean) / λ2) * 100, 2))
    rsd_vcmp = abs(round(((λ3 - mean) / λ3) * 100, 2))

    if not (rsd_temp < 20 and rsd_ucmp < 20 and rsd_vcmp < 20):
        return []

    # 定义正弦波模型函数
    def sine_wave(height, A, B, omega):
        omega = 2 * np.pi / mean
        return A * np.sin(omega * height + B)

    def fit_sine_wave(height, data, mean):
        # 提供初始猜测值，例如 A=1, B=0, omega=2 * np.pi / mean
        initial_guess = [1, 0, 2 * np.pi / mean]
        # 设置参数的边界：A的下限为0（确保A为正值），B和omega没有边界限制
        bounds = ([0, -np.inf, -np.inf], [np.inf, np.inf, np.inf])
        # curve_fit 函数通过调整 sine_wave 函数的参数（A、B 和 omega）来最小化模型预测值与实际数据 data 之间的残差平方和。
        params, covariance = curve_fit(
            sine_wave, height, data, p0=initial_guess, bounds=bounds)
        return params

    # 对 u 风、v 风、温度扰动量进行拟合
    params_u = fit_sine_wave(height, residual_ucmp, mean)  # 输出3个参数分别是振幅，相位，角频率
    u_fit = sine_wave(height, *params_u)

    params_v = fit_sine_wave(height, residual_vcmp, mean)
    v_fit = sine_wave(height, *params_v)

    params_T = fit_sine_wave(height, residual_temp, mean)
    T_fit = sine_wave(height, *params_T)

    params_w = fit_sine_wave(height, residual_wcmp, mean)
    w_fit = sine_wave(height, *params_w)

    # 将高度、u_fit、v_fit创建一个数据框df1
    df1 = pd.DataFrame({"Height": height, "u_fit": u_fit, "v_fit": v_fit})

    # 计算假设的水平传播方向
    rad_seta = calculate_seta(
        params_v[0], params_u[0], params_u[1], params_v[1])  # 输入、输出都为弧度
    # 转换为角度并保留2位小数
    seta_degrees = round(np.degrees(rad_seta), 2)  # 38.41°

    # 计算水平风
    uh = u_fit * np.sin(seta_degrees * np.pi / 180) + \
        v_fit * np.cos(seta_degrees * np.pi / 180)  # 水平风
    # temp = T_fit
    # 将高度、水平风、温度创建一个数据框df2
    df2 = pd.DataFrame({"Height": height, "uh": uh, "temp": T_fit})
    a, b = analyze_direction(df1, df2, seta_degrees)  # 计算垂直、水平传播方向

    Axial_ratio = round(calculate_axial_ratio(u_fit, v_fit),
                        2)  # 计算长短轴之比并保留2位小数，后续计算本征频率会用到

    # 计算剩余几个参数
    # 常量
    # f = 1.32e-4  # rad/s
    # N = 2.13e-2  # 浮力频数，单位 rad/s

    # 根据站点纬度，计算惯性频率f
    ou = 7.27e-5  # 单位rad/s，地球自转频率
    # lat = 64.5
    f = 2 * ou * np.sin(lat * np.pi / 180)
    # print(f"f = {f:.3e} rad/s")

    # 求浮力频率N
    df3 = pd.DataFrame({"height": height, "temp": temp,
                       "temp_bar": poly_temp(height)})

    Cp = 1005  # 单位：J/kgK
    df3["temp_bar_derivative"] = np.gradient(
        df3["temp_bar"], (df3["height"]) * 1000)  # 求背景温度对高度的偏导
    gen = (g / temp) * (df3["temp_bar_derivative"] + (g / Cp))
    sqrt_gen = np.sqrt(np.abs(gen))  # np.sqrt()对列表求根号，np.sqrt()对数字求根号
    N = sqrt_gen.mean()  # 计算平均值，得到浮力频率的均值

    # 1. 本征频率omega_upper
    omega_upper = sp.symbols("omega_upper")
    eq1 = sp.Eq(Axial_ratio, omega_upper / f)
    omega_upper = (sp.solve(eq1, omega_upper))[0]
    omega_upper = sp.N(omega_upper)

    # 求本征频率与惯性频率的比值
    w_f = omega_upper / f  # todo:总是大于1

    # 本征周期
    zhou_qi = 2 * np.pi / omega_upper / 3600

    # 2. 垂直波长λ_z、水平波长λ_h

    # 垂直波数
    λ_z = mean  # 保留两位小数
    # 在北半球，重力波向上传播，垂直波数k_z＜0
    if a == 1:
        k_z = -(2 * np.pi / λ_z / 1000)
    else:
        k_z = 2 * np.pi / λ_z / 1000

    # 水平波数
    k_h = sp.symbols("k_h")
    eq2 = sp.Eq(k_h**2, (omega_upper**2 - f**2) * (k_z**2) / N**2)
    k_h = sp.solve(eq2, k_h)
    k_h = [sol.evalf() for sol in k_h if sol.evalf() > 0][0]
    k_h = sp.N(k_h)  # 正值

    # 在北半球，重力波向上传播，垂直波数k_h＜0
    if b > 180:  # 说明水平传播方向与假定方向相反
        k_h = -k_h
    else:
        k_h = k_h

    λ_h = abs(round((2 * np.pi / k_h / 1000), 2))  # 水平波长，取绝对值，因为波数可能为负数

    # 纬向、经向水平波长λ_x, λ_y
    # λ_x = abs(λ_h * np.sin(b * np.pi / 180))
    # λ_y = abs(λ_h * np.cos(b * np.pi / 180))

    # 本征相速度
    c_x = omega_upper / (abs(k_h) * np.sin(b * np.pi / 180))
    c_y = omega_upper / (abs(k_h) * np.cos(b * np.pi / 180))
    c_z = omega_upper / k_z

    """
    # 群速度
    C_gz = -((omega_upper ** 2 - f ** 2)) / (omega_upper * k_z)
    C_gh = sp.symbols('C_gh')
    eq3 = sp.Eq(C_gz / C_gh, sp.sqrt(omega_upper ** 2 - f ** 2) / N)
    C_gh = (sp.solve(eq3, C_gh))[0]
    C_gh = sp.N(C_gh) # 转成数值
    """

    # 波动能E_k
    u2_mean = np.mean(np.square(u_fit))
    v2_mean = np.mean(np.square(v_fit))
    Ek = (1 / 2) * (u2_mean + v2_mean)  # 3.11 J/kg

    # 势能E_p
    T_normalized = T_fit / poly_temp(height)  # 归一化处理
    T2_mean = np.mean(np.square(T_normalized))
    E_p = ((g**2) * T2_mean) / (2 * (N**2))  # 1.02J/kg

    # 动量通量
    # 经、纬向风动量通量
    P = press  # 单位Pa
    R = 287  # 单位：J/(kgK)，气体常量
    T = temp  # 温度已经是开尔文
    densi = P / (R * T)  # 向量
    c = (omega_upper * g) / (N**2)
    S = 1 - (f**2) / (omega_upper**2)  # 常数

    # 归一化温度90°相位变换

    hilb = hilbert(T_normalized)
    hilb1 = np.imag(hilb)  # 获取虚部作为90°相位变换的结果

    uT = u_fit * hilb1  # 向量
    uT1 = np.mean(uT)  # 均值
    MFu = (np.mean(-densi * c * uT1 * S)) * 1000  # 单位mPa

    vT = v_fit * hilb1
    vT1 = np.mean(vT)
    MFv = (np.mean(-densi * c * vT1 * S)) * 1000

    # 调整动量通量符号与本征相速度符号一致
    # 调整b的符号与a一致
    MFu1 = abs(MFu) if c_x >= 0 else -abs(MFu)
    MFv1 = abs(MFv) if c_y >= 0 else -abs(MFv)

    omega_upper = float(omega_upper)
    # make wf, c_x, c_y, c_z, MFu, MFv, zhouqi float
    w_f = float(w_f)
    c_x = float(c_x)
    c_y = float(c_y)
    c_z = float(c_z)
    MFu1 = float(MFu1)
    MFv1 = float(MFv1)
    zhou_qi = float(zhou_qi)

    return [
        a,
        b,
        omega_upper,
        w_f,
        λ_z,
        λ_h,
        c_x,
        c_y,
        c_z,
        Ek,
        E_p,
        MFu1,
        MFv1,
        params_u[0],
        params_v[0],
        params_T[0],
        ucmp,
        vcmp,
        temp,
        height,
        poly_ucmp,
        poly_vcmp,
        poly_wcmp,
        poly_temp,
        residual_ucmp,
        residual_vcmp,
        residual_temp,
        u_fit,
        v_fit,
        T_fit,
        uh,
        zhou_qi,
    ]


def extract_wave(data: pd.DataFrame, lat: float, g: float):
    wave = calculate_wave(
        data[(data["alt"] >= 15) & (data["alt"] <= 25)], lat, g)
    if len(wave) == 0:
        return []
    return [
        wave[0],
        wave[1],
        round(wave[2], 6),
        round(wave[3], 2),
        round(wave[4], 2),
        round(wave[5], 2),
        round(wave[6], 2),
        round(wave[7], 2),
        round(wave[8], 3),
        round(wave[9], 4),
        round(wave[10], 4),
        round(wave[11], 4),
        round(wave[12], 4),
        round(wave[13], 2),
        round(wave[14], 2),
        round(wave[15], 2),
        round(wave[31], 2),
    ]


def is_terrain_wave(data: pd.DataFrame, lat: float, g: float):
    wave = calculate_wave(
        data[(data["alt"] >= 2) & (data["alt"] <= 10)], lat, g)
    if len(wave) == 0:
        return False

    # 计算观测u风v风的均值，为后边合成背景速度
    mean_ucmp = np.mean(wave[16])  # 11.34
    mean_vcmp = np.mean(wave[17])  # 0.38

    x1, y1 = mean_ucmp, mean_vcmp  # 背景风的分量
    x2, y2 = float(wave[6]), float(wave[7])  # 相速度的分量
    # wave[18](wave[19])

    """
    # 二阶拟合的u风v风
    mean_ucmp = np.mean(poly_ucmp(height))
    mean_vcmp = np.mean(poly_vcmp(height))

    x1, y1 = mean_ucmp, mean_vcmp  # 背景风的分量
    x2, y2 = float(c_x), float(c_y)  # 相速度的分量
    """
    # 计算两个风的方向
    direction1 = np.arctan2(y1, x1)
    direction2 = np.arctan2(y2, x2)
    # 转化为角度（该角度是以x轴逆时针旋转），取特殊值用于验证。
    # direction1_degrees = np.degrees(direction1)
    # direction2_degrees = np.degrees(direction2)

    # 计算夹角（弧度）
    angle_radians = np.abs(direction1 - direction2)
    # 确保夹角在0到π之间
    if angle_radians > np.pi:
        angle_radians = 2 * np.pi - angle_radians
    # 将夹角转换为度
    angle_degrees = np.degrees(angle_radians)

    # 判断是否超过150，或者165
    if round(angle_degrees, 2) > 150:
        return True
    return False