import plotstyle_plt
import numpy as np
import matplotlib.pyplot as plt
from thunderhopper.filetools import search_files
from thunderhopper.modeltools import load_data
from thunderhopper.filtertools import find_kern_specs
from color_functions import load_colors, shade_colors
from plot_functions import hide_axis, ylimits, xlabel, ylabel, super_ylabel,\
                           plot_line, plot_barcode, strip_zeros, time_bar,\
                           letter_subplot, letter_subplots, hide_ticks,\
                           super_xlabel, super_ylabel, assign_colors
from IPython import embed

# GENERAL SETTINGS:
target = 'Omocestus_rufipes'
search_kwargs = dict(
    incl='subset',
    excl='noise',
    dir='../data/inv/thresh_lp/'
)
pure_paths = search_files(target, **search_kwargs)
load_kwargs = dict(
    keywords=['scales', 'measure', 'thresh']
)
save_path = None#'../figures/fig_invariance_thresh_lp_subset.pdf'

# GRAPH SETTINGS:
fig_kwargs = dict(
    figsize=(32/2.54, 16/2.54),
)
super_grid_kwargs = dict(
    nrows=1,
    ncols=1,
    wspace=0,
    hspace=0,
    left=0,
    right=1,
    bottom=0,
    top=1
)
grid_kwargs = dict(
    nrows=2,
    ncols=1,
    wspace=0,
    hspace=0.1,
    left=0.15,
    right=0.95,
    bottom=0.1,
    top=0.85
)
inset_bounds = [0.2, 1.01, 0.6, 0.4]

# PLOT SETTINGS:
colors = load_colors('../data/stage_colors.npz')
color_factors = [-0.5, 0.5]
lw = dict(
    one=3,
    kern=3,
    all=1,
)
ax_labels = dict(
    x='scale $\\alpha$',
    y='$\\mu_f$',
)
xlab_kwargs = dict(
    y=0.005,
    fontsize=16,
    ha='center',
    va='bottom',
)
ylab_kwargs = dict(
    x=0,
    fontsize=20,
    ha='left',
    va='center',
)
yloc = 0.2

# EXECUTION:
for pure_path in pure_paths:
    print(f'Processing {pure_path}')
    noise_path = pure_path.replace('.npz', '_noise.npz')

    # Load kernel invariance data:
    pure_data, config = load_data(pure_path, **load_kwargs)
    noise_data, _ = load_data(noise_path, **load_kwargs)
    scales = pure_data['scales']

    # Adjust grid parameters:
    n_columns = config['k_specs'].shape[0] + 1
    super_grid_kwargs['ncols'] = n_columns

    # Prepare overall graph:
    fig = plt.figure(**fig_kwargs)
    super_grid = fig.add_gridspec(**super_grid_kwargs)

    # Prepare axes:
    all_axes = np.zeros((grid_kwargs['nrows'], n_columns), dtype=object)
    subfigs = []
    for i in range(n_columns):
        subfig = fig.add_subfigure(super_grid[0, i])
        grid = subfig.add_gridspec(**grid_kwargs)
        subfigs.append(subfig)
        for j in range(grid_kwargs['nrows']):
            ax = subfig.add_subplot(grid[j, 0])
            ax.set_xlim(scales[0], scales[-1])
            ax.set_ylim(0, 1)
            ax.set_xscale('symlog', linthresh=scales[1], linscale=0.5)
            ax.yaxis.set_major_locator(plt.MultipleLocator(yloc))
            if i > 0:
                hide_ticks(ax, side='left')
            all_axes[j, i] = ax
        hide_ticks(all_axes[0, i], side='bottom')
    super_xlabel(ax_labels['x'], fig, all_axes[-1, 0], all_axes[-1, -1], **xlab_kwargs)
    super_ylabel(ax_labels['y'], fig, all_axes[0, 0], all_axes[1, 0], **ylab_kwargs)

    # Plot kernel-specific results:
    in_min, in_high = ylimits(config['kernels'], pad=0.05)
    for i in range(config['k_specs'].shape[0]):
        pure_ax, noise_ax = all_axes[:, i]
        # Plot results of pure-song analysis:
        pure_ax.plot(scales, pure_data['measure_feat'][:, i, :],
                     c=colors['feat'], lw=lw['one'])
        # Plot results of noise-song analysis:
        noise_ax.plot(scales, noise_data['measure_feat'][:, i, :],
                      c=colors['feat'], lw=lw['one'])
        # Indicate kernel waveform:
        inset = pure_ax.inset_axes(inset_bounds)
        inset.plot(config['k_times'], config['kernels'][:, i], c='k', lw=lw['kern'])
        inset.set_xlim(config['k_times'][0], config['k_times'][-1])
        inset.set_ylim(in_min, in_high)
        inset.axis('off')

    # Load population invariance data:
    pure_data, config = load_data(pure_path.replace('_subset', ''), **load_kwargs)
    noise_data, _ = load_data(noise_path.replace('_subset', ''), **load_kwargs)
    scales = pure_data['scales']

    # Get kernel type-specific colors:
    types, ind = np.unique(config['k_specs'][:, 0], return_index=True)
    types = types[np.argsort(ind)].astype(int)
    factors = np.linspace(*color_factors, types.size)
    kern_colors = shade_colors(colors['feat'], factors)
    kern_colors = dict(zip(types.astype(str), kern_colors))

    # Plot population-wide results:
    pure_ax, noise_ax = all_axes[:, -1]
    handles = pure_ax.plot(scales, pure_data['measure_feat'], c='k', lw=lw['all'])
    assign_colors(handles, config['k_specs'][:, 0], kern_colors)

    handles = noise_ax.plot(scales, noise_data['measure_feat'], c='k', lw=lw['all'])
    assign_colors(handles, config['k_specs'][:, 0], kern_colors)


    if save_path is not None:
        fig.savefig(save_path)
    plt.show()

print('Done.')
embed()