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 plot_functions import ylabel, super_xlabel, super_ylabel, title_subplot, time_bar
from color_functions import load_colors
from misc_functions import shorten_species
from IPython import embed

# GENERAL SETTINGS:
target_species = [
    'Chorthippus_biguttulus',
    'Chorthippus_mollis',
    'Chrysochraon_dispar',
    # 'Euchorthippus_declivus',
    'Gomphocerippus_rufus',
    'Omocestus_rufipes',
    'Pseudochorthippus_parallelus',
]
data_path = '../data/inv/rect_lp/condensed/'
save_path = '../figures/fig_invariance_rect-lp_appendix.pdf'

# ANALYSIS SETTINGS:
relate_to_noise = True
exclude_zero = True
cutoffs = np.array([np.nan, 2500, 250, 25])
search_kwargs = dict(
    incl=['noise', 'norm-base' if relate_to_noise else 'unnormed'],
    dir=data_path,
)

# GRAPH SETTINGS:
fig_kwargs = dict(
    figsize=(32/2.54, 16/2.54),
    nrows=cutoffs.size,
    ncols=len(target_species),
    sharex=True,
    sharey=True,
    gridspec_kw=dict(
        wspace=0.4,
        hspace=0.2,
        left=0.12,
        right=0.98,
        bottom=0.1,
        top=0.95,
    )
)

# PLOT SETTINGS:
species_colors = load_colors('../data/species_colors.npz')
line_kwargs = dict(
    lw=2,
    alpha=0.5,
    zorder=2,
)
fill_kwargs = dict(
    alpha=0.3,
    zorder=1,
)
mean_kwargs = dict(
    lw=2,
    alpha=1,
    zorder=3,
    ls='--'
)
mean_colors = {
    'Chorthippus_biguttulus': (1,) * 3,
    'Chorthippus_mollis': (0,) * 3,
    'Chrysochraon_dispar': (0,) * 3,
    'Euchorthippus_declivus': (0,) * 3,
    'Gomphocerippus_rufus': (0,) * 3,
    'Omocestus_rufipes': (0,) * 3,
    'Pseudochorthippus_parallelus': (1,) * 3,
}
xlab = 'scale $\\alpha$'
ylabs = ['$\\text{unfiltered}$'] + [f'${int(cutoff)}\\,\\text{{Hz}}$' for cutoff in cutoffs[1:]]
super_ylab = '$\\sigma_{\\text{env}}\\,/\\,\\sigma_{\\eta}$' if relate_to_noise else '$\\sigma_{\\text{env}}$'
xlab_kwargs = dict(
    y=0,
    fontsize=16,
    ha='center',
    va='bottom',
)
ylab_kwargs = dict(
    x=0.05,
    fontsize=16,
    ha='center',
    va='top',
)
ylab_super_kwargs = dict(
    x=0,
    fontsize=20,
    ha='left',
    va='center',
)
title_kwargs = dict(
    x=0.5,
    yref=0.99,
    ha='center',
    va='top',
    fontsize=16,
    fontstyle='italic',
)
letter_kwargs = dict(
    x=0.005,
    y=0.99,
    fontsize=22,
    ha='left',
    va='top',
)

# Prepare graph:
fig, axes = plt.subplots(**fig_kwargs)
[ylabel(ax, lab, transform=fig.transFigure, **ylab_kwargs) for ax, lab in zip(axes[:, 0], ylabs)]
super_xlabel(xlab, fig, axes[-1, 0], axes[-1, -1], **xlab_kwargs)
super_ylabel(super_ylab, fig, axes[0, 0], axes[-1, 0], **ylab_super_kwargs)

# Run through species:
for i, (species, spec_axes) in enumerate(zip(target_species, axes.T)):
    title_subplot(spec_axes[0], shorten_species(species), ref=fig, **title_kwargs)
    
    # Load species data:
    path = search_files(species, **search_kwargs)[0]
    data, config = load_data(path, files=['scales', 'mean_env', 'sd_env'])
    scales = data['scales']
    means = data['mean_env']
    sds = data['sd_env']

    if exclude_zero:
        # Exclude zero scale:
        inds = scales > 0
        scales = scales[inds]
        means = means[inds, ...]
        sds = sds[inds, ...]

    # Run through cutoffs:
    for j, ax in enumerate(spec_axes):
        # Plot recording-specific traces:
        for k in range(means.shape[-1]):
            ax.plot(scales, means[:, j, k], c=species_colors[species], **line_kwargs)
            spread = (means[:, j, k] - sds[:, j, k], means[:, j, k] + sds[:, j, k])
            ax.fill_between(scales, *spread, color=species_colors[species], **fill_kwargs)
        # Plot cutoff-specific mean trace:
        ax.plot(scales, means[:, j, :].mean(axis=-1), c=mean_colors[species], **mean_kwargs)

# Posthocs:
sylog_kwargs = dict(linthresh=scales[scales > 0][0], linscale=0.5)
axes[0, 0].set_xscale('symlog', **sylog_kwargs)
axes[0, 0].set_yscale('symlog', **sylog_kwargs)
axes[0, 0].set_xlim(scales[0], scales[-1])
axes[0, 0].set_ylim(0.9, scales[-1])
axes[0, 0].xaxis.set_major_locator(plt.LogLocator(base=10, subs=[1]))

# Save graph:
fig.savefig(save_path)
plt.show()