import plotstyle_plt
import numpy as np
import matplotlib.pyplot as plt
from plot_functions import xlabel, ylabel, strip_zeros, letter_subplots

# GENERAL SETTINGS:
data_path = '../data/inv/noise_env/sd_conversion.npz'
save_path = '../figures/fig_noise_env_sd_conversion.pdf'

# PLOT SETTINGS:
fig_kwargs = dict(
    figsize=(32/2.54, 16/2.54),
    nrows=2,
    ncols=1,
    sharex=True,
    sharey=True,
    gridspec_kw=dict(
        wspace=0,
        hspace=0.1,
        left=0.065,
        right=0.98,
        bottom=0.08,
        top=0.95,
    )
)
grid_line_kwargs = dict(
    visible=True,
    which='major',
    axis='both',
    color='k',
    lw=0.5,
)
trial_kwargs = dict(
    color='k',
    alpha=0.5,
    lw=0.5,
)
line_kwargs = dict(
    color='black',
    lw=1,
)
fill_kwargs = dict(
    color='k',
    alpha=0.5,
)
xlabels = dict(
    bottom='$\\text{scale }\\alpha$',
)
ylabels = dict(
    top='$\\sigma_{\\eta}\\,(PLACEHOLDER \\,\\text{realizations})$',
    bottom='$\\sigma_{\\eta}\\,(\\text{mean}\\,\\pm\\,\\text{SD})$',
)
xlab_kwargs = dict(
    y=0,
    fontsize=20,
    ha='center',
    va='bottom',
)
ylab_kwargs = dict(
    x=0,
    fontsize=20,
    ha='center',
    va='top',
)
title_kwargs = dict(
    t='$\\sigma_{\\text{filt}}\\,=$',
    x=0.5,
    y=1,
    ha='center',
    va='top',
    fontsize=20,
)
letter_kwargs = dict(
    x=0.005,
    y=0.99,
    fontsize=22,
    ha='left',
    va='top',
)

# Fetch data:
data = dict(np.load('../data/inv/noise_env/sd_conversion.npz'))
n = data['n_trials']

# Adjust parameters:
ylabels['top'] = f'$\\sigma_{{\\eta}}\\,({data["n_trials"]}\\text{{ realizations}})$'
title_kwargs['t'] += f'$\\,{strip_zeros(data["sd_factor"])}$'

# Prepare graph:
fig, (ax1, ax2) = plt.subplots(**fig_kwargs)
fig.suptitle(**title_kwargs)
ax1.grid(**grid_line_kwargs)
ax1.set_xlim(data['scales'][0], data['scales'][-1])
ax1.set_xscale('symlog', linthresh=data['scales'][1], linscale=0.5)
ax1.set_ylim(0.4, 1.2)
ylabel(ax1, ylabels['top'], transform=fig.transFigure, **ylab_kwargs)
ax2.grid(**grid_line_kwargs)
xlabel(ax2, xlabels['bottom'], transform=fig.transFigure, **xlab_kwargs)
ylabel(ax2, ylabels['bottom'], transform=fig.transFigure, **ylab_kwargs)
letter_subplots((ax1, ax2), **letter_kwargs)

# Plot individual trials:
ax1.plot(data['scales'], data['trials'], **trial_kwargs)

# Plot mean and spread across trials:
ax2.plot(data['scales'], data['mean'], **line_kwargs)
ax2.fill_between(data['scales'], data['mean'] - data['spread'], data['mean'] + data['spread'], **fill_kwargs)

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