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

assert sys.version_info >= (3, 0)

import numpy as np

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/potential_single.h5'
    pars['output_label'] = 'potential_single'
    dataset_energy = {
        'potential_single': {
            'rates_Ax': np.arange(1., 500.01, 1.),  # range of rA for which to calculate U
            'dr_grid_samples': 1000
        },
    }

    modularities = np.round(np.arange(0.6, 1.001, 0.01).astype(np.float32), decimals=4)
    U_offset = 1e6  # arbitrary offset
    recompute = False  # whether to force recomputation or try to load previous results instead
    rates_Ax = dataset_energy[pars['output_label']]['rates_Ax']
    dr_grid = np.linspace(min(rates_Ax), max(rates_Ax),
                          dataset_energy[pars['output_label']]['dr_grid_samples'])
    potentials = calc.energy_single_trapezoid(dr_grid, modularities, pars, recompute, data_full_path)

    mft_plotting.plot_energy_single_input_specific(potentials, dr_grid)
