# 探空气球只有重力波

from dataclasses import dataclass
from io import BytesIO
import matplotlib.pyplot as plt
import numpy as np

from modules.balloon.extract_wave import calculate_wave, is_terrain_wave

lat = 52.21
g = 9.76


def plot_polynomial_fit(ucmp, vcmp, temp, height, poly_ucmp, poly_vcmp, poly_temp):
    plt.figure(figsize=(16, 5))

    # u 风扰动量
    plt.subplot(1, 3, 1)
    plt.plot(ucmp, height, linestyle="-")
    plt.plot(poly_ucmp(height), height)
    plt.ylabel("高度(km)")
    plt.xlabel("纬向风(m/s)")
    plt.grid(True)

    # v 风扰动量
    plt.subplot(1, 3, 2)
    plt.plot(vcmp, height, linestyle="-")
    plt.plot(poly_vcmp(height), height)
    plt.ylabel("高度(km)")
    plt.xlabel("经向风 (m/s)")
    plt.grid(True)

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

    # 温度扰动量
    plt.subplot(1, 3, 3)
    plt.plot(temp, height, linestyle="-")
    plt.plot(poly_temp(height), height)
    plt.ylabel("高度(km)")
    plt.xlabel("温度(K)")
    plt.grid(True)

    plt.tight_layout()


def plot_sine_fit(residual_ucmp, residual_vcmp, residual_temp, height, u_fit, v_fit, T_fit):
    plt.figure(figsize=(16, 5))

    # u 风扰动量
    plt.subplot(1, 3, 1)
    plt.plot(residual_ucmp, height, marker="o", linestyle="None")
    plt.plot(u_fit, height)
    plt.ylabel("高度(km)")
    plt.xlabel("纬向风 (m/s)")
    plt.grid(True)

    # v 风扰动量
    plt.subplot(1, 3, 2)
    plt.plot(residual_vcmp, height, marker="o", linestyle="None")
    plt.plot(v_fit, height)
    plt.ylabel("高度(km)")
    plt.xlabel("经向风(m/s)")
    plt.grid(True)

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

    # 温度扰动量
    plt.subplot(1, 3, 3)
    plt.plot(residual_temp, height, marker="o", linestyle="None")
    plt.plot(T_fit, height)
    plt.ylabel("高度(km)")
    plt.xlabel("温度(K)")
    plt.grid(True)

    plt.tight_layout()


def plot_uv_vector(u_fit, v_fit, height, specified_heights, markers):
    plt.figure(figsize=(8, 6))
    plt.plot(u_fit, v_fit)
    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.axvline(0, color="gray", linestyle="--")
    plt.axhline(0, color="gray", linestyle="--")
    plt.xlabel("经向风 (m/s)")
    plt.ylabel("纬向风(m/s)")
    plt.legend()
    plt.grid(True)
    plt.title("径向风-纬向风矢量图")


def plot_temp_horizontal_wind(uh, T_fit, height, specified_heights, markers):
    plt.figure(figsize=(8, 6))
    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.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.title("温度-水平风矢量图")


def render_by_mode_single(data, mode):
    wave = calculate_wave(
        data[(data["alt"] >= 15) & (data["alt"] <= 25)], lat, g)

    @dataclass
    class ReturnValue:
        image: BytesIO
        is_terrain_wave: bool
        has_data: bool

    if len(wave) == 0:
        return ReturnValue(image=BytesIO(), is_terrain_wave=False, has_data=False)

    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:]

    if mode == "观测的二阶多项式拟合":
        plot_polynomial_fit(ucmp, vcmp, temp, height, poly_ucmp,
                            poly_vcmp, poly_temp)
    elif mode == "扰动分量的正弦波拟合":
        plot_sine_fit(residual_ucmp, residual_vcmp, residual_temp,
                      height, u_fit, v_fit, T_fit)
    elif mode == "径向风-纬向风矢量图":
        plot_uv_vector(u_fit, v_fit, height, [15, 20, 25], [
                       "o", "s", "D"])
    elif mode == "温度-水平风矢量图":
        plot_temp_horizontal_wind(uh, T_fit, height, [15, 20, 25], [
                                  "o", "s", "D"])
    else:
        raise ValueError("Unknown mode")
    plt.rcParams["font.sans-serif"] = ["SimHei"]  # 显示中文
    buff = BytesIO()
    plt.savefig(buff, format="png")
    plt.close()

    buff.seek(0)

    return ReturnValue(
        image=buff,
        is_terrain_wave=c,
        has_data=True
    )