import plotstyle_plt
import glob
import numpy as np
import matplotlib.pyplot as plt
from itertools import product
from thunderhopper.modeltools import load_data
from color_functions import load_colors
from plot_functions import hide_axis, letter_subplots, ylimits, ylabel, super_xlabel, plot_line
from IPython import embed

def strip_zeros(num, right_digits=5):
    if isinstance(num, int):
        return num
    num = f'{num:.{right_digits}f}'
    left, right = num.split('.')
    right = right.rstrip('0')
    if right:
        return f'{left}.{right}'
    return left

def plot_snippets(axes, time, snippets, label, scales=None,
                  ymin=None, ymax=None, **kwargs):
    ymin, ymax = ylimits(snippets, minval=ymin, maxval=ymax, pad=0.05)
    ylabel(axes[0], label, x=-0.7, rotation=0, ha='left', va='center')
    for i, (ax, snippet) in enumerate(zip(axes, snippets.T)):
        plot_line(ax, time, snippet, ymin=ymin, ymax=ymax, **kwargs)
        if scales is not None:
            ax.set_title(f'$\\alpha={strip_zeros(scales[i])}$')
    return None

# GENERAL SETTINGS:
target = 'Omocestus_rufipes'
data_paths = glob.glob(f'../data/inv/log_hp/{target}*.npz')
stages = ['env', 'log', 'inv']
files = stages + ['scales', 'plot_scales', 'snr_env', 'snr_log', 'snr_inv']
save_path = '../figures/fig_invariance_log_hp.pdf'

# PLOT SETTINGS:
fig_kwargs = dict(
    figsize=(32/2.54, 16/2.54),
)
grid1_kwargs = dict(
    nrows=len(stages),
    ncols=None,
    wspace=0.05,
    hspace=0.3,
    left=0.1,
    right=0.985,
    bottom=0.17,
    top=0.9
)
grid2_kwargs = dict(
    nrows=1,
    ncols=3,
    wspace=0.35,
    hspace=0,
    left=0.1,
    right=0.985,
    bottom=0.18,
    top=0.95
)
colors = load_colors('../data/stage_colors.npz')
lw_snippets = 0.25
lw_analysis = 3
ylabels = dict(
    env=r'$x_{\text{env}}$',
    log=r'$x_{\text{dB}}$',
    inv=r'$x_{\text{adapt}}$',
    abs='abs. SNR',#'abs. ' + r'$\text{SNR}$',
    norm='norm. SNR',#'norm. ' + r'$\text{SNR}$',
    gain='rel. SNR gain'
)
xloc = dict(
    abs=5,
)
yloc = dict(
    env=100,
    log=10,
    inv=5,
    abs=50,
)
letter_kwargs = dict(
    x=0.03,
    y=0.99,
    ha='left',
    va='top'
)
fit_kwargs = dict(
    c='darkred',
    ls='--',
    zorder=1.9
)
fit_lw = dict(
    abs=6,
    norm=3,
    gain=3
)
text_kwargs = dict(
    abs=dict(s='$\\alpha^2+1$', fontsize=16, c=fit_kwargs['c'],
             x=0.85, y=0.9, ha='right', va='center'),
    norm=dict(s='$\\alpha^2+1$', fontsize=16, c=fit_kwargs['c'],
              x=0.85, y=0.9, ha='right', va='center'),
    gain=dict(s='$\\frac{1}{\\alpha}$', fontsize=24, c=fit_kwargs['c'],
              x=0.75, y=0.8, ha='left', va='center')
)
calculated_floor = False
floor_kwargs = dict(
    xmin=0,
    color=3 * (0.85,),
    zorder=0,
    lw=0
)

# EXECUTION:
for data_path in data_paths:
    print(f'Processing {data_path}')

    # Load invariance data:
    data, config = load_data(data_path, files)
    t_full = np.arange(data['env'].shape[0]) / config['env_rate']
    nonzero_scales = data['scales'][data['scales'] > 0]
    floor_kwargs['xmax'] = data['floor_scale'] if calculated_floor else 1

    # Prepare overall graph:
    fig = plt.figure(**fig_kwargs)
    super_grid = fig.add_gridspec(2, 1, wspace=0, hspace=1,
                                  left=0, right=1, bottom=0, top=1)

    # Prepare snippet axes:
    subfig1 = fig.add_subfigure(super_grid[0, :])
    grid1_kwargs['ncols'] = data['plot_scales'].size
    grid1 = subfig1.add_gridspec(**grid1_kwargs)
    axes = np.zeros((grid1.nrows, grid1.ncols), dtype=object)
    for i, j in product(range(grid1.nrows), range(grid1.ncols)):
        axes[i, j] = subfig1.add_subplot(grid1[i, j])
    [hide_axis(ax, 'bottom') for ax in axes[:-1, :].flatten()]
    [hide_axis(ax, 'left') for ax in axes[:, 1:].flatten()]
    super_xlabel('time [s]', subfig1, axes[-1, 0], axes[-1, -1], y=0)
    letter_subplots(axes[:, 0], labels='abc', **letter_kwargs)

    # Prepare analysis axes:
    subfig2 = fig.add_subfigure(super_grid[1, :])
    grid2 = subfig2.add_gridspec(**grid2_kwargs)
    symlog_kwargs = dict(linthresh=nonzero_scales.min(), linscale=0.2)

    ax_abs_snr = subfig2.add_subplot(grid2[:, 0])
    ax_abs_snr.set_ylabel(ylabels['abs'])

    ax_norm_snr = subfig2.add_subplot(grid2[:, 1])
    ax_norm_snr.set_ylabel(ylabels['norm'])
    ax_norm_snr.set_xscale('symlog', **symlog_kwargs)
    ax_norm_snr.set_yscale('log')

    ax_gain = subfig2.add_subplot(grid2[:, 2])
    ax_gain.set_ylabel(ylabels['gain'])
    ax_gain.set_xscale('symlog', **symlog_kwargs)
    ax_gain.set_yscale('log')
    super_xlabel('song scale $\\alpha$', subfig2, ax_abs_snr, ax_gain)
    letter_subplots([ax_abs_snr, ax_norm_snr, ax_gain], labels='def', **letter_kwargs)

    # Plot envelope snippets:
    plot_snippets(axes[0, :], t_full, data['env'], ylabels['env'], ymin=0,
                  scales=data['plot_scales'], c=colors['env'], lw=lw_snippets)
    axes[0, 0].yaxis.set_major_locator(plt.MultipleLocator(yloc['env']))
    
    # Plot logarithmic snippets:
    plot_snippets(axes[1, :], t_full, data['log'], ylabels['log'], ymax=0,
                  c=colors['log'], lw=lw_snippets)
    axes[1, 0].yaxis.set_major_locator(plt.MultipleLocator(yloc['log']))
    
    # Plot invariant snippets:
    plot_snippets(axes[2, :], t_full, data['inv'], ylabels['inv'],
                  c=colors['inv'], lw=lw_snippets)
    axes[2, 0].yaxis.set_major_locator(plt.MultipleLocator(yloc['inv']))

    # Plot in-representation SNRs (absolute):
    ax_abs_snr.plot(data['scales'], data['snr_env'], c=colors['env'], lw=lw_analysis)
    ax_abs_snr.plot(data['scales'], data['snr_log'], c=colors['log'], lw=lw_analysis)
    ax_abs_snr.plot(data['scales'], data['snr_inv'], c=colors['inv'], lw=lw_analysis)
    ax_abs_snr.axvspan(**floor_kwargs)
    ax_abs_snr.xaxis.set_major_locator(plt.MultipleLocator(xloc['abs']))
    ax_abs_snr.yaxis.set_major_locator(plt.MultipleLocator(yloc['abs']))
    ax_abs_snr.set_ylim(0, data['snr_env'].max())

    # Plot envelope SNR fit:
    ax_abs_snr.plot(data['scales'], data['scales'] ** 2 + 1,
                    lw=fit_lw['abs'], **fit_kwargs)
    ax_abs_snr.text(transform=ax_abs_snr.transAxes, **text_kwargs['abs'])

    # Get normalized SNRs:
    norm_snr_env = data['snr_env'] / data['snr_env'].max()
    norm_snr_log = data['snr_log'] / data['snr_log'].max()
    norm_snr_inv = data['snr_inv'] / data['snr_inv'].max()

    # Plot in-representation SNRs (normalized):
    ax_norm_snr.plot(data['scales'], norm_snr_env, c=colors['env'], lw=lw_analysis)
    ax_norm_snr.plot(data['scales'], norm_snr_log, c=colors['log'], lw=lw_analysis)
    ax_norm_snr.plot(data['scales'], norm_snr_inv, c=colors['inv'], lw=lw_analysis)
    ax_norm_snr.axvspan(**floor_kwargs)
    ax_norm_snr.set_ylim(norm_snr_env.min(), 1)

    # # Plot envelope SNR fit:
    # ax_norm_snr.plot(nonzero_scales, nonzero_scales / nonzero_scales.max(),
    #                  lw=fit_lw['norm'], **fit_kwargs)
    # ax_norm_snr.text(transform=ax_norm_snr.transAxes, **text_kwargs['norm'])

    # Get relative SNR gain:
    gain_log = norm_snr_log / norm_snr_env
    gain_inv = norm_snr_inv / norm_snr_env

    # Plot across-representation gain:
    ax_gain.plot(data['scales'], gain_log, c=colors['log'], lw=lw_analysis)
    ax_gain.plot(data['scales'], gain_inv, c=colors['inv'], lw=lw_analysis)
    ax_gain.axvspan(**floor_kwargs)
    ax_gain.set_ylim(1, 10)

    # Plot amplification fit:
    ax_gain.plot(nonzero_scales, nonzero_scales.max() / nonzero_scales,
                 lw=fit_lw['gain'], **fit_kwargs)
    ax_gain.text(transform=ax_gain.transAxes, **text_kwargs['gain'])
    
    if save_path is not None:
        fig.savefig(save_path)
    plt.show()

print('Done.')
embed()