# -*- coding: utf-8 -*-
"""
Created on Tue Aug 18 13:01:00 2020

@author: SHEN WANXIANG
"""



import warnings
warnings.filterwarnings("ignore")
import matplotlib.pyplot as plt
import seaborn as sns





def imshow(x_arr,  ax, mode = 'dark',  color_list = ['#1300ff','#ff0c00','#25ff00', '#d000ff','#e2ff00', 
                                                     '#00fff6', '#ff8800', '#fccde5','#178b66', '#8a0075'],
           x_max = 255, vmin = -1, vmax = 1,):
    
    
    assert x_arr.ndim == 3, 'input must be 3d array!'
    w, h, c = x_arr.shape
    assert len(color_list) >= c, 'length of the color list should equal or larger than channel numbers'
    
    x = x_arr.copy()
    x[x == 0] = 'nan'

    xxx = x_arr.sum(axis=-1)
    xxx[xxx != 0] = 'nan'

    if mode == 'dark':
        cmaps = [sns.dark_palette(color, n_colors =  50, reverse=False) for color in color_list]

    else:
        cmaps = [sns.light_palette(color, n_colors =  50, reverse=False) for color in color_list]
        
    for i in range(c):
        data = x[:,:,i]/x_max
        sns.heatmap(data, cmap = cmaps[i],  vmin = vmin, vmax = vmax,  
                    yticklabels=False, xticklabels=False, cbar=False, ax=ax, ) # linewidths=0.005, linecolor = '0.9'

    if mode == 'dark':
        sns.heatmap(xxx, vmin=-10000, vmax=1, cmap = 'Greys', yticklabels=False, xticklabels=False, cbar=False, ax=ax)
    else:
        sns.heatmap(xxx, vmin=0, vmax=1, cmap = 'Greys', yticklabels=False, xticklabels=False, cbar=False, ax=ax)
        ax.axhline(y=0, color='grey',lw=2, ls =  '--')
        ax.axhline(y=data.shape[0], color='grey',lw=2, ls =  '--')
        ax.autoscale()
        ax.axvline(x=data.shape[1], color='grey',lw=2, ls =  '--')  
        ax.axvline(x=0, color='grey',lw=2, ls =  '--')


def imshow_wrap(x,  mode = 'dark', color_list = ['#1300ff','#ff0c00','#25ff00', '#d000ff','#e2ff00', 
                                                 '#00fff6', '#ff8800', '#fccde5','#178b66', '#8a0075'], 
                x_max = 255, vmin = -1, vmax = 1,):
    
    fig, ax = plt.subplots(figsize=(4,4))
    imshow(x.astype(float),mode = mode, color_list = color_list, ax=ax, x_max = x_max, vmin = vmin, vmax=vmax)