import numpy as np

import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
import seaborn as sns
import datetime
from .plot import image_dB

plt.rcParams['image.origin'] = 'lower'

# Save a palette to a variable:\

colors = ["#003FFF", "#E8000B", "#03ED3A", "#1A1A1A" ] # B, R, G, Black 
colors = ["#E8000B", "#03ED3A", "#003FFF", "#1A1A1A" ] # R, G, B,Black 
colors = ["#003FFF", "#E8000B", "#03ED3A",  "#1A1A1A" ] # Blue, Red, Green, Black

palette = sns.color_palette(colors, 4)


def plot_t_c_matrix(c_t_matrix, title="T", vmin=2, vmax=98, aspect=None, origin="lower", *args, **kwargs):
    C_T_11 = c_t_matrix[:, :, 0, 0]
    im_perc = np.percentile(np.abs(C_T_11), [vmin, vmax])
    if aspect is None:
        aspect = c_t_matrix.shape[1]/c_t_matrix.shape[0]*2

    def plot_sections(c_t_matrix):
        
        fig, axs = plt.subplots(3, 4, constrained_layout=True)
        for idx, ax in zip(range(len(c_t_matrix)), axs[0, :3].flat):
            # print(idx)
            c_idx_abs = np.abs(c_t_matrix[:, :, idx, idx])
            pcm = ax.imshow(c_idx_abs, interpolation='hamming', norm=LogNorm(vmin=im_perc[0], vmax=im_perc[1]), # vmin=.9*c11_db.min() , vmax= 1/3*c11_db.max(), 
                              cmap='gray', aspect=aspect, origin=origin, *args, **kwargs)
            ax.axis('off')
            ax.axis('off')
            ax.set_title("${0}_{{{1}{1}}}$".format(title, idx+1))
            
            # plot hisogram
            ax_hist = axs[0, -1]
            ax_hist.hist(image_dB(c_idx_abs.ravel()), bins='auto', color=palette[idx], histtype='step', label="${0}_{{{1}{1}}}$".format(title, idx+1), density=True, linestyle=('solid'))
            ax_hist.legend(loc='upper left' ) #, bbox_to_anchor=(0.5, 1.05),
            #  ncol=3, fancybox=True, shadow=True)
            ax_hist.set_title("Histogram")
        
        cbar_1 = fig.colorbar(pcm, ax=axs[0, :3], shrink=0.6, location='left')
        cbar_1.set_label('$dB $ ', rotation=0)
        
        # plot of images
        il_lower = np.tril_indices(c_t_matrix.shape[3], -1)
        for idx, nx, ny, ax in zip(range(4), il_lower[0], il_lower[1], axs[1, :3].flat):
            # 
            # print(nx, ny)
            c_idx_abs = np.abs(c_t_matrix[:, :, nx, ny])
            c_idx_angle = np.angle(c_t_matrix[:, :, nx, ny])
            
            pcm = ax.imshow(c_idx_abs, interpolation='hamming', norm=LogNorm(vmin=im_perc[0], vmax=im_perc[1]), # vmin=.9*c11_db.min() , vmax= 1/3*c11_db.max(), 
                              cmap='gray', aspect=aspect, origin=origin, *args, **kwargs)
            ax.axis('off')
            ax.axis('off')
            ax.set_title("${0}_{{{1}{2}}}$".format(title, nx+1, ny+1))
            
            # Histogram
            ax_hist = axs[1, -1]
            ax_hist.hist(image_dB(c_idx_abs.ravel()),
                         bins='auto', color=palette[idx], histtype='step',
                         label="${0}_{{{1}{2}}}$".format(title, nx+1, ny+1),
                         density=True, linestyle=('solid'))

            ax_hist.legend(loc='upper left' ) #, bbox_to_anchor=(0.5, 1.05),
            #  ncol=3, fancybox=True, shadow=True)
            ax_hist.set_title("Histogram")
            
            ax_phase = axs[2, idx]
            pcm_phase = ax_phase.imshow(c_idx_angle, interpolation='hamming', vmin=-np.pi, vmax=np.pi, # vmin=.9*c11_db.min() , vmax= 1/3*c11_db.max(), 
                              cmap='jet', aspect=aspect, origin=origin, *args, **kwargs)
            ax_phase.axis('off')
            ax_phase.axis('off')
            ax_phase.set_title("$\phi_{{ {0}_{{{1}{2}}} }}$".format(title, nx+1, ny+1))
            
            # Histogram
            ax_phase_hist = axs[2, -1]
            ax_phase_hist.hist(c_idx_angle.ravel(),
                         bins='auto', color=palette[idx], histtype='step',
                         label="$\phi_{{ {0}_{{{1}{2}}} }}$".format(title, nx+1, ny+1),
                         density=True, linestyle=('solid')) 
            
            ax_phase_hist.legend(loc='upper left' ) #, bbox_to_anchor=(0.5, 1.05),
            #  ncol=3, fancybox=True, shadow=True)
            ax_phase_hist.set_title("Histogram")
            
        ax_phase_hist.set_xlim([-np.pi, np.pi])
        ax_phase_hist.set_xticks([-np.pi, 0, np.pi])
        ax_phase_hist.set_xticklabels(['$- \pi$', '0', '$\pi$']) 
        
        cbar_2 = fig.colorbar(pcm, ax=axs[1, :3], shrink=0.6, location='left')
        cbar_2.set_label('$dB $ ', rotation=0)
        cbar_3 = fig.colorbar(pcm_phase, ax=axs[2, :3], shrink=0.6, location='left', ticks=[-np.pi, 0, np.pi])
        cbar_3.ax.set_yticklabels(['$- \pi$', '0', '$\pi$']) 
        cbar_3.set_label('$\phi$', rotation=0)
        return fig
    fig = plot_sections(c_t_matrix)
    return fig



