import plotstyle_plt
import numpy as np
import matplotlib.pyplot as plt
from thunderhopper.filetools import search_files
from plot_functions import ylabel, super_xlabel, title_subplot
from color_functions import load_colors
from misc_functions import shorten_species

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

# ANALYSIS SETTINGS:
exclude_zero = True

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

# PLOT SETTINGS:
colors = load_colors('../data/species_colors.npz')
line_kwargs = dict(
    lw=2,
)
fill_kwargs = dict(
    alpha=0.3,
    zorder=1,
)
mean_kwargs = dict(
    # c=(0.5,) * 3,
    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': (0,) * 3,
}
xlab = 'scale $\\alpha$'
ylab = '$\\sigma_{\\text{adapt}}\\,/\\,\\sigma_{\\eta}$'
xlab_kwargs = dict(
    y=0,
    fontsize=16,
    ha='center',
    va='bottom',
)
ylab_kwargs = dict(
    x=0,
    fontsize=20,
    ha='center',
    va='top',
)
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)
axes[0].set_ylim(0.9, 20)
axes[0].set_xscale('log')
axes[0].set_yscale('log')
super_xlabel(xlab, fig, axes[0], axes[-1], **xlab_kwargs)
ylabel(axes[0], ylab, **ylab_kwargs, transform=fig.transFigure)

# Run through species:
for species, ax in zip(target_species, axes):
    title_subplot(ax, shorten_species(species), ref=fig, **title_kwargs)
    color = colors[species]
    
    # Load species data:
    path = search_files(species, incl=['noise', 'norm-base'], dir=data_path)[0]
    data = dict(np.load(path))
    scales = data['scales']
    means = data['mean_inv']
    sds = data['sd_inv']

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

    # Plot recording-specific traces:
    for mean, sd in zip(means.T, sds.T):
        ax.plot(scales, mean, c=color, **line_kwargs)
        ax.fill_between(scales, mean - sd, mean + sd, color=color, **fill_kwargs)

    # Plot species mean trace:
    ax.plot(scales, means.mean(axis=-1), c=mean_colors[species], **mean_kwargs)

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