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",
        "每月上传/下传重力波占比",
        "动能和势能分布情况",
    ]