import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import numpy as np
from modules.balloon.extract_wave import calculate_wave, is_terrain_wave


def plot_once(data: pd.DataFrame, path: str, lat: float, g: float):
    wave = calculate_wave(
        data[(data["alt"] >= 15) & (data["alt"] <= 25)], lat, g)
    if len(wave) == 0:
        return []

    c = is_terrain_wave(data, lat, g)
    a, b, omega_upper, w_f, λ_z, λ_h, c_x, c_y, c_z, Ek, E_p, MFu1, MFv1, params_u, params_v, params_T = wave[
        :16]
    ucmp, vcmp, temp, height, poly_ucmp, poly_vcmp, _, poly_temp, residual_ucmp, residual_vcmp, residual_temp, u_fit, v_fit, T_fit, uh, _ = wave[
        16:]

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

    # 二阶多项式拟合曲线

    # u 风扰动量
    plt.subplot(2, 6, 1)
    plt.plot(ucmp, height, label="", linestyle="-")
    plt.plot(poly_ucmp(height), height, label="")
    plt.ylabel("Height(km)")
    plt.xlabel("Zonal wind (m/s)")
    plt.grid(True)

    # v 风扰动量
    plt.subplot(2, 6, 2)
    plt.plot(vcmp, height, label="", linestyle="-")
    plt.plot(poly_vcmp(height), height, label="")
    plt.ylabel("Height(km)")
    plt.xlabel("Meridional wind (m/s)")
    plt.grid(True)

    plt.title("观测的二阶多项式拟合", fontsize=16)

    # 温度扰动量
    plt.subplot(2, 6, 3)
    plt.plot(temp, height, label="", linestyle="-")
    plt.plot(poly_temp(height), height, label="")
    plt.ylabel("Height(km)")
    plt.xlabel("Temperature(K)")
    plt.grid(True)

    # 绘制正弦拟合图

    # u 风扰动量
    plt.subplot(2, 6, 4)
    plt.plot(residual_ucmp, height, label="", marker="o", linestyle="None")
    plt.plot(u_fit, height, label="")
    plt.ylabel("Height(km)")
    plt.xlabel("Zonal wind (m/s)")
    plt.grid(True)

    # v 风扰动量
    plt.subplot(2, 6, 5)
    plt.plot(residual_vcmp, height, label="", marker="o", linestyle="None")
    plt.plot(v_fit, height, label="")
    plt.ylabel("Height(km)")
    plt.xlabel("Meridional wind (m/s)")
    plt.grid(True)

    plt.title("扰动分量的正弦波拟合", fontsize=16)

    # 温度扰动量
    plt.subplot(2, 6, 6)
    plt.plot(residual_temp, height, label="", marker="o", linestyle="None")
    plt.plot(T_fit, height, label="")
    plt.ylabel("Height(km)")
    plt.xlabel("Temperature(K)")
    plt.grid(True)

    # 标记3个特定高度
    specified_heights = [15, 15.05, 15.1]
    markers = ["*", "D", "^"]  # 星号，菱形，三角形

    # a. U-V矢量曲线
    plt.subplot(2, 2, 3)
    plt.plot(u_fit, v_fit)  # 绘制U-V曲线
    for h, marker in zip(specified_heights, markers):
        index = np.abs(height - h).argmin()  # 找到离指定高度最近的索引
        plt.scatter(u_fit[index], v_fit[index],
                    marker=marker, s=100, label=f"{h} km")

    # 设置坐标轴范围和比例
    plt.xlim(-8, 8)
    plt.ylim(-4, 4)
    # plt.gca().set_aspect("equal")  # 确保横纵坐标刻度长短一致
    plt.axvline(0, color="gray", linestyle="--")
    plt.axhline(0, color="gray", linestyle="--")
    plt.xlabel("Zonal Wind (m/s)")
    plt.ylabel("Meridional Wind (m/s)")
    plt.legend()  # 显示图例
    plt.grid(True)  # 显示网格线
    plt.text(0.05, 1, "(a)", transform=plt.gca().transAxes,
             fontsize=14, verticalalignment="top")
    plt.title("径向风-纬向风矢量图")

    # b. 水平风-温度的矢端曲线
    plt.subplot(2, 2, 4)
    plt.plot(uh, T_fit)  # 绘制水平风-温度曲线
    for h, marker in zip(specified_heights, markers):
        index = np.abs(height - h).argmin()  # 找到离指定高度最近的索引
        plt.scatter(uh[index], T_fit[index],
                    marker=marker, s=100, label=f"{h} km")

    # 设置坐标轴范围和比例
    plt.xlim(-4, 4)
    plt.ylim(-2, 2)
    # plt.gca().set_aspect("equal")  # 确保横纵坐标刻度长短一致
    plt.axvline(0, color="gray", linestyle="--")
    plt.axhline(0, color="gray", linestyle="--")
    plt.xlabel("Horizontal wind (m/s)")
    plt.ylabel("Temp (K)")
    plt.legend()
    plt.grid(True)
    plt.text(0.05, 1, "(b)", transform=plt.gca().transAxes,
             fontsize=14, verticalalignment="top")
    plt.title("温度-水平风矢量图")

    # 设置中文显示和负号正常显示
    plt.rcParams["font.sans-serif"] = ["SimHei"]  # 显示中文
    plt.rcParams["axes.unicode_minus"] = False  # 正常显示负号

    plt.tight_layout()
    # plt.savefig(path, transparent=True)
    plt.savefig(path)
    plt.close()

    """
    地形重力波 c 是否是地形重力波
    垂直传播方向 a 1上
    水平传播方向 b degree
    本征（固有）频率 round(omega_upper, 6) rad/s
    本征周期 round(2 * np.pi / omega_upper / 3600, 2) h
    本征频率与固有频率的比值（即长短轴之比） round(w_f, 2)
    垂直波长 round(λ_z, 2) km
    水平波长 round(λ_h, 2) km
    纬向本征相速度 round(c_x, 2) m/s
    经向本征相速度 round(c_y, 2) m/s
    垂直本征相速度 round(c_z, 3) m/s
    波动能 round(Ek, 4) J/kg
    势能 round(E_p, 4) J/kg
    纬向动量通量 round(MFu1, 4) mPa
    经向动量通量 round(MFv1, 4) mPa
    纬向风扰动振幅 round(params_u, 2) m/s
    经向风扰动振幅 round(params_v, 2) m/s
    温度扰动振幅 round(params_T, 2) K
    """
    return [
        c,
        a,
        b,
        round(omega_upper, 6),
        round(2 * np.pi / omega_upper / 3600, 2),
        round(w_f, 2),
        round(λ_z, 2),
        round(λ_h, 2),
        round(c_x, 2),
        round(c_y, 2),
        round(c_z, 3),
        round(Ek, 4),
        round(E_p, 4),
        round(MFu1, 4),
        round(MFv1, 4),
        round(params_u, 2),
        round(params_v, 2),
        round(params_T, 2),
    ]