# 探空气球只有重力波

from io import BytesIO
import numpy as np
import pandas as pd
import seaborn as sns
from windrose import WindroseAxes

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec


def plot_w_f(filtered_df, ax):
    min_val = 1
    max_val = 10
    bin_width = 0.5
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["w_f"], bins=bins, kde=False,
                 edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("w/f值统计结果", fontsize=24)
    ax.set_xlabel("w/f")
    ax.set_ylabel("Occurrence(%)")


def plot_zhou_qi(filtered_df, ax):
    min_val = np.floor(filtered_df["zhou_qi"].min())
    max_val = np.ceil(filtered_df["zhou_qi"].max())
    bin_width = 1
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["zhou_qi"], bins=bins,
                 kde=False, edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("周期统计结果", fontsize=24)
    ax.set_xlabel("h")
    ax.set_ylabel("Occurrence(%)")


def plot_ver_wave_len(filtered_df, ax):
    min_val = np.floor(filtered_df["ver_wave_len"].min())
    max_val = np.ceil(filtered_df["ver_wave_len"].max())
    bin_width = 0.5
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["ver_wave_len"], bins=bins,
                 kde=False, edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("垂直波长分布", fontsize=24)
    ax.set_xlabel("Vertical wavelength(km)")
    ax.set_ylabel("Occurrence(%)")


def plot_hori_wave_len(filtered_df, ax):
    min_val = np.floor(filtered_df["hori_wave_len"].min())
    max_val = np.ceil(filtered_df["hori_wave_len"].max())
    bin_width = 100
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["hori_wave_len"], bins=bins,
                 kde=False, edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("水平波长分布", fontsize=24)
    ax.set_xlabel("Horizontal wavelength(km)")
    ax.set_ylabel("Occurrence(%)")


def plot_c_x(filtered_df, ax):
    min_val = np.floor(filtered_df["c_x"].min())
    max_val = np.ceil(filtered_df["c_x"].max())
    bin_width = 10
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["c_x"], bins=bins, kde=False,
                 edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("纬向本征相速度", fontsize=24)
    ax.set_xlabel("Zonal phase speed (m/s)")
    ax.set_ylabel("Occurrence (%)")


def plot_c_y(filtered_df, ax):
    min_val = np.floor(filtered_df["c_y"].min())
    max_val = np.ceil(filtered_df["c_y"].max())
    bin_width = 10
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["c_y"], bins=bins, kde=False,
                 edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("经向本征相速度", fontsize=24)
    ax.set_xlabel("Meridional phase speed (m/s)")
    ax.set_ylabel("Occurrence (%)")


def plot_c_z(filtered_df, ax):
    min_val = filtered_df["c_z"].min()
    max_val = filtered_df["c_z"].max()
    bin_width = 0.1
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["c_z"], bins=bins, kde=False,
                 edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("垂直本征相速度", fontsize=24)
    ax.set_xlabel("Vertical phase speed (m/s)")
    ax.set_ylabel("Occurrence (%)")


def plot_u1(filtered_df, ax):
    min_val = np.floor(filtered_df["u1"].min())
    max_val = np.ceil(filtered_df["u1"].max())
    bin_width = 0.5
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["u1"], bins=bins, kde=False,
                 edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title(" ", fontsize=24)
    ax.set_xlabel("Zonal wind amplitude (m/s)")
    ax.set_ylabel("Occurrence (%)")


def plot_v1(filtered_df, ax):
    min_val = np.floor(filtered_df["v1"].min())
    max_val = np.ceil(filtered_df["v1"].max())
    bin_width = 0.5
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["v1"], bins=bins, kde=False,
                 edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("扰动振幅统计结果", fontsize=24)
    ax.set_xlabel("Meridional wind amplitude (m/s)")
    ax.set_ylabel("Occurrence (%)")


def plot_T1(filtered_df, ax):
    min_val = np.floor(filtered_df["T1"].min())
    max_val = np.ceil(filtered_df["T1"].max())
    bin_width = 0.5
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df["T1"], bins=bins, kde=False,
                 edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title(" ", fontsize=24)
    ax.set_xlabel("Temperature amplitude (K)")
    ax.set_ylabel("Occurrence (%)")


def plot_MFu(filtered_df1, ax):
    min_val = np.floor(filtered_df1["MFu"].min())
    max_val = np.ceil(filtered_df1["MFu"].max())
    bin_width = 0.1
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df1["MFu"], bins=bins, kde=False,
                 edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("纬向动量通量统计结果", fontsize=24)
    ax.set_xlabel("Zonal momentum flux(mPa)")
    ax.set_ylabel("Occurrence(%)")


def plot_MFv(filtered_df1, ax):
    min_val = np.floor(filtered_df1["MFv"].min())
    max_val = np.ceil(filtered_df1["MFv"].max())
    bin_width = 0.1
    bins = np.arange(min_val, max_val + bin_width, bin_width)
    sns.histplot(filtered_df1["MFv"], bins=bins, kde=False,
                 edgecolor="black", stat="percent", ax=ax)
    ax.set_xlim(min_val, max_val)
    ax.set_title("经向动量通量统计结果", fontsize=24)
    ax.set_xlabel("Meridional momentum flux(mPa)")
    ax.set_ylabel("Occurrence(%)")


def plot_horizontal_propagation(filtered_df: pd.DataFrame,  season):
    def get_season(date):
        month = date.month
        if month in [12, 1, 2]:
            return "Winter"
        elif month in [3, 4, 5]:
            return "Spring"
        elif month in [6, 7, 8]:
            return "Summer"
        else:
            return "Fall"
    fig = plt.figure(figsize=(36, 20))
    gs = gridspec.GridSpec(2, 2)
    ax10 = []
    for i in range(4):
        ax10.append(fig.add_subplot(gs[i//2, i % 2], projection="windrose"))

    data = filtered_df.copy()

    data.loc[:, "date"] = data["file_name"].str[-28:-16]
    filtered_df = data[[
        "date"] + [col for col in data.columns if col != "file_name" and col != "date"]]
    filtered_df.reset_index(drop=True, inplace=True)

    filtered_df = filtered_df.drop_duplicates(
        subset="date", keep="last")  # 使用 drop_duplicates 函数，保留每组重复中的最后一个记录

    filtered_df["date1"] = filtered_df["date"].str.split(
        "T").str[0]  # 再加一列，只保留日期部分
    filtered_df["date1"] = pd.to_datetime(
        filtered_df["date1"], format="%Y%m%d")  # 确保 'date1' 列是日期格式
    filtered_df["season"] = filtered_df["date1"].apply(get_season)  # 添加季节列
    seasons = ["Winter", "Spring", "Summer", "Fall"]
    for ax, season in zip(ax10, seasons):
        season_data = filtered_df[filtered_df["season"] == season]
        windrose = WindroseAxes.from_ax(ax)
        ax.set_title(season, fontsize=18)
        windrose.bar(season_data["b"], np.ones_like(
            season_data["b"]), normed=False)  # normed=True表示占每个季节的占比
    # season_data = filtered_df[filtered_df["season"] == season]
    # windrose = WindroseAxes.from_ax(ax)
    # ax.set_title(season, fontsize=18)
    # windrose.bar(season_data["b"], np.ones_like(
    #     season_data["b"]), normed=False)


def plot_vertical_propagation(filtered_df, ax):
    data = filtered_df.copy()

    data.loc[:, "date"] = data["file_name"].str[-28:-16]
    filtered_df = data[[
        "date"] + [col for col in data.columns if col != "file_name" and col != "date"]]
    filtered_df.reset_index(drop=True, inplace=True)

    filtered_df = filtered_df.drop_duplicates(
        subset="date", keep="last")  # 使用 drop_duplicates 函数，保留每组重复中的最后一个记录

    filtered_df["date1"] = filtered_df["date"].str.split(
        "T").str[0]  # 再加一列，只保留日期部分
    filtered_df["date1"] = pd.to_datetime(
        filtered_df["date1"], format="%Y%m%d")  # 确保 'date1' 列是日期格式

    filtered_df.set_index("date1", inplace=True)
    monthly_stats_df = (
        filtered_df.groupby([filtered_df.index.month, filtered_df.index.year])
        .apply(lambda x: pd.Series({"Upload (%)": (x["a"] == 1).mean() * 100, "Downward (%)": (x["a"] == -1).mean() * 100}))
        .reset_index(level=1, drop=True)
    )
    monthly_avg_stats_df = monthly_stats_df.groupby(level=0).mean()
    dates = monthly_avg_stats_df.index.to_numpy()
    ax.plot(dates, monthly_avg_stats_df["Upload (%)"].to_numpy(
    ), marker="o", label="Up (%)")
    ax.plot(dates, monthly_avg_stats_df["Downward (%)"].to_numpy(
    ), marker="o", label="Down (%)")
    ax.set_xticks(
        ticks=np.arange(1, 13), labels=["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"], rotation=0
    )
    ax.legend()
    ax.set_title("每月上传/下传重力波占比", fontsize=24)
    ax.set_xlabel("Month")
    ax.set_ylabel("Percentage (%)")


def plot_energy_distribution(filtered_df, ax):
    data = filtered_df.copy()

    data.loc[:, "date"] = data["file_name"].str[-28:-16]
    filtered_df = data[[
        "date"] + [col for col in data.columns if col != "file_name" and col != "date"]]
    filtered_df.reset_index(drop=True, inplace=True)

    filtered_df = filtered_df.drop_duplicates(
        subset="date", keep="last")  # 使用 drop_duplicates 函数，保留每组重复中的最后一个记录

    filtered_df["date1"] = filtered_df["date"].str.split(
        "T").str[0]  # 再加一列，只保留日期部分
    filtered_df["date1"] = pd.to_datetime(
        filtered_df["date1"], format="%Y%m%d")  # 确保 'date1' 列是日期格式

    filtered_df.reset_index(inplace=True)
    filtered_df["year_month"] = filtered_df["date1"].dt.to_period("M")
    monthly_avg = filtered_df.groupby("year_month")[["Ek", "E_p"]].mean()
    full_range = pd.period_range(
        start=monthly_avg.index.min(), end=monthly_avg.index.max(), freq="M")
    full_monthly_avg = pd.DataFrame(index=full_range).join(monthly_avg)
    full_monthly_avg["Ek"] = pd.to_numeric(
        full_monthly_avg["Ek"], errors="coerce")
    full_monthly_avg["E_p"] = pd.to_numeric(
        full_monthly_avg["E_p"], errors="coerce")
    ax.plot(full_monthly_avg.index.values.astype(
        str), full_monthly_avg["Ek"].values, marker="o", linestyle="-", color="r", label="动能月平均值")
    ax.plot(full_monthly_avg.index.values.astype(
        str), full_monthly_avg["E_p"].values, marker="o", linestyle="-", color="b", label="势能月平均值")
    ax.set_title("动能和势能分布情况", fontsize=24)
    ax.set_xlabel("Month", fontsize=14)
    ax.set_ylabel("Wave energy (J/kg)", fontsize=14)
    months = full_monthly_avg.index.values.astype(str)
    june_indices = [i for i, date in enumerate(months) if date.endswith("-06")]
    december_indices = [i for i, date in enumerate(
        months) if date.endswith("-12")]
    selected_indices = june_indices + december_indices
    ax.set_xticks(ticks=selected_indices, labels=[
                  months[i] for i in selected_indices], rotation=45)
    ax.legend()


# def render_based_on_mode(df, mode, seaon = None):
def render_based_on_mode(df, mode, season=None):
    buf = BytesIO()
    fig, ax = plt.subplots(figsize=(10, 6))
    if mode == "w/f值统计结果":
        plot_w_f(df, ax)
    elif mode == "周期统计结果":
        plot_zhou_qi(df, ax)
    elif mode == "垂直波长分布":
        plot_ver_wave_len(df, ax)
    elif mode == "水平波长分布":
        plot_hori_wave_len(df, ax)
    elif mode == "纬向本征相速度":
        plot_c_x(df, ax)
    elif mode == "经向本征相速度":
        plot_c_y(df, ax)
    elif mode == "垂直本征相速度":
        plot_c_z(df, ax)
    elif mode == "Zonal wind amplitude (m/s)":
        plot_u1(df, ax)
    elif mode == "扰动振幅统计结果":
        plot_v1(df, ax)
    elif mode == "Temperature amplitude (K)":
        plot_T1(df, ax)
    elif mode == "纬向动量通量统计结果":
        plot_MFu(df, ax)
    elif mode == "经向动量通量统计结果":
        plot_MFv(df, ax)
    elif mode == "horizontal propagation":
        plot_horizontal_propagation(df, season="Fall")
    elif mode == "每月上传/下传重力波占比":
        plot_vertical_propagation(df, ax)
    elif mode == "动能和势能分布情况":
        plot_energy_distribution(df, ax)
    else:
        raise ValueError("Invalid mode")
    plt.rcParams["font.sans-serif"] = ["SimHei"]  # 显示中文
    plt.savefig(buf)
    buf.seek(0)
    plt.close()
    return buf
#     if mode == ""


def get_all_modes():
    return [
        "w/f值统计结果",
        "周期统计结果",
        "垂直波长分布",
        "水平波长分布",
        "纬向本征相速度",
        "经向本征相速度",
        "垂直本征相速度",
        "Zonal wind amplitude (m/s)",
        "扰动振幅统计结果",
        "Temperature amplitude (K)",
        "纬向动量通量统计结果",
        "经向动量通量统计结果",
        "horizontal propagation",
        "每月上传/下传重力波占比",
        "动能和势能分布情况",
    ]
