import plotstyle_plt
import numpy as np
import matplotlib.pyplot as plt
from thunderhopper.filetools import search_files
from plot_functions import xlabel, super_ylabel
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/log_hp/saturation/'
save_path = '../figures/fig_saturation_log-hp_appendix.pdf'

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

# PLOT SETTINGS:
colors = load_colors('../data/species_colors.npz')
bar_kwargs = dict(
    ec='w',
)
mean_kwargs = dict(
    c='k',
    lw=3,
    ls='--'
)
xlab = 'scale $\\alpha$'
ylab = '$\\text{PDF}_{\\alpha}$'
xlab_kwargs = dict(
    y=0,
    fontsize=16,
    ha='center',
    va='bottom',
)
ylab_kwargs = dict(
    x=0.005,
    fontsize=16,
    ha='left',
    va='center',
)
leg_x = fig_kwargs['gridspec_kw']['left']
leg_y = fig_kwargs['gridspec_kw']['top']
leg_box = [
    leg_x,
    leg_y,
    fig_kwargs['gridspec_kw']['right'] - leg_x,
    1 - leg_y
]
leg_kwargs = dict(
    ncols=len(target_species),
    loc='upper center',
    bbox_to_anchor=leg_box,
    frameon=False,
    prop=dict(
        size=15,
        style='italic',
    ),
    borderpad=0,
    borderaxespad=0,
    handlelength=1,
    columnspacing=1,
)
text_kwargs = dict(
    x=1,
    y=1,
    fontsize=14,
    ha='right',
    va='top',
)

# Prepare graph:
fig, axes = plt.subplots(**fig_kwargs)
xlabel(axes[-1], xlab, **xlab_kwargs, transform=fig.transFigure)
super_ylabel(ylab, fig, axes[0], axes[-1], **ylab_kwargs)

# Run through species:
handles = []
for species, ax in zip(target_species, axes):
    color = colors[species]
    
    # Load species data:
    path = search_files(species, dir=data_path)[0]
    data = dict(np.load(path))
    hist = data['hist']
    bins = data['bins']
    n_songs = data['crit_scales'].size

    # Plot distribution of saturation points:
    handles.append(ax.bar(bins, hist, width=bins[1] - bins[0], fc=color, **bar_kwargs))
    ax.set_ylim(0, hist.max() * 1.05)

    # Indicate mean of distribution:
    ax.axvline(data['crit_scales'].mean(), **mean_kwargs)

    # Indicate number of songs:
    ax.text(**text_kwargs, s=f'n = {n_songs}', transform=ax.transAxes)

# Posthocs:
labels = [shorten_species(species) for species in target_species]
fig.legend(handles, labels, **leg_kwargs)
ax.set_xlim(0, bins[-1])

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