import matplotlib as mpl
from .indicator import envelope_of_expected_error, root_mean_square_error, r2_score
import numpy as np
from mtmtool.io import auto_make_dirs
from matplotlib import pyplot as plt

style_dict = {
    "single_normal": {"figsize": (3.538, 3.538), "fontsize": 10.5, "ticksize": 9},
    "single_bold": {"figsize": (5, 4.5), "fontsize": 12, "ticksize": 10.5},
    # "double_normal": {"figsize": (6.7, 5)},
    # "double_bold": {"figsize": (7.48, 6)},
}

def density_chart(x, y, xlabel="AOD Measured", ylabel="AOD Retrieved", title=None, type="kernels", bins=25, xy_max_edge=2, save_path=None, dpi=300, figsize=(3, 3), style=None, fontsize=9, ticksize=8):
    """绘制密度图, 来比较标签和预测值的分布情况, 使用高斯核密度(kernels)来绘制密度图 或 使用直方图(scatter or grid)来统计单位面积内的点的数量

    Parameters
    ----------
    x : list
        label
    y : list
        predict
    xlabel : str, optional
        x轴标签, by default "AOD Measured"
    ylabel : str, optional
        y轴标签, by default "AOD Retrieved"
    title : str, optional
        标题, by default None
    type : str, {"kernels", "scatter", "grid"}, optional
        要绘制的图的类型, by default "kernels"
    bins : int, optional
        直方图统计时, 每1个单位有多少个箱子, 每个箱子单位刻度为(1/bins), by default 25
    xy_max_edge : int, optional
        最大的xy显示边界, by default 2
    save_path : _type_, optional
        保存路径, 如果为None, 则直接显示, by default None
    dpi : int, optional
        分辨率, 每英寸下多少个像元, by default 300
    figsize : tuple, optional
        画布大小, 单位是英寸, by default (3, 3)
    style : _type_, optional
        风格样式, by default None
    fontsize : int, optional
        标题等的字体大小, by default 9
    ticksize : int, optional
        tick等的字体大小, by default 8

    """
    if style in style_dict:
        figsize = style_dict[style].get("figsize", figsize)
        fontsize = style_dict[style].get("fontsize", fontsize)
        ticksize = style_dict[style].get("ticksize", ticksize)
    with mpl.rc_context():
        # 设置字体
        font = {'family': 'serif', 'serif': 'Times New Roman', 'weight': 'normal', 'size': fontsize}
        mpl.rc('font', **font)
        mpl.rcParams['mathtext.fontset'] = 'stix'

        # 设置布局
        mpl.rcParams["savefig.bbox"] = "tight"
        mpl.rcParams['figure.constrained_layout.use'] = True

        # Set x axis
        mpl.rcParams['xtick.direction'] = "in"
        mpl.rcParams['xtick.major.size'] = 3
        mpl.rcParams['xtick.major.width'] = 0.5
        mpl.rcParams['xtick.minor.size'] = 1.5
        mpl.rcParams['xtick.minor.width'] = 0.5
        mpl.rcParams['xtick.minor.visible'] = True
        mpl.rcParams['xtick.top'] = True

        # Set y axis
        mpl.rcParams['ytick.direction'] = "in"
        mpl.rcParams['ytick.major.size'] = 3
        mpl.rcParams['ytick.major.width'] = 0.5
        mpl.rcParams['ytick.minor.size'] = 1.5
        mpl.rcParams['ytick.minor.width'] = 0.5
        mpl.rcParams['ytick.minor.visible'] = True
        mpl.rcParams['ytick.right'] = True

        # 设置分辨率
        mpl.rcParams['figure.dpi'] = dpi  # plt.show显示分辨率
        mpl.rcParams["savefig.dpi"] = dpi # plt.savefig保存分辨率

        # 设置画布
        fig, axes = plt.subplots(1, 1, figsize=figsize)

        # 设置其它参数
        axes.set_xlabel(xlabel, labelpad=2)
        axes.set_ylabel(ylabel, labelpad=2)
        axes.set_title(title)
        axes.set_xlim(-0.05, xy_max_edge + 0.05)
        axes.set_ylim(-0.05, xy_max_edge + 0.05)
        axes.set_aspect("equal")
        axes.grid(linestyle=":", color="r", alpha=0.1)
        xyticks = [i/2 for i in range(int(xy_max_edge//0.5) + 1)]
        xyticks = [i for i in xyticks if i <= xy_max_edge]
        xytickslables = [f"{format(i, '.1f')}" for i in xyticks]
        axes.set_xticks(xyticks)
        axes.set_yticks(xyticks)
        axes.set_xticklabels(xytickslables)
        axes.set_yticklabels(xytickslables)
        axes.tick_params(labelsize=ticksize)

        # 将数据转为numpy数组
        x, y = np.array(x).reshape(-1), np.array(y).reshape(-1) # x means label, y means predict
        mask = (x <= xy_max_edge) & (y <= xy_max_edge)
        print(f"数据总数: {len(x)}, 有效数据总数: {mask.sum()}, 越界数据总数: {(~mask).sum()}")

        if type in ["kernels"]:
            from scipy.stats import gaussian_kde
            xy = np.vstack([x, y])
            z = gaussian_kde(xy)(xy)
            ax_data = axes.scatter(x, y, marker = 'o', c = z, s = 20, cmap = 'jet', vmin = 0)

        elif type in ["scatter", "grid"]:
            # 计算数据的直方图分布
            xy_max_edge = xy_max_edge
            nbins = xy_max_edge * bins
            x = x[mask]
            y = y[mask]
            H, xedges, yedges = np.histogram2d(x, y, bins=nbins, range=[[0, xy_max_edge], [0, xy_max_edge]])

            # 生成散点图图层
            if type == "scatter":
                x_index = (x * bins).astype("int")
                y_index = (y * bins).astype("int")
                z = H[x_index, y_index]
                # 获取散点大小
                axis_height = axes.xaxis.clipbox.height # 获取坐标轴高度
                axis_length = axes.get_xlim()[1] - axes.get_xlim()[0] # 获取坐标轴长度
                scatter_point_size = int(axis_height/(bins*(axis_length)))
                s = (scatter_point_size**2)/4
                # 画散点图
                sort_index = np.argsort(z)
                x = x[sort_index]
                y = y[sort_index]
                z = z[sort_index]
                ax_data = axes.scatter(x, y, marker='.', c=z, s=s, label="AOD scatter", cmap="jet", linewidths=0)
            elif type == "grid": # 绘制方格图
                # H needs to be rotated and flipped
                H = np.rot90(H)
                H = np.flipud(H)
                # Mask zeros
                Hmasked = np.ma.masked_where(H == 0, H)  # Mask pixels with a value of zero
                ax_data = axes.pcolormesh(xedges, yedges, Hmasked, cmap="jet", vmin=1, linewidth=0)
        else:
            raise ValueError(f"没有type: {type}参数对应的画图模式")
        
        # 生成包络线图层
        axes.plot([-1 / 3, xy_max_edge+1], [-1 / 3, (xy_max_edge+1) - (0.05 + 0.15 * (xy_max_edge+1))], linestyle='--', color='black', label='EE envelopes', linewidth = 0.5)
        axes.plot([-1 / 3, xy_max_edge], [-1 / 3, xy_max_edge + (0.05 + 0.15 * xy_max_edge)], linestyle='--', color='black', linewidth = 0.6)
        axes.plot([-1 / 3, xy_max_edge+1], [-1 / 3, xy_max_edge+1], linestyle='-', color='black', linewidth = 0.5)

        # 生成文本图层
        # # 计算指标
        above_percent, below_percent, within_percent = envelope_of_expected_error(x, y)
        rmse = root_mean_square_error(x, y)
        r2 = r2_score(x, y) if len(x) > 1 else 0
        add_text = "\n".join([
            f"$EE$ envelopes: ±(0.05+15%)",
            f"     {format(within_percent, '.1f')}% within $EE$",
            f"     {format(above_percent, '.1f')}% above $EE$",
            f"     {format(below_percent, '.1f')}% below $EE$",
            f"$RMSE$ = {format(rmse, '.3f')}",
            f"$R^2$ = {format(r2, '.3f')}",
            f"$N$ = {len(x)}",
        ])
        axes.text(x=0, y=xy_max_edge, s=add_text, fontsize=ticksize, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0, linewidth=0.2))

        # 生成图例
        lg = axes.legend(loc=4, alignment="right", fontsize=ticksize, labelspacing=0.2, borderpad=0.2, handletextpad=0.2, handlelength=1)
        lg.get_frame().set(linewidth=0.2, edgecolor="k", alpha=0.5)

        # 生成colorbar
        cb = fig.colorbar(ax_data, ax=axes, pad=0.02, fraction=0.05)
        if type in ["scatter", "grid"]:
            cb.ax.set_title(label='Num', loc='center', fontdict={"size": ticksize}, pad=0) #loc参数
            yticks = list(filter(lambda x: x == int(x), cb.ax.get_yticks()))
            cb.ax.set_yticks(yticks)
        cb.ax.tick_params(labelsize=ticksize)
        
        # 画图
        if save_path is not None:
            auto_make_dirs(save_path, isdir = False)
            fig.savefig(save_path, dpi=dpi)
            plt.close(fig)
        else:
            plt.show()