"""
Figure 7g.
"""
import sys
sys.path.append('../')
sys.path.append('../../')
sys.path.append('../../../')

assert sys.version_info >= (3, 0)

import numpy as np
import pickle as pkl
import h5py_wrapper as h5w

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

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

    data_full_path = './data/fig7_g.pkl'

    initial_rate_active = 8.5
    na_rates = np.arange(1., 8.1, 0.4)
    # na_rates = [1., 2.]
    results_all = []
    high_activity_state_th = 15.

    if not recompute:
        with open(data_full_path, 'rb') as f:
            data = pkl.load(f)
    else:
        data = {'results': {}}

    for na_r in na_rates:
        # initialized to rates in first module
        initial_rates = np.array([initial_rate_active, na_r])
        res_key = tuple(initial_rates.tolist())

        if not recompute:
            self_const_states = data['results'][res_key]
        else:
            print(f"Recomputing data for input rates {res_key}")
            self_const_states, modularity_values, limits = calc.nu_vs_layers(
                pars, rates_initial_layer=initial_rates, modularity_values=np.arange(0.76, 0.93, 0.002),
                multithreaded=True)
            data['results'][res_key] = self_const_states

        for m, rates in self_const_states.items():
            if np.max(rates) > high_activity_state_th:  # checks if activity converges to high activity state
                results_all.append(m)
                break

    if recompute:  # save new data
        with open(data_full_path, 'wb') as f:
            data['initial_rate_active'] = initial_rate_active
            data['na_rates'] = na_rates
            pkl.dump(data, f)

    mft_plotting.plot_noise_vs_layers(results_all, initial_rate_active, na_rates, pars)