def plot_eigenvalues_hist_dB(array_dict, suptitle, vmin=2, vmax=98, aspect=None, origin="lower", *args, **kwargs):
    colors = ["#003FFF", "#E8000B", "#03ED3A",  "#1A1A1A" ] # Blue, Red, Green, Black
    palette = sns.color_palette(colors, 4)
    # sns.palplot(palette)
        
    n_images = len(array_dict.keys())
    fig = plt.figure(constrained_layout=False, figsize=(4 * n_images, 10))
    gs = fig.add_gridspec(nrows=2, ncols=n_images)
    f_img = {}
    f_hist = {}
    ax_img = {}
    f_hist = fig.add_subplot(gs[1, :])

    for idx, key in enumerate(array_dict):
        legend = key
        array = array_dict[key]
        if aspect is None:
            aspect = array.shape[1]/array.shape[0]*2
        
        f_img[idx] = fig.add_subplot(gs[0, idx])
        im_perc = np.percentile(array, [vmin, vmax])

        ax_img[idx] = f_img[idx].imshow(array, aspect=aspect, cmap='jet',
                            vmin=0, vmax=im_perc[1], origin=origin)
        f_img[idx].axis('off')
        f_img[idx].set_title(legend)
        # fig.colorbar(ax_ang, ax=f_ax2, shrink=0.6)
        cbar = fig.colorbar(ax_img[idx], ax=f_img[idx],  # ticks=[-np.pi*.99, 0, np.pi*.99],
                            orientation='vertical', shrink=.8, pad=0.1)
        cbar.ax.set_xlabel("", rotation=0)

        f_hist.hist(image_dB(array.flatten()),color=palette[idx], bins='auto', histtype='step', label=legend, density=True) # , linestyle=(':')
        f_hist.legend()
        # ax = (sns.kdeplot((array.flatten()), shade=False, ax=f_cc, legend=True, label=title))
    f_hist.set_xlabel("$dB$")
    f_hist.set_ylabel("Normalized Frequency")
    fig.suptitle(suptitle)
    # now = datetime.datetime.now()
    # fig.savefig("eigenvalues_{}_{}.tiff".format(title, now.strftime("%Y-%m-%d-%H:%M:%S")))
    return fig

from matplotlib.ticker import PercentFormatter

def plot_image_hist(array_dict, suptitle, aspect=None, origin="lower", *args, **kwargs):
    colors = ["#003FFF", "#E8000B", "#03ED3A",  "#1A1A1A" ] # Blue, Red, Green, Black
    palette = sns.color_palette(colors, 4)

    n_images = len(array_dict.keys())
    fig = plt.figure(constrained_layout=False, figsize=(9 * n_images, 10))
    gs = fig.add_gridspec(nrows=2, ncols=n_images + 1)

    for idx, key in enumerate(array_dict):
        title = key
        array = array_dict[key]
        ax_im = fig.add_subplot(gs[0, idx])
        if (np.max(array) - np.min(array)) < 1.1:
            ax_cc = ax_im.imshow(array,  cmap='jet',
                                vmin=0.0, vmax=1.)
        elif (np.max(array) - np.min(array)) < 100 and (np.max(array) - np.min(array)) > 10:
            ax_cc = ax_im.imshow(array, cmap='jet',
                                vmin=0.0, vmax=90.)
        else:
            ax_cc = ax_im.imshow(array, cmap='jet',
                                vmin=np.min(array), vmax=np.max(array))

        ax_im.axis('off')
        ax_im.set_title(title)
        # ax_im.clim(0, 10)
        ax_cc.set_clim(ax_cc.get_clim()[0],ax_cc.get_clim()[1])

        # fig.colorbar(ax_ang, ax=f_ax2, shrink=0.6)
        cbar = fig.colorbar(ax_cc, ax=ax_im,  # ticks=[-np.pi*.99, 0, np.pi*.99],
                            orientation='vertical', shrink=1.0, pad=0.05)
        cbar.ax.set_xlabel("", rotation=0)

        ax_hist = fig.add_subplot(gs[1, idx])
        ax_hist.hist((array.ravel()), color=palette[idx], bins=1000, histtype='step', label=title, density=False, weights=100*np.ones(len(array.ravel())) / len(array.ravel())) # , linestyle=(':')
        ax_hist.set_xlim(ax_cc.get_clim())
        # Now we format the y-axis to display percentage
        ax_hist.yaxis.set_major_formatter(PercentFormatter(xmax=1))
        fig.suptitle(suptitle)

        # ax = (sns.kdeplot((array.flatten()), shade=False, ax=f_cc, legend=True, label=title))
        
    return fig