import copy

import matplotlib
import matplotlib.cm as cm
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import numpy as np

import emout.utils as utils

_r = 0.9
_d = 0.5
mycmap = mcolors.LinearSegmentedColormap(
    "gray-jet",
    {
        "red": (
            (0.00, 0.2, 0.2),
            (_d * (1 - _r), 0.3, 0.3),
            (0.35 * _r + (1 - _r), 0, 0),
            (0.66 * _r + (1 - _r), 1, 1),
            (0.89 * _r + (1 - _r), 1, 1),
            (1.00, 0.5, 0.5),
        ),
        "green": (
            (0.00, 0.2, 0.2),
            (_d * (1 - _r), 0.3, 0.3),
            (0.125 * _r + (1 - _r), 0, 0),
            (0.375 * _r + (1 - _r), 1, 1),
            (0.640 * _r + (1 - _r), 1, 1),
            (0.910 * _r + (1 - _r), 0, 0),
            (1.000, 0, 0),
        ),
        "blue": (
            (0.00, 0.2, 0.2),
            (_d * (1 - _r), 0.3, 0.3),
            (0.00 * _r + (1 - _r), 0.5, 0.5),
            (0.11 * _r + (1 - _r), 1, 1),
            (0.34 * _r + (1 - _r), 1, 1),
            (0.65 * _r + (1 - _r), 0, 0),
            (1.00, 0, 0),
        ),
    },
)


def figsize_with_2d(data2d, dpi=10):
    """2次元データから図のサイズを計算する.

    Parameters
    ----------
    data2d : numpy.ndarray
        2次元データ
    dpi : int, optional
        1データを何pixelで表すか, by default 10

    Returns
    -------
    (float, float)
        図のサイズ
    """
    px = 1 / plt.rcParams["figure.dpi"] * dpi
    figsize = (data2d.shape[1] * px, data2d.shape[0] * px)
    return figsize


def plot_2dmap(
    data2d,
    mesh=None,
    savefilename=None,
    cmap=mycmap,
    mask_color="gray",
    vmin=None,
    vmax=None,
    figsize=None,
    xlabel=None,
    ylabel=None,
    title=None,
    interpolation="bilinear",
    dpi=10,
    colorbar_label="",
    cbargs={},
    add_colorbar=True,
    **kwargs
):
    """2次元カラーマップをプロットする.

    Parameters
    ----------
    data2d : numpy.ndarray
        2次元データ
    mesh : (numpy.ndarray, numpy.ndarray), optional
        メッシュ, by default None
    savefilename : str, optional
        保存するファイル名(Noneの場合保存しない), by default None
    cmap : matplotlib.Colormap or str or None, optional
        カラーマップ, by default cm.coolwarm
    mask_color : str
        マスクされた位置の色, by default 'gray'
    vmin : float, optional
        最小値, by default None
    vmax : float, optional
        最大値, by default None
    figsize : (float, float), optional
        図のサイズ, by default None
    xlabel : str, optional
        x軸のラベル, by default None
    ylabel : str, optional
        y軸のラベル, by default None
    title : str, optional
        タイトル, by default None
    interpolation : str, optional
        用いる補間方法, by default 'bilinear'
    dpi : int, optional
        解像度(figsizeが指定された場合は無視される), by default 10
    add_colorbar: bool, optional
        カラーバーを追加する場合True, by default True

    Returns
    -------
    AxesImage or None
        プロットしたimageデータ(保存した場合None)
    """
    if mesh is None:
        x = list(range(data2d.shape[1]))
        y = list(range(data2d.shape[0]))
        mesh = np.meshgrid(x, y)

    if cmap is not None:
        if isinstance(cmap, str):
            cmap = copy.copy(plt.get_cmap(str(cmap)))
        else:
            cmap = copy.copy(cmap)
        cmap.set_bad(color=mask_color)

    extent = [mesh[0][0, 0], mesh[0][-1, -1], mesh[1][0, 0], mesh[1][-1, -1]]
    img = plt.imshow(
        data2d,
        interpolation=interpolation,
        cmap=cmap,
        origin="lower",
        vmin=vmin,
        vmax=vmax,
        extent=extent,
        aspect="auto",
        **kwargs
    )
    if add_colorbar:
        if "cb" in cbargs:
            cb = plt.colorbar(label=colorbar_label, **cbargs["cb"])
        else:
            cb = plt.colorbar(label=colorbar_label)

    if "others" in cbargs and "yticklabels" in cbargs["others"]:
        cb.ax.set_yticklabels(cbargs["others"]["yticklabels"])

    if title is not None:
        plt.title(title)
    if xlabel is not None:
        plt.xlabel(xlabel)
    if ylabel is not None:
        plt.ylabel(ylabel)

    if savefilename is not None:
        plt.gcf().savefig(savefilename)
        plt.close(plt.gcf())
        return None
    else:
        return img


def plot_2d_contour(
    data2d,
    mesh=None,
    levels=None,
    colors=["black"],
    cmap=None,
    alpha=1,
    vmin=None,
    vmax=None,
    savefilename=None,
    figsize=None,
    xlabel=None,
    ylabel=None,
    title=None,
    dpi=10,
    fmt="%1.1f",
    fontsize=12,
    **kwargs
):
    """2次元等高線をプロットする.

    Parameters
    ----------
    data2d : numpy.ndarray
        2次元データ
    mesh : (numpy.ndarray, numpy.ndarray), optional
        メッシュ, by default None
    levels : int
        等高線数, by default None
    alpha : float
        透明度(0.0~1.0), by default 1
    savefilename : str, optional
        保存するファイル名(Noneの場合保存しない), by default None
    cmap : matplotlib.Colormap or str or None, optional
        カラーマップ, by default None
    mask_color : str
        マスクされた位置の色, by default 'gray'
    vmin : float, optional
        最小値, by default None
    vmax : float, optional
        最大値, by default None
    figsize : (float, float), optional
        図のサイズ, by default None
    xlabel : str, optional
        x軸のラベル, by default None
    ylabel : str, optional
        y軸のラベル, by default None
    title : str, optional
        タイトル, by default None
    interpolation : str, optional
        用いる補間方法, by default 'bilinear'
    dpi : int, optional
        解像度(figsizeが指定された場合は無視される), by default 10
    fmt : str
        clabelの形式, by default '%1.1f'
    fontsize : str
        clabelのフォントサイズ, by default 12

    Returns
    -------
    AxesImage or None
        プロットしたimageデータ(保存した場合None)
    """
    if mesh is None:
        x = list(range(data2d.shape[1]))
        y = list(range(data2d.shape[0]))
        mesh = np.meshgrid(x, y)

    kwargs = {
        "alpha": alpha,
        "vmin": vmin,
        "vmax": vmax,
    }
    if cmap is None:
        kwargs["colors"] = colors
    else:
        kwargs["cmap"] = cmap
    if levels is not None:
        kwargs["levels"] = levels
    cont = plt.contour(*mesh, data2d, **kwargs)
    cont.clabel(fmt=fmt, fontsize=fontsize)

    if title is not None:
        plt.title(title)
    if xlabel is not None:
        plt.xlabel(xlabel)
    if ylabel is not None:
        plt.ylabel(ylabel)

    if savefilename is not None:
        plt.gcf().savefig(savefilename)
        plt.close(plt.gcf())
        return None
    else:
        return cont


def plot_surface(
    x,
    y,
    z,
    value,
    ax3d=None,
    add_colorbar=False,
    savefilename=None,
    cmap=cm.jet,
    mask_color="gray",
    vmin=None,
    vmax=None,
    figsize=None,
    xlabel=None,
    ylabel=None,
    zlabel=None,
    title=None,
    ninterp=1,
    function="linear",
    dpi=10,
    colorbar_label="",
):
    """3次元表面プロットをする.

    Parameters
    ----------
    x : (numpy.ndarray, numpy.ndarray), optional
        x座標のメッシュ
    y : (numpy.ndarray, numpy.ndarray), optional
        y座標のメッシュ
    z : (numpy.ndarray, numpy.ndarray), optional
        z座標のメッシュ
    val : (numpy.ndarray, numpy.ndarray), optional
        値のメッシュ
    ax3d : Axes3D
        Axes3Dオブジェクト, by default None
    savefilename : str, optional
        保存するファイル名(Noneの場合保存しない), by default None
    cmap : matplotlib.Colormap or str or None, optional
        カラーマップ, by default cm.coolwarm
    vmin : float, optional
        最小値, by default None
    vmax : float, optional
        最大値, by default None
    figsize : (float, float), optional
        図のサイズ, by default None
    xlabel : str, optional
        x軸のラベル, by default None
    ylabel : str, optional
        y軸のラベル, by default None
    zlabel : str, optional
        z軸のラベル, by default None
    title : str, optional
        タイトル, by default None
    dpi : int, optional
        解像度(figsizeが指定された場合は無視される), by default 10

    Returns
    -------
    AxesImage or None
        プロットしたimageデータ(保存した場合None)
    """
    if savefilename is not None:
        if figsize is None:
            fig = plt.figure()
        else:
            if figsize == "auto":
                figsize = figsize_with_2d(x, dpi=dpi)
            fig = plt.figure(figsize=figsize)
    else:
        fig = plt.gcf()

    if ax3d is None:
        ax3d = fig.gca(projection="3d")

    if cmap is not None:
        if isinstance(cmap, str):
            cmap = copy.copy(cm.get_cmap(str(cmap)))
        else:
            cmap = copy.copy(cmap)
        cmap.set_bad(color=mask_color)

    if ninterp is not None:
        x = utils.interp2d(x, ninterp, method=function)
        y = utils.interp2d(y, ninterp, method=function)
        z = utils.interp2d(z, ninterp, method=function)
        value = utils.interp2d(value, ninterp)

    if vmin is None:
        vmin = value.min()
    if vmax is None:
        vmax = value.max()

    norm = matplotlib.colors.Normalize(vmin, vmax)
    mappable = cm.ScalarMappable(cmap=cmap, norm=norm)
    mappable.set_array([])
    value_colors = mappable.to_rgba(value)

    surf = ax3d.plot_surface(
        x, y, z, facecolors=value_colors, vmin=vmin, vmax=vmax, shade=False
    )
    if add_colorbar:
        plt.colorbar(mappable, ax=ax3d, label=colorbar_label)

    if title is not None:
        ax3d.set_title(title)
    if xlabel is not None:
        ax3d.set_xlabel(xlabel)
    if ylabel is not None:
        ax3d.set_ylabel(ylabel)
    if zlabel is not None:
        ax3d.set_zlabel(zlabel)

    if savefilename is not None:
        fig.savefig(savefilename)
        plt.close(fig)
        return None
    else:
        return surf


def plot_line(
    data1d,
    x=None,
    savefilename=None,
    vmin=None,
    vmax=None,
    figsize=None,
    xlabel=None,
    ylabel=None,
    label=None,
    title=None,
    **kwargs
):
    """1次元データをプロットする.

    Parameters
    ----------
    data1d : array-like or scalar
        プロットする1次元データ
    x : array-like or scalar
        横軸となる1次元データ, by default None
    savefilename : str, optional
        保存するファイル名, by default None
    vmin : float, optional
        最小値, by default None
    vmax : float, optional
        最大値, by default None
    figsize : (float, float), optional
        図のサイズ, by default None
    xlabel : str, optional
        横軸のラベル, by default None
    ylabel : str, optional
        縦軸のラベル, by default None
    label : str, optional
        ラベル, by default None
    title : str, optional
        タイトル, by default None

    Returns
    -------
    Line2D or None
        プロットデータを表す線オブジェクト(保存した場合None)
    """
    if savefilename is not None:
        if figsize is None:
            fig = plt.figure()
        else:
            fig = plt.figure(figsize=figsize)

    if x is None:
        line = plt.plot(data1d, label=label, **kwargs)
    else:
        line = plt.plot(x, data1d, label=label, **kwargs)
    plt.ylim([vmin, vmax])

    if title is not None:
        plt.title(title)
    if xlabel is not None:
        plt.xlabel(xlabel)
    if ylabel is not None:
        plt.ylabel(ylabel)

    if savefilename is not None:
        fig.savefig(savefilename)
        plt.close(fig)
        return None
    else:
        return line


def plot_2d_vector(
    x_data2d,
    y_data2d,
    mesh=None,
    savefilename=None,
    color=None,
    scale=1,
    scaler="standard",
    skip=1,
    easy_to_read=True,
    figsize=None,
    xlabel=None,
    ylabel=None,
    title=None,
    dpi=10,
    cmap=None,
):
    """2次元ベクトル図をプロットする.

    Parameters
    ----------
    x_data2d, y_data2d : numpy.ndarray
        2次元データ
    mesh : (numpy.ndarray, numpy.ndarray), optional
        メッシュ, by default None
    savefilename : str, optional
        保存するファイル名(Noneの場合保存しない), by default None
    color : str
        ベクトルの色, by default None
    scale : float
        ベクトルの大きさ係数(最終的な大きさにこの値を掛ける), by default 1
    skip : int
        プロットするデータ間隔, by default 1
    easy_to_read : bool
        ベクトルを見やすい大きさにスケーリングするならTrue, by default True
    figsize : (float, float), optional
        図のサイズ, by default None
    xlabel : str, optional
        x軸のラベル, by default None
    ylabel : str, optional
        y軸のラベル, by default None
    title : str, optional
        タイトル, by default None
    interpolation : str, optional
        用いる補間方法, by default 'bilinear'
    dpi : int, optional
        解像度(figsizeが指定された場合は無視される), by default 10

    Returns
    -------
    AxesImage or None
        プロットしたimageデータ(保存した場合None)
    """
    fig = None
    if savefilename is not None:
        if figsize is None:
            fig = plt.figure()
        else:
            if figsize == "auto":
                figsize = figsize_with_2d(x_data2d, dpi=dpi)
            fig = plt.figure(figsize=figsize)

    if mesh is None:
        x = list(range(x_data2d.shape[1]))
        y = list(range(x_data2d.shape[0]))
        mesh = np.meshgrid(x, y)

    x = mesh[0]
    y = mesh[1]
    U = np.array(x_data2d)
    V = np.array(y_data2d)

    x_skip = skip if type(skip) == int else skip[0]
    y_skip = skip if type(skip) == int else skip[1]
    x = x[::y_skip, ::x_skip]
    y = y[::y_skip, ::x_skip]
    U = U[::y_skip, ::x_skip]
    V = V[::y_skip, ::x_skip]

    norm = np.sqrt(U**2 + V**2)

    if scaler == "standard":
        norm_max = np.nanmax(np.abs(norm))
        U /= norm_max
        V /= norm_max

    elif scaler == "normal":
        U /= norm
        V /= norm

    elif scaler == "log":
        U = U / norm * np.log(norm + 1)
        V = V / norm * np.log(norm + 1)

    # 見やすい大きさに線形スケーリングを行う
    if easy_to_read:
        dx = (x.max() - x.min()) / x.shape[0]
        multiplier = dx * 1.2
        norm_mean = np.nanmean(np.sqrt(U**2 + V**2))
        U *= scale / norm_mean * multiplier
        V *= scale / norm_mean * multiplier

    if cmap is None:
        img = plt.quiver(
            x,
            y,
            U,
            V,
            angles="xy",
            scale_units="xy",
            scale=1,
        )
    else:
        img = plt.quiver(
            x,
            y,
            U,
            V,
            np.sqrt(U**2 + V**2),
            angles="xy",
            scale_units="xy",
            scale=1,
            cmap=cmap,
        )

    if title is not None:
        plt.title(title)
    if xlabel is not None:
        plt.xlabel(xlabel)
    if ylabel is not None:
        plt.ylabel(ylabel)

    if savefilename is not None:
        fig.savefig(savefilename)
        plt.close(fig)
        return None
    else:
        return img


def plot_2d_streamline(
    x_data2d,
    y_data2d,
    mesh=None,
    savefilename=None,
    skip=1,
    figsize=None,
    xlabel=None,
    ylabel=None,
    title=None,
    dpi=10,
    color=None,
    cmap=None,
    norm="linear",
    vmin=None,
    vmax=None,
    density=1,
):
    """2次元ベクトル図をプロットする.

    Parameters
    ----------
    x_data2d, y_data2d : numpy.ndarray
        2次元データ
    mesh : (numpy.ndarray, numpy.ndarray), optional
        メッシュ, by default None
    savefilename : str, optional
        保存するファイル名(Noneの場合保存しない), by default None
    color : str
        ベクトルの色, by default None
    scale : float
        ベクトルの大きさ係数(最終的な大きさにこの値を掛ける), by default 1
    skip : int
        プロットするデータ間隔, by default 1
    easy_to_read : bool
        ベクトルを見やすい大きさにスケーリングするならTrue, by default True
    figsize : (float, float), optional
        図のサイズ, by default None
    xlabel : str, optional
        x軸のラベル, by default None
    ylabel : str, optional
        y軸のラベル, by default None
    title : str, optional
        タイトル, by default None
    interpolation : str, optional
        用いる補間方法, by default 'bilinear'
    dpi : int, optional
        解像度(figsizeが指定された場合は無視される), by default 10

    Returns
    -------
    AxesImage or None
        プロットしたimageデータ(保存した場合None)
    """
    fig = None
    if savefilename is not None:
        if figsize is None:
            fig = plt.figure()
        else:
            if figsize == "auto":
                figsize = figsize_with_2d(x_data2d, dpi=dpi)
            fig = plt.figure(figsize=figsize)

    if mesh is None:
        x = list(range(x_data2d.shape[1]))
        y = list(range(x_data2d.shape[0]))
        mesh = np.meshgrid(x, y)

    x = mesh[0]
    y = mesh[1]
    U = np.array(x_data2d)
    V = np.array(y_data2d)

    x_skip = skip if type(skip) == int else skip[0]
    y_skip = skip if type(skip) == int else skip[1]
    x = x[::y_skip, ::x_skip]
    y = y[::y_skip, ::x_skip]
    U = U[::y_skip, ::x_skip]
    V = V[::y_skip, ::x_skip]

    if cmap:
        length = np.sqrt(U**2 + V**2)
        vmin = vmin or length.min()
        vmax = vmax or length.max()

        if norm == "linear":
            norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
        elif norm == "log":
            norm = mcolors.LogNorm(vmin=vmin, vmax=vmax)
        elif norm == "centered":
            norm = mcolors.CenteredNorm()
        elif norm == "symlog":
            norm = mcolors.SymLogNorm(vmin=vmin, vmax=vmax)

        img = plt.streamplot(
            x,
            y,
            U,
            V,
            color=length,
            cmap=cmap,
            norm=norm,
            density=density,
        )
    else:
        img = plt.streamplot(
            x,
            y,
            U,
            V,
            color=color,
            density=density,
        )

    if title is not None:
        plt.title(title)
    if xlabel is not None:
        plt.xlabel(xlabel)
    if ylabel is not None:
        plt.ylabel(ylabel)

    if savefilename is not None:
        fig.savefig(savefilename)
        plt.close(fig)
        return None
    else:
        return img


def plot_3d_quiver(
    x_data3d,
    y_data3d,
    z_data3d,
    ax3d=None,
    mesh=None,
    savefilename=None,
    color=None,
    scale=1,
    scaler="standard",
    skip=1,
    easy_to_read=True,
    figsize=None,
    xlabel=None,
    ylabel=None,
    title=None,
    dpi=10,
    cmap=None,
):
    """2次元ベクトル図をプロットする.

    Parameters
    ----------
    x_data2d, y_data2d : numpy.ndarray
        2次元データ
    mesh : (numpy.ndarray, numpy.ndarray), optional
        メッシュ, by default None
    savefilename : str, optional
        保存するファイル名(Noneの場合保存しない), by default None
    color : str
        ベクトルの色, by default None
    scale : float
        ベクトルの大きさ係数(最終的な大きさにこの値を掛ける), by default 1
    skip : int
        プロットするデータ間隔, by default 1
    easy_to_read : bool
        ベクトルを見やすい大きさにスケーリングするならTrue, by default True
    figsize : (float, float), optional
        図のサイズ, by default None
    xlabel : str, optional
        x軸のラベル, by default None
    ylabel : str, optional
        y軸のラベル, by default None
    title : str, optional
        タイトル, by default None
    interpolation : str, optional
        用いる補間方法, by default 'bilinear'
    dpi : int, optional
        解像度(figsizeが指定された場合は無視される), by default 10

    Returns
    -------
    AxesImage or None
        プロットしたimageデータ(保存した場合None)
    """
    fig = None
    if savefilename is not None:
        if figsize is None:
            fig = plt.figure()
        else:
            if figsize == "auto":
                figsize = figsize_with_2d(x_data3d[:, 0, :], dpi=dpi)
            fig = plt.figure(figsize=figsize)

    if ax3d is None:
        ax3d = fig.add_subplot(projection="3d")

    if mesh is None:
        x = list(range(x_data3d.shape[1]))
        y = list(range(x_data3d.shape[0]))
        z = list(range(x_data3d.shape[0]))
        mesh = np.meshgrid(x, y, z)

    x = mesh[0]
    y = mesh[1]
    z = mesh[2]
    U = np.array(x_data3d)
    V = np.array(y_data3d)
    W = np.array(z_data3d)

    x_skip = skip if type(skip) == int else skip[0]
    y_skip = skip if type(skip) == int else skip[1]
    z_skip = skip if type(skip) == int else skip[2]
    x = x[::z_skip, ::y_skip, ::x_skip]
    y = y[::z_skip, ::y_skip, ::x_skip]
    z = z[::z_skip, ::y_skip, ::x_skip]
    U = U[::z_skip, ::y_skip, ::x_skip]
    V = V[::z_skip, ::y_skip, ::x_skip]
    W = W[::z_skip, ::y_skip, ::x_skip]

    norm = np.sqrt(U**2 + V**2)

    if scaler == "standard":
        norm_max = np.nanmax(np.abs(norm))
        U /= norm_max
        V /= norm_max
        W /= norm_max

    elif scaler == "normal":
        U /= norm
        V /= norm
        W /= norm

    elif scaler == "log":
        U = U / norm * np.log(norm + 1)
        V = V / norm * np.log(norm + 1)
        W = W / norm * np.log(norm + 1)

    # 見やすい大きさに線形スケーリングを行う
    if easy_to_read:
        dx = (x.max() - x.min()) / x.shape[0]
        multiplier = dx * 1.2
        norm_mean = np.nanmean(np.sqrt(U**2 + V**2))
        U *= scale / norm_mean * multiplier
        V *= scale / norm_mean * multiplier
        W *= scale / norm_mean * multiplier

    if cmap is None:
        img = ax3d.quiver(
            x,
            y,
            z,
            U,
            V,
            W,
            angles="xy",
            scale_units="xy",
            scale=1,
        )
    else:
        img = ax3d.quiver(
            x,
            y,
            z,
            U,
            V,
            W,
            np.sqrt(U**2 + V**2 + W**2),
            angles="xy",
            scale_units="xy",
            scale=1,
            cmap=cmap,
        )

    if title is not None:
        plt.title(title)
    if xlabel is not None:
        plt.xlabel(xlabel)
    if ylabel is not None:
        plt.ylabel(ylabel)

    if savefilename is not None:
        fig.savefig(savefilename)
        plt.close(fig)
        return None
    else:
        return img
