# Autogenerated using:
# - Template: 'src/cubnm/sim/simgroup.py.mako'
# - Model definition: 'codegen/recipes/rww.yaml'
# Do not modify this autogenerated code. Instead modify template
# or model definition.
from .base import SimGroup
import sys
import logging
import numpy as np

logging.basicConfig(stream=sys.stdout, level=logging.WARNING)
logger = logging.getLogger(__name__)

class rWWSimGroup(SimGroup):
    model_name = "rWW"
    global_param_names = ['G']
    regional_param_names = ['w_p', 'J_N', 'wIE', 'sigma']
    state_names = ['I_E', 'I_I', 'r_E', 'r_I', 'S_E', 'S_I']
    sel_state_var = "r_E"
    n_noise = 2
    default_params = {
        "G": None,  # required
        "w_p": 1.4,
        "J_N": 0.15,
        "wIE": 1.0,
        "sigma": 0.01,
    }

    def __init__(
        self, 
        *args, 
        do_fic = True, 
        max_fic_trials = 0, 
        I_SAMPLING_START = 1000, 
        I_SAMPLING_END = 10000, 
        init_delta = 0.02, 
        fic_penalty_scale = 0.5, 
        **kwargs
    ):
        r"""
        Group of reduced Wong-Wang (excitatory-inhibitory) simulations that are executed in parallel.

        Parameters
        ---------
        do_fic: :obj:`bool`
            whether to apply feedback inhibition control. If provided wIE parameters will be ignored
        max_fic_trials: :obj:`int`
            maximum number of numerical FIC trials. If set to 0, FIC will be done only analytically
        I_SAMPLING_START: :obj:`int`
            starting time of numerical FIC I_E sampling (msec)
        I_SAMPLING_END: :obj:`int`
            end time of numerical FIC I_E sampling (msec)
        init_delta: :obj:`float`
            initial delta for numerical FIC adjustment
        fic_penalty_scale: :obj:`float`
            how much deviation from FIC target mean rE of 3 Hz is penalized. Set to 0 to disable FIC penalty.
        *args, **kwargs:
            see :class:`cubnm.sim.SimGroup` for details

        Attributes
        ----------
        param_lists: :obj:`dict` of :obj:`np.ndarray`
            dictionary of parameter lists, including
                - ``'G'``: global coupling strength. Shape: (N_SIMS,)
                - ``'w_p'``: local excitatory recurrence. Shape: (N_SIMS, nodes)
                - ``'J_N'``: synaptic coupling strength. Shape: (N_SIMS, nodes)
                - ``'wIE'``: inhibitory to excitatory weight. Shape: (N_SIMS, nodes)
                - ``'sigma'``: noise amplitude. Shape: (N_SIMS, nodes)
                - ``'v'``: conduction velocity. Shape: (N_SIMS,)

        Equations
        ---------
        .. math::

            \begin{gather}
            I_i^E=W^EI_b+w_i^{p} J_i^{N} S_i^E+GJ_i^{N}\sum_{j}{C_{ij}S_j^E}-w_i^{IE}S_i^I \\
            I_i^I=W^II_b+J_i^{N}S_i^E-w^{II}S_i^I \\
            r_i^E=H^E(I_i^E)\ =\frac{a^EI_i^E-b^E}{1\ -\ e^{-d^E(a^EI_i^E-b^E)}} \\
            r_i^I=H^I(I_i^I)\ =\frac{a^II_i^I-b^I}{1\ -\ e^{-d^I(a^II_i^I-b^I)}} \\
            \dot{S_i^E}=-\frac{S_i^E}{\tau_E}+(1\ -\ S_i^E)\gamma r_i^E(t)+\sigma_i\epsilon_i^E \\
            \dot{S_i^I}=-\frac{S_i^I}{\tau_I}+r_i^I+\sigma_i\epsilon_i^I \\
            \end{gather}

        References
        ----------        
        * Deco et al. 2014 Journal of Neuroscience (10.1523/JNEUROSCI.5068-13.2014)
        * Demirtas et al. 2019 Neuron (10.1016/j.neuron.2019.01.017)
        * Deco et al. 2018 Current Biology (10.1016/j.cub.2018.07.083)
        """
        self.do_fic = do_fic
        self.max_fic_trials = max_fic_trials
        self.I_SAMPLING_START = I_SAMPLING_START
        self.I_SAMPLING_END = I_SAMPLING_END
        self.init_delta = init_delta
        self.fic_penalty_scale = fic_penalty_scale
        # parent init must be called after setting
        # model-specific attributes because
        # it sets .last_config which may include some
        # model-specific attributes
        super().__init__(*args, **kwargs)

    def get_config(self, *args, **kwargs):
        config = super().get_config(*args, **kwargs)
        config.update({
            'do_fic': self.do_fic,
            'max_fic_trials': self.max_fic_trials,
            'I_SAMPLING_START': self.I_SAMPLING_START,
            'I_SAMPLING_END': self.I_SAMPLING_END,
            'init_delta': self.init_delta,
            'fic_penalty_scale': self.fic_penalty_scale,
        })
        return config
    
    @property
    def _model_config(self):
        """
        Internal model configuration used in the simulation
        """
        model_config = super()._model_config
        model_config.update({
            'do_fic': str(int(self.do_fic)),
            'max_fic_trials': str(self.max_fic_trials),
            'I_SAMPLING_START': str(self.I_SAMPLING_START),
            'I_SAMPLING_END': str(self.I_SAMPLING_END),
            'init_delta': str(self.init_delta),
            'fic_penalty_scale': str(self.fic_penalty_scale),
        })
        return model_config

    @property
    def labels(self):
        """
        Labels of parameters and state variables
        to use in plots and reports
        """
        labels = super().labels
        labels.update({
            'G': r'G',
            'w_p': r'$w^{p}$',
            'J_N': r'$J^{N}$',
            'wIE': r'$w^{IE}$',
            'sigma': r'$\sigma$',
            'I_E': r'$I^E$',
            'I_I': r'$I^I$',
            'r_E': r'$r^E$',
            'r_I': r'$r^I$',
            'S_E': r'$S^E$',
            'S_I': r'$S^I$',
        })
        return labels
    
    @property
    def _to_clear(self):
        """
        List of attributes to be cleared when clear() is called
        """
        attrs = super()._to_clear + [
            "fic_unstable",
            "fic_failed",
            "fic_ntrials",
        ]
        return [attr for attr in attrs if hasattr(self, attr)]

    @SimGroup.N.setter
    def N(self, N):
        super(rWWSimGroup, rWWSimGroup).N.__set__(self, N)
        if self.do_fic:
            self.param_lists["wIE"] = np.zeros((self._N, self.nodes), dtype=float)

    def _process_out(self, out):
        super()._process_out(out)
        if self.do_fic:
            # FIC stats
            self.fic_unstable = self._global_bools[0]
            self.fic_failed = self._global_bools[1]
            self.fic_ntrials = self._global_ints[0]
            # write FIC (+ numerical adjusted) wIE to param_lists
            self.param_lists["wIE"] = self._regional_params[2].reshape(self.N, -1)

    def score(self, emp_fc_tril=None, emp_fcd_tril=None, emp_bold=None, **kwargs):
        """
        Calculates individual goodness-of-fit terms and aggregates them.
        In FIC models also calculates fic_penalty. Note that while 
        FIC penalty is included in the cost function of optimizer, 
        it is not included in the aggregate GOF.

        Parameters
        --------
        emp_fc_tril: :obj:`np.ndarray`
            1D array of empirical FC lower triangle. Shape: (edges,)
        emp_fcd_tril: :obj:`np.ndarray`
            1D array of empirical FCD lower triangle. Shape: (window_pairs,)
        emp_bold: :obj:`np.ndarray` or None
            cleaned and parcellated empirical BOLD time series. Shape: (nodes, volumes)
            Motion outliers should either be excluded (not recommended as it disrupts
            the temporal structure) or replaced with zeros.
            If provided emp_fc_tril and emp_fcd_tril will be ignored.
        **kwargs:
            keyword arguments passed to `cubnm.sim.SimGroup.score`

        Returns
        -------
        scores: :obj:`pd.DataFrame`
            The goodness of fit measures (columns) of each simulation (rows)
        """
        scores = super().score(emp_fc_tril, emp_fcd_tril, emp_bold, **kwargs)
        # calculate FIC penalty
        if self.do_fic & (self.fic_penalty_scale > 0):
            # calculate time-averaged r_E in each simulation-node
            if self.states_ts:
                mean_r_E = self.sim_states["r_E"][:, self.n_states_samples_remove:].mean(axis=1)
            else:
                mean_r_E = self.sim_states["r_E"]
            for idx in range(self.N):
                # absolute deviation from target of 3 Hz
                diff_r_E = np.abs(mean_r_E[idx, :] - 3)
                if (diff_r_E > 1).sum() > 0:
                    # at least one node has a deviation greater than 1
                    # only consider deviations greater than 1
                    # in the penalty
                    diff_r_E[diff_r_E <= 1] = np.NaN
                    # within each node the FIC penalty decreases
                    # exponentially as the deviations gets smaller
                    # the penalty max value in each node is 1
                    # which is then scaled by fic_penalty_scale
                    # and averaged across nodes (with 0s/NaNs)
                    # in node with no deviation > 1 Hz
                    scores.loc[idx, "-fic_penalty"] = (
                        -np.nansum(1 - np.exp(-0.05 * (diff_r_E - 1)))
                        * self.fic_penalty_scale / self.nodes
                    )
                else:
                    scores.loc[idx, "-fic_penalty"] = 0
        return scores

    def _get_save_data(self):
        """
        Get the simulation outputs and parameters to be saved to disk

        Returns
        -------
        out_data: :obj:`dict`
            dictionary of simulation outputs and parameters
        """
        out_data = super()._get_save_data()
        if self.do_fic:
            out_data.update(
                fic_unstable=self.fic_unstable,
                fic_failed=self.fic_failed,
                fic_ntrials=self.fic_ntrials,
            )
        return out_data

    def _problem_init(self, problem):
        """
        Extends BNMProblem initialization and includes FIC penalty
        if indicated.

        Parameters
        ----------
        problem: :obj:`cubnm.optimize.BNMProblem`
            optimization problem object
        """
        if (self.do_fic) & ("wIE" in problem.het_params):
            raise ValueError(
                "In rWW wIE should not be specified as a heterogeneous parameter when FIC is done"
            )
        if problem.multiobj:
            if self.do_fic & (self.fic_penalty_scale > 0):
                problem.obj_names.append("-fic_penalty")
                problem.n_obj += 1

    def _problem_evaluate(self, problem, X, out, *args, **kwargs):
        """
        Extends BNMProblem evaluation and includes FIC penalty
        in the cost function if indicated.

        Parameters
        ----------
        X: :obj:`np.ndarray`
            the normalized parameters of current population in range [0, 1]. 
            Shape: (N, ndim)
        out: :obj:`dict`
            the output dictionary to store the results with keys ``'F'`` and ``'G'``.
            Currently only ``'F'`` (cost) is used.
        *args, **kwargs
        """
        # Note: scores (inidividual GOF measures) is passed on in kwargs 
        # in the internal mechanism of pymoo evaluation function
        scores = kwargs["scores"][-1]
        if self.do_fic & (self.fic_penalty_scale > 0):
            if problem.multiobj:
                out["F"] = np.concatenate(
                    [out["F"], -scores.loc[:, ["-fic_penalty"]].values], axis=1
                )
            else:
                out["F"] -= scores.loc[:, "-fic_penalty"].values

    @classmethod
    def _get_test_configs(cls, cpu_gpu_identity=False):
        """
        Get configs for testing the simulations

        Parameters
        ----------
        cpu_gpu_identity: :obj:`bool`
            indicates whether configs are for CPU/GPU identity tests
            in which case force_cpu is not included in the configs
            since tests will be done on both CPU and GPU

        Returns
        -------
        configs: :obj:`dict` of :obj:`list`
        """
        configs = super()._get_test_configs(cpu_gpu_identity)
        configs.update({
            # 0: no FIC
            # 1: analytical FIC
            # 2: analytical + numerical FIC
            'do_fic': [0, 1, 2]
        })
        return configs

    @classmethod
    def _get_test_instance(cls, opts):
        """
        Initializes an instance that is used in tests

        Parameters
        ----------
        opts: :obj:`dict`
            dictionary of test options

        Returns
        -------
        sim_group: :obj:`cubnm.sim.rWWSimGroup`
            simulation group object of the test simulation
            which is not run yet
        """
        # initialze sim group
        sim_group = super()._get_test_instance(opts)
        # set do_fic
        sim_group.do_fic = opts['do_fic'] > 0
        if opts['do_fic'] == 2:
            sim_group.max_fic_trials = 5
        return sim_group
