import os
# from tkinter import N
import matplotlib.pyplot as plt
import matplotlib
matplotlib.rc('ytick', labelsize=20)
matplotlib.rc('xtick', labelsize=20)
matplotlib.rcParams['lines.linewidth'] = 4
matplotlib.rcParams['axes.spines.right'] = False
matplotlib.rcParams['axes.spines.top'] = False
matplotlib.rcParams['axes.linewidth'] = 2
matplotlib.rcParams['xtick.major.width'] = 2
matplotlib.rcParams['xtick.minor.width'] = 2
matplotlib.rcParams['ytick.major.width'] = 2
matplotlib.rcParams['ytick.minor.width'] = 2
matplotlib.rcParams["font.family"] = "DeJavu Serif"
matplotlib.rcParams["font.serif"] = ["Times New Roman"]
import seaborn as sns
import numpy as np
from anal.util import transpose_np
from anal.pmap_th import get_pmaps_theory, get_pmaps_theory_grid
from plot.util import sigma_pcolor, get_pal, get_color_gradient as get_c
from plot.const import BLUE, RED, DARKBLUE, DARKRED, GREEN, \
                       RUN, SIGMA_W, SIGMA_B
from tqdm import tqdm


# line plot, fig 2.b,c, fig 7.a
def line_plots(args, ws, bs, ets, lenT1, lenT2=None, lenthm=None):
    """
    (legend, x)
    stat
        zd lst (안에 loss)
        theory
        wb lst (안에 loss)
    다양한 len에서는 loss 제외
    nw, nb, T, L
    nw, nb, E, L-1
    """
    stat_name = ['z', 'zh', 'h', 'd', 'loss_inf', 'dz_ratio', 'w_grad', 'b_grad', 'loss_tr']
    stat_dt={'z':0, 'd':1, 'w_grad':2, 'b_grad':3}
    thm_name = ['p', 'pt', 'ph', 'q']
    Ts = [10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000, 50000, 100000, 200000]
    T = args.T
    idx_T = Ts.index(T)
    Ts = Ts[:idx_T+1]
    ls = np.arange(2,args.n_layers+1)
    early_l = args.n_layers//6
    mid_l = args.n_layers//2
    late_l = (args.n_layers//6)*5
    short_ls = [i for i in range(args.n_layers-1)] # [early_l, mid_l, late_l]
    # zd_stat_lst, wb_stat_lst, wb_inf_stat_lst = lenT1.get_stat()
    # # 44-45 spits error
    # zd = transpose_np([(zd_stat_lst[0][0][:,:,:,:,1:],
                        # zd_stat_lst[0][1][:,:,:,:,1:])]
                       # +zd_stat_lst[1:-2]) # excluding loss_inf and qpratio
    # wb = transpose_np(wb_stat_lst[:-1])
    # wb_inf = transpose_np(wb_inf_stat_lst)
    zd_stat_lst, wb_stat_lst, wb_inf_stat_lst, zd, wb, wb_inf = lenT1.get_stat_plot()
    pal_lst = [get_c(BLUE, RED, len(lst)) for lst in \
               [zd_stat_lst[:-1], wb_stat_lst[:-1], ws, bs, ets, Ts, ls]]
    pal_zd, pal_wb, pal_ws, pal_bs, pal_ets, pal_Ts, pal_ls = pal_lst
    plot_dir = 'line_plots_'
    # shape: n_sws, n_sbs, n_ets, n_iters, n_ls(-1)
    if lenT2 is not None:
        zd_stat_lst2, wb_stat_lst2, wb_inf_stat_lst2, zd2, wb2, wb_inf2 = lenT2.get_stat_plot()
        stat_lst2 = zd_stat_lst2 + wb_inf_stat_lst2
    for i_stat, stat in enumerate(tqdm(zd_stat_lst + wb_inf_stat_lst)):
        pre_name = f'{stat_name[i_stat]}/{stat_name[i_stat]}'
        with_thm = stat_name[i_stat] in ['z', 'd', 'w_grad', 'b_grad'] \
                   and lenthm is not None
        ## note. 더 나은 방식:
        ## for loop을 먼저 중첩시키고, if == 0 조건을 그리고자하는 변수에 대해서 걸어줌
        # for et, T, w, b, l
        for i_et, et in enumerate(ets):
            for i_T, cur_T in enumerate(Ts):
                for i_w, w in enumerate(ws): # (sigma_b, layer)
                    plot_name = pre_name + f'_sw:{w}_et:{et}_T:{cur_T}'
                    ## T exponential
                    ms, ss = stat[0][i_w,:,i_et,cur_T-1], stat[1][i_w,:,i_et,cur_T-1]
                    ps = lenthm[stat_dt[stat_name[i_stat]]][i_w,:,i_et,cur_T-1] if with_thm else None
                    plot_xs(args, plot_dir+'lyr_b', plot_name, bs, pal_bs, ms, ss, ps)
                    for i_b, b in enumerate(bs):
                        if i_T == 0: # (layer, iteration)
                            ls = ls-1 if i_stat == 0 else ls
                            plot_name = pre_name + f'_sw:{w}_sb:{b}_et:{et}'
                            ms = stat[0][i_w,i_b,i_et].T; ss = stat[1][i_w,i_b,i_et].T
                            ps = lenthm[stat_dt[stat_name[i_stat]]][i_w,i_b,i_et].T if with_thm else None
                            plot_xs(args, plot_dir+'iter_l', plot_name, ls, pal_ls, ms, ss, ps)
                            ls = ls+1 if i_stat == 0 else ls
                        if i_w == 0: # (sigma_w, layer)
                            plot_name = pre_name + f'_sb:{b}_et:{et}_T:{cur_T}'
                            ms = stat[0][:,i_b,i_et,cur_T-1]; ss = stat[1][:,i_b,i_et,cur_T-1]
                            ps = (lenthm[stat_dt[stat_name[i_stat]]][:,i_b,i_et,cur_T-1] if with_thm else None)
                            ms2 = None; ss2 = None
                            if lenT2 is not None:
                                stat2 = stat_lst2[i_stat]
                                ms2 = stat2[0][:,i_b,i_et,cur_T-1]; ss2 = stat2[1][:,i_b,i_et,cur_T-1]
                            plot_xs(args, plot_dir+'lyr_w', plot_name, ws, pal_ws, ms, ss, ps, ms2, ss2)
                        for i_l, l in enumerate(short_ls):
                            lb = 1 if stat_name[i_stat] == 'z' else 2
                            if i_stat == 0: # (다양한 zd wb lens, iteration)
                                if i_T == 0:
                                    plot_name = f'whole/whole(zd)_sw:{w}_sb:{b}_et:{et}_l:{l+lb}'
                                    plot_xs(args, plot_dir+'iter_a', plot_name, stat_name[:4],
                                        pal_zd, zd[0][:,i_w,i_b,i_et,:,l], zd[1][:,i_w,i_b,i_et,:,l])
                                    plot_name = f'whole/whole(wb)_sw:{w}_sb:{b}_et:{et}_l:{l+lb}'
                                    plot_xs(args, plot_dir+'iter_a', plot_name, stat_name[6:-1],
                                        pal_wb, wb_inf[0][:,i_w,i_b,i_et,:,l], wb_inf[1][:,i_w,i_b,i_et,:,l])
                                if i_l == 0: # (다양한 zd wb lens, layer, last iter)
                                    plot_name = f'whole/whole(zd)_sw:{w}_sb:{b}_et:{et}_T:{cur_T}'
                                    ms = zd[0][:,i_w,i_b,i_et,cur_T-1,:] # zd 생성시 tr로 첫 :가 z, zh, h등을 가리킴
                                    ss = zd[1][:,i_w,i_b,i_et,cur_T-1,:]
                                    plot_xs(args, plot_dir+'lyr_a', plot_name, stat_name[:4],
                                                pal_zd, ms, ss)
                                    plot_name = f'whole/whole(wb)_sw:{w}_sb:{b}_et:{et}_T:{cur_T}'
                                    ms = wb[0][:,i_w,i_b,i_et,-1,:]
                                    ss = wb[1][:,i_w,i_b,i_et,-1,:]
                                    plot_xs(args, plot_dir+'lyr_a', plot_name, stat_name[6:-1],
                                                pal_wb, ms, ss)
                                if i_w == 0: # (다양한 zd wb lens, sigma_w, last iter)
                                    plot_name = f'whole/whole(zd)_sb:{b}_et:{et}_l:{l+lb}_T:{cur_T}'
                                    ms = zd[0][:,:,i_b,i_et,cur_T-1,l]
                                    ss = zd[1][:,:,i_b,i_et,cur_T-1,l]
                                    plot_xs(args, plot_dir+'sw_a', plot_name, stat_name[:4],
                                            pal_zd, ms, ss, is_line=False, xl=ws)
                                    plot_name = f'whole/whole(wb)_sb:{b}_et:{et}_l:{l+lb}_T:{cur_T}'
                                    ms = wb[0][:,:,i_b,i_et,-1,l]
                                    ss = wb[1][:,:,i_b,i_et,-1,l]
                                    plot_xs(args, plot_dir+'sw_a', plot_name, stat_name[6:-1],
                                            pal_zd, ms, ss, is_line=False, xl=ws)
                                if i_b == 0: # (다양한 zd wb lens, sigma_b, last iter)
                                    plot_name = f'whole/whole(zd)_sw:{w}_et:{et}_l:{l+lb}_T:{cur_T}'
                                    ms = zd[0][:,i_w,:,i_et,cur_T-1,l]
                                    ss = zd[1][:,i_w,:,i_et,cur_T-1,l]
                                    plot_xs(args, plot_dir+'sb_a', plot_name, stat_name[:4],
                                            pal_zd, ms, ss, is_line=False, xl=bs)
                                    plot_name = f'whole/whole(wb)_sw:{w}_et:{et}_l:{l+lb}_T:{cur_T}'
                                    ms = wb[0][:,i_w,:,i_et,-1,l]
                                    ss = wb[1][:,i_w,:,i_et,-1,l]
                                    plot_xs(args, plot_dir+'sb_a', plot_name, stat_name[6:-1],
                                            pal_zd, ms, ss, is_line=False, xl=bs)
                            # l = l-1 if i_stat == 0 else l
                            if i_w == 0:
                                plot_name = pre_name + f'_sb:{b}_et:{et}_T:{cur_T}_l:{l+lb}'
                                ms = stat[0][:,i_b,i_et,:cur_T,l]; ss = stat[1][:,i_b,i_et,:cur_T,l]
                                ps = lenthm[stat_dt[stat_name[i_stat]]][:,i_b,i_et,:cur_T,l] if with_thm else None
                                plot_xs(args, plot_dir+'iter_w', plot_name, ws, pal_ws,
                                        ms, ss, ps)
                                if i_et == 0: # (eta, sigma_w) : for b, 3l, T
                                    plot_name = pre_name + f'_sb:{b}_l:{l+lb}_T:{cur_T}'
                                    ms = stat[0][:,i_b,:,cur_T-1,l].T; ss = stat[1][:,i_b,:,cur_T-1,l].T
                                    ps = lenthm[stat_dt[stat_name[i_stat]]][:,i_b,:,cur_T-1,l].T if with_thm else None
                                    plot_xs(args, plot_dir+'sw_eta', plot_name, ets, pal_ets,
                                            ms, ss, ps, is_line=False, xl=ws)
                            if i_b == 0:
                                if i_T == 0: # (sigma_b, iteration)
                                    plot_name = pre_name + f'_sw:{w}_et:{et}_l:{l+lb}'
                                    ms = stat[0][i_w,:,i_et,:,l]; ss = stat[1][i_w,:,i_et,:,l]
                                    ps = lenthm[stat_dt[stat_name[i_stat]]][i_w,:,i_et,:,l] if with_thm else None
                                    plot_xs(args, plot_dir+'iter_b', plot_name, bs, pal_bs,
                                            ms, ss, ps)
                                if i_et == 0: # (eta, sigma_b) for w, 3l, T
                                    plot_name = pre_name + f'_sw:{w}_l:{l+lb}_T:{cur_T}'
                                    ms = stat[0][i_w,:,:,cur_T-1,l].T; ss = stat[1][i_w,:,:,cur_T-1,l].T
                                    ps = lenthm[stat_dt[stat_name[i_stat]]][i_w,:,:,cur_T-1,l].T if with_thm else None
                                    plot_xs(args, plot_dir+'sb_eta', plot_name, ets, pal_ets,
                                            ms, ss, ps, is_line=False, xl=bs)
                            if i_et == 0:
                                if i_l == 0: # (eta, l) for w, b, T
                                    plot_name = pre_name + f'_sw:{w}_sb:{b}_T:{cur_T}'
                                    ms = stat[0][i_w,i_b,:,cur_T-1]; ss = stat[1][i_w,i_b,:,cur_T-1]
                                    ps = lenthm[stat_dt[stat_name[i_stat]]][i_w,i_b,:,cur_T-1] if with_thm else None
                                    plot_xs(args, plot_dir+'lyr_eta', plot_name, ets, pal_ets,
                                            ms, ss, ps)
                                if i_T == 0: # (eta, iter) for w, b, 3l
                                    plot_name = pre_name + f'_sw:{w}_sb:{b}_l:{l+lb}'
                                    ms = stat[0][i_w,i_b,:,:,l]; ss = stat[1][i_w,i_b,:,:,l]
                                    ps = lenthm[stat_dt[stat_name[i_stat]]][i_w,i_b,:,:,l] if with_thm else None
                                    plot_xs(args, plot_dir+'iter_eta', plot_name, ets, pal_ets,
                                            ms, ss, ps)
                            if i_l == 0 and i_T == 0: #(T, l) for w, b, eta
                                Tpt = [T-1 for T in Ts]
                                plot_name = pre_name + f'_sw:{w}_sb:{b}_et:{et}'
                                ms = stat[0][i_w,i_b,i_et,Tpt,:]; ss = stat[1][i_w,i_b,i_et,Tpt,:]
                                plot_xs(args, plot_dir+'lyr_iter', plot_name, Ts, pal_Ts,
                                        ms, ss, None)


def lineplot(plot_dir, legends, x, ms, ss, pal, marker='-'):
    for i, (m, s, c) in enumerate(zip(ms, ss, pal)):
        if np.any(m==0):
            continue
        plt.plot(x, m.T, marker, markersize=3, color=c, alpha=1)
        label=f'${plot_dir[-1]}=$'.replace('$a=$','').replace('r=','T=')\
            .replace('w=','\sigma_w=').replace('b=','\sigma_b=')
        label = label + f'${legends[i]}$'.replace('_','\_')
        if ss is not None:
            plt.fill_between(x, m-s, m+s, alpha=0.2, color=c,
                label=label)


def lineplot_ps(x, ps, pal, marker='o'):
    for i, (p, c) in enumerate(zip(ps, pal)):
        plt.plot(x, p.T, 'o', markersize=5, color=c, alpha=0.8)


def plot_xs(args, plot_dir, plot_name, legends, pal, ms, ss, ps=None,
            ms2=None, ss2=None, is_line=True, xl=None):
    """
    ms: n_m, n_iteration (e.g. n_layer, etc.)
    if legend is layer: figsize = (10, 6)
    if plotting dz_ratio: y-axis is percentage
    if x-axis is sigma_w or sigma_b: bar plot
    """
    if plot_dir[-1] == 'l':
        plt.figure(figsize=(10, 6))
    else: plt.figure(figsize=(5, 3))
    if 'dz_ratio' in plot_name:
        plt.ylim(0, 100)
        ms *= 100; ss *= 100
    else: plt.yscale('log')
    x = np.arange(1, ms.shape[-1]+1)

    def barplot(ps):
        n_bars=len(ms); width=0.8/n_bars
        x_offset = x - width * (n_bars - 1) / 2
        for i, (m, s, c) in enumerate(zip(ms, ss, pal)):
            if np.any(m==0):
                continue
            plt.bar(x_offset, m.T, width=width, color=c)
            label=f'${plot_dir[-1]}=$'.replace('$a=$','').replace('r=','T=')\
                .replace('w=','\sigma_w=').replace('b=','\sigma_b=').replace('et=', '\eta=')
            label = label + f'${legends[i]}$'.replace('_','\_')
            if ss is not None:
                plt.errorbar(x_offset, m, s, fmt='o', alpha=0.5, color=c,
                    label=label)
            x_offset += width
        if ps is not None:
            if ps.shape[-1] != ms.shape[-1]: ps = ps[:,1:]
            for i, (p, c) in enumerate(zip(ps, pal)):
                plt.bar(x_offset, p.T, width=width, color=c, alpha=0.6)
                x_offset += width
        plt.xticks(x, xl)

    lineplot(plot_dir, legends, x, ms, ss, pal, marker='-')
    if ms2 is not None:
        legends2 = [lgnd + '_reg' for lgnd in legends]
        lineplot(plot_dir, legends2, x, ms2, ss2, pal, marker='-')
    if ps is not None:
        if ps.shape[-1] != ms.shape[-1]:
            ps = ps[:,1:]
        lineplot_ps(x, ps, pal) if is_line else barplot(ps)
    if 'dz_ratio' in plot_name:
        ms /= 100; ss /= 100
    plt.tick_params(axis='x', labelsize=20)
    plt.tick_params(axis='y', labelsize=20)
    # plt.tick_params(top='off', right='off')
    if is_line: plt.xlim(1, ms.shape[-1])
    else: plt.xlim(0, ms.shape[-1]+1)
    # optionally label x-axis, y-axis, title, legend
    if args.label:
        plt.xlabel(plot_dir.split('_')[-2], fontsize=20)
        prefix = 'len of ' if 'loss' not in plot_name else ''
        yl = prefix + plot_name.split('/')[1].split('_')[0] \
            if not 'dz_ratio' in plot_name else 'd/z percentage (%)'
        plt.ylabel(yl, fontsize=20)
        title = plot_name.split('/')[1].split('_')[-3:] if 'whole' in plot_name \
            else plot_name.split('/')[1].split('_')[-2:]
        title = ' '.join(title).replace('sw:', '$\sigma_w$:')\
                               .replace('sb:', '$\sigma_b$:')\
                               .replace('et:', '$\eta$:')
        plt.title(title, fontsize=20)
        if len(legends) < 7: plt.legend()
        else: plt.legend(bbox_to_anchor=(1.1, 1))
    # plt.tight_layout()
    p = f'{args.fig_dir}/{plot_dir}/'
    if not os.path.exists(p+plot_name.split('/')[0]):
        os.makedirs(p+plot_name.split('/')[0])
    plt.savefig(os.path.join(p, f'{plot_name}.png'), bbox_inches='tight')
    plt.cla()
    plt.clf()
    plt.close()



def old_plot_xs(args, plot_dir, plot_name, legends, pal, ms, ss, ps=None,
            ms2=None, ss2=None, is_line=True, xl=None):
    """
    ms: n_m, n_iteration (e.g. n_layer, etc.)
    if legend is layer: figsize = (10, 6)
    if plotting dz_ratio: y-axis is percentage
    if x-axis is sigma_w or sigma_b: bar plot
    """
    if plot_dir[-1] == 'l':
        plt.figure(figsize=(10, 6))
    else: plt.figure(figsize=(5, 3))
    if 'dz_ratio' in plot_name:
        plt.ylim(0, 100)
        ms *= 100; ss *= 100
    else: plt.yscale('log')
    x = np.arange(1, ms.shape[-1]+1)

    def lineplot(ps):
        for i, (m, s, c) in enumerate(zip(ms, ss, pal)):
            if np.any(m==0):
                continue
            plt.plot(x, m.T, '-', markersize=3, color=c, alpha=1)
            label=f'${plot_dir[-1]}=$'.replace('$a=$','')\
                .replace('w=','\sigma_w=').replace('b=','\sigma_b=')
            label = label + f'${legends[i]}$'.replace('_','\_')
            if ss is not None:
                plt.fill_between(x, m-s, m+s, alpha=0.2, color=c,
                    label=label)
        if ps is not None:
            if ps.shape[-1] != ms.shape[-1]: ps = ps[:,1:]
            for i, (p, c) in enumerate(zip(ps, pal)):
                plt.plot(x, p.T, 'o', markersize=5, color=c, alpha=0.8)
        # plt.xticks(x)

    def barplot(ps):
        n_bars=len(ms); width=0.8/n_bars
        x_offset = x - width * (n_bars - 1) / 2
        for i, (m, s, c) in enumerate(zip(ms, ss, pal)):
            if np.any(m==0):
                continue
            plt.bar(x_offset, m.T, width=width, color=c)
            label=f'${plot_dir[-1]}=$'.replace('$a=$','')\
                .replace('w=','\sigma_w=').replace('b=','\sigma_b=').replace('et=', '\eta=')
            label = label + f'${legends[i]}$'.replace('_','\_')
            if ss is not None:
                plt.errorbar(x_offset, m, s, fmt='o', alpha=0.5, color=c,
                    label=label)
            x_offset += width
        if ps is not None:
            if ps.shape[-1] != ms.shape[-1]: ps = ps[:,1:]
            for i, (p, c) in enumerate(zip(ps, pal)):
                plt.bar(x_offset, p.T, width=width, color=c, alpha=0.6)
                x_offset += width
        plt.xticks(x, xl)

    lineplot(ps) if is_line else barplot(ps)
    if 'dz_ratio' in plot_name:
        ms /= 100; ss /= 100
    plt.tick_params(axis='x', labelsize=20)
    plt.tick_params(axis='y', labelsize=20)
    # plt.tick_params(top='off', right='off')
    if is_line: plt.xlim(1, ms.shape[-1])
    else: plt.xlim(0, ms.shape[-1]+1)
    # optionally label x-axis, y-axis, title, legend
    if args.label:
        plt.xlabel(plot_dir.split('_')[-2], fontsize=20)
        prefix = 'len of ' if 'loss' not in plot_name else ''
        yl = prefix + plot_name.split('/')[1].split('_')[0] \
            if not 'dz_ratio' in plot_name else 'd/z percentage (%)'
        plt.ylabel(yl, fontsize=20)
        title = plot_name.split('/')[1].split('_')[-3:] if 'whole' in plot_name \
            else plot_name.split('/')[1].split('_')[-2:]
        title = ' '.join(title).replace('sw:', '$\sigma_w$:')\
                               .replace('sb:', '$\sigma_b$:')\
                               .replace('et:', '$\eta$:')
        plt.title(title, fontsize=20)
        if len(legends) < 7: plt.legend()
        else: plt.legend(bbox_to_anchor=(1.1, 1))
    # plt.tight_layout()
    p = f'{args.fig_dir}/{plot_dir}/'
    if not os.path.exists(p+plot_name.split('/')[0]):
        os.makedirs(p+plot_name.split('/')[0])
    plt.savefig(os.path.join(p, f'{plot_name}.png'), bbox_inches='tight')
    plt.cla()
    plt.clf()
    plt.close()


# scatter, fig 2.a
def scatter_plots(lenlog, args, run_cond_vs):
    '''
    example of the shapes
    lenlog.z_len, lenlog.d_len: (n_runs, nw, nb, net, T, L, bsz) = (2, 3, 2, 100, 15, 64)
    p_iter: (nw, T, L, n_runs*bsz) = (3, 100, 15, 2*64)
    '''
    ps, qs = lenlog.z_len, lenlog.d_len # (n_runs, nw, nb, net, T, L, bsz) -> (nw, T, L, nb*net*n_runs*bsz)
    p_iter, q_iter = ps.transpose(1,4,5,2,3,0,6), qs.transpose(1,4,5,2,3,0,6)
    pconds = ['z_', 'd_']
    pal = get_c(BLUE, RED, args.n_conds)
    # ps_iter = [ps.transpose(0,2,3,1,4) for ps in ps_result]
    # p_iter, ph_iter, pt_iter, pc_iter, q_iter = ps_iter
    p_iter = np.reshape(p_iter, (p_iter.shape[0], p_iter.shape[1], p_iter.shape[2],
                                p_iter.shape[3]*p_iter.shape[4]*p_iter.shape[5]*p_iter.shape[6]))
    q_iter = np.reshape(q_iter, (q_iter.shape[0], q_iter.shape[1], q_iter.shape[2],
                                q_iter.shape[3]*q_iter.shape[4]*q_iter.shape[5]*q_iter.shape[6]))
    for l in range(1, args.n_layers-1):
        plt.figure(figsize=(5, 3))
        for i_w in range(args.n_conds):
            # plot_qmaps(qrange, qmaps[..., 1], widxs, bidxs, lw=2) # from theory
            plt.scatter([], [], label=f'$\sigma_w$={run_cond_vs[i_w]}', color=pal[i_w])
            for t in range(args.T-1):
                alpha = 0.003 + 0.047*(min((5**t) / (5**5),1))
                plt.scatter(p_iter[i_w, t, l, :], p_iter[i_w, t+1, l, :],
                    s=20, color=pal[i_w], alpha=alpha)
        plt.plot((1e-2, 1e+4), (1e-2, 1e+4), '--', color='k', zorder=900, linewidth=4)
        plt.xlim(1e-2, 1e+4)
        plt.ylim(1e-2, 1e+5)
        plt.tick_params(axis='x')
        plt.tick_params(axis='y')
        plt.xscale('log')
        plt.yscale('log')
        if args.label:
            plt.legend()
            plt.xlabel(f'${{length}}^t$', fontsize=20)
            plt.ylabel(f'${{length}}^{{t+1}}$', fontsize=20)
            plt.title(f'$\eta$: {args.eta} $\sigma_b$: {args.sigma_b} l:{l}', fontsize=20)
        # plt.tight_layout()
        if not os.path.exists(f'{args.fig_dir}/scatter_plots_len_w'):
            os.makedirs(f'{args.fig_dir}/scatter_plots_len_w')
        plt.savefig(os.path.join(f'{args.fig_dir}/scatter_plots_len_w',
                    f'{pconds[0]}l:{l+1}.png'), bbox_inches='tight')
        plt.cla()
        plt.clf()
        plt.close()
        plt.figure(figsize=(5, 3))
        for i_w in range(args.n_conds):
            # plot_qmaps(qrange, qmaps[..., 1], widxs, bidxs, lw=2) # from theory
            plt.scatter([], [], label=f'$\sigma_w$={run_cond_vs[i_w]}', color=pal[i_w])
            for t in range(args.T-1):
                alpha = 0.003 + 0.047*(min((5**t) / (5**5),1))
                plt.scatter(q_iter[i_w, t, l, :], q_iter[i_w, t+1, l, :],
                    s=20, color=pal[i_w], alpha=alpha)
        plt.plot((1e-2, 1e+4), (1e-2, 1e+4), '--', color='k', zorder=900, linewidth=4)
        plt.xlim(1e-2, 1e+4)
        plt.ylim(1e-2, 1e+5)
        plt.tick_params(axis='x')
        plt.tick_params(axis='y')
        plt.xscale('log')
        plt.yscale('log')
        if args.label:
            plt.legend()
            plt.xlabel(f'${{length}}^t$', fontsize=20)
            plt.ylabel(f'${{length}}^{{t+1}}$', fontsize=20)
            plt.title(f'$\eta$: {args.eta} $\sigma_b$: {args.sigma_b} l:{l}', fontsize=20)
        # plt.tight_layout()
        if not os.path.exists(f'{args.fig_dir}/scatter_plots_len_w'):
            os.makedirs(f'{args.fig_dir}/scatter_plots_len_w')
        plt.savefig(os.path.join(f'{args.fig_dir}/scatter_plots_len_w',
                    f'{pconds[1]}l:{l+1}.png'), bbox_inches='tight')
        plt.cla()
        plt.clf()
        plt.close()


# heatmap, fig 2.d
def heatmap_plots(lenlog, args, sigma_ws, sigma_bs):
    '''
    example of the shapes
    pmaps_grid: (n_ps, nw, nb, L)
    qmaps_grid: (nw, nb, L)
    T is deleted since got the last time step
    ps_result: list
        z, d: (nw, n_runs, nb, L, bsz)
        w: (nw, n_runs, nb, L)
    '''
    zd_stat, wb_stat, _ = lenlog.get_stat()
    ps_result = [zd_stat[0][0][:,:,-1], zd_stat[3][0][:,:,-1],
                 wb_stat[0][0][:,:,-1], wb_stat[1][0][:,:,-1]]
    # TODO: check dimensions and mean over runs and bsz
    pconds = ['z', 'd', 'w_grad', 'b_grad']
    # ps_result = [np.mean(p_result,(-1,-4)) if p_result.ndim == 5 else \
    #              np.mean(p_result,(-3)) for p_result in ps_result]
    n_ws = len(sigma_ws); n_bs = len(sigma_bs)
    max_sw = max(sigma_ws); max_sb = max(sigma_bs)
    xticks = np.array([0, 0.25, 0.5, 0.75, 1.0]) * max_sb
    yticks = np.array([0, 0.25, 0.5, 0.75, 1.0]) * max_sw
    xticks = np.round(xticks, 2); yticks = np.round(yticks, 2)
    if not args.skip_theory:
        pmaps_grid, qmaps_grid = get_pmaps_theory_grid(args, sigma_ws, sigma_bs)
    for i, p_result in enumerate(ps_result):
        pcond = pconds[i]
        for l in range(1, args.n_layers-1):
            plt.figure(figsize=(10, 6))
            if not args.skip_theory:
                sigma_pcolor(pmaps_grid[i,:,:,l], sigma_ws, sigma_bs)
            else:
                sigma_pcolor(p_result[:,:,l], sigma_ws, sigma_bs)
            if args.label:
                plt.title(f'{pcond}, $\eta$:{args.eta}, l:{l}', fontsize=40)
                plt.xlabel(f'$\sigma_b$', fontsize=40)
                plt.ylabel(f'$\sigma_w$', fontsize=40)
            # TODO: change ticks according to the range of sigmas
            plt.xticks(xticks, fontsize=40)
            plt.yticks(yticks, fontsize=40)
            plt.xlim(0, max_sb)
            plt.ylim(0, max_sw)
            # plt.tight_layout()
            if not os.path.exists(f'{args.fig_dir}/heatmaps_sbsw/{pcond}'):
                os.makedirs(f'{args.fig_dir}/heatmaps_sbsw/{pcond}')
            plt.savefig(os.path.join(f'{args.fig_dir}/heatmaps_sbsw/{pcond}',
                        f'{pcond}_l:{l+1}.png'), bbox_inches='tight')
            plt.cla()
            plt.clf()
            plt.close()
