"""
Figure 7b.
"""
import sys
import os
sys.path.append(os.environ.get('NEST_PYTHON_PREFIX'))
sys.path.append('../')
sys.path.append('../../')
sys.path.append('../../../')

assert sys.version_info >= (3, 0), "Please run this script using Python3"
import numpy as np
import pickle as pkl

from mft_helpers import mft_parameters as parameters
from mft_helpers import calc
from mft_helpers import mft_plotting

recompute = False  # change to run fixed point search again, otherwise use stored data
method = 'nu0_fb433'

round_rates = False
stability_analysis = 'matrix'  # use proper stability analysis as defined by Eq. 11

# stability_analysis = 'manual'  # use manual stability analysis; to be used with line below
# round_rates = 'manual'  # round output rates to given precision


if __name__ == "__main__":
    pars = parameters.set_parameter_space()  # create parameter set
    pars = calc.derived_parameters(pars)

    pars['method'] = method
    pars['nuX'] = 12.
    pars['g'] = 12.

    modularity_values = np.arange(.615, 1.01, 0.02)

    assert stability_analysis == round_rates or (stability_analysis == 'matrix' and round_rates is False)

    filename = 'data/fig7_b.pkl'
    print(f'Processing Fig 7b for dataset: {filename}')

    if recompute:
        ## compute self-consistent stationary state (firing rate, mean and SD of input; "working point")
        results = calc.fixed_points(pars, modularity_values=np.arange(0.6, 1.001, 0.01),
                                    stability_analysis=stability_analysis, multithreaded=False)
    else:
        with open(filename, 'rb') as f:
            results = pkl.load(f)

    filename_suffix = f'not_important'
    fig, ax, fig_filename = mft_plotting.plot_fixed_points(results, pars, stability_analysis=stability_analysis,
                                                           filename_suffix=filename_suffix, save=False,
                                                           loc_legend_stability='center right',
                                                           loc_legend_stimulated='upper left',
                                                           frameon=True, legendsize=12, narrowlegend=True)

    ax.text(0.62,5e1,'fading regime', fontsize=14)
    ax.text(0.85,5e1,'active regime', fontsize=14)
    ax.fill_between([0.6,0.76],1e-5*np.ones(2),1e3*np.ones(2),color='0.75')

    mft_plotting.sns.despine(ax=ax)
    ax.set_xlim(0.6, 1.)
    ax.set_ylim(3e-1, 600)
    ax.set_title(None)
    fig.savefig(f"{pars['figure_root_path']}/fig7_b.pdf")
