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, letter_subplots, 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_species.pdf'

# 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,
)
xlab = 'scale $\\alpha$'
ylab = '$\\sigma_{\\alpha}\\,/\\,\\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)
# letter_subplots(axes, **letter_kwargs)

# 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, dir=data_path)[0]
    data = dict(np.load(path))
    scales = data['scales']
    means = data['mean']
    sds = data['sd']

    # 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)

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