# Autogenerated using:
# - Template: 'src/cubnm/sim/simgroup.py.mako'
# - Model definition: 'codegen/recipes/wc.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 WCSimGroup(SimGroup):
    model_name = "WC"
    global_param_names = ['G']
    regional_param_names = ['c_EE', 'c_IE', 'c_EI', 'c_II', 'P_E', 'P_I', 'sigma_E', 'sigma_I']
    state_names = ['E', 'I']
    sel_state_var = "E"
    n_noise = 2
    default_params = {
        "G": None,  # required
        "c_EE": 16,
        "c_IE": 12,
        "c_EI": 15,
        "c_II": 3,
        "P_E": 0.5,
        "P_I": 0.0,
        "sigma_E": 5e-05,
        "sigma_I": 5e-05,
    }

    def __init__(self, *args, **kwargs):
        r"""
        Group of Wilson-Cowan simulations that are executed in parallel.

        Parameters
        ---------
        *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,)
                - ``'c_EE'``: local E to E coupling. Shape: (N_SIMS, nodes)
                - ``'c_IE'``: local I to E coupling. Shape: (N_SIMS, nodes)
                - ``'c_EI'``: local E to I coupling. Shape: (N_SIMS, nodes)
                - ``'c_II'``: local I to I coupling. Shape: (N_SIMS, nodes)
                - ``'P_E'``: excitatory background drive. Shape: (N_SIMS, nodes)
                - ``'P_I'``: inhibitory background drive. Shape: (N_SIMS, nodes)
                - ``'sigma_E'``: excitatory noise strength. Shape: (N_SIMS, nodes)
                - ``'sigma_I'``: inhibitory noise strength. Shape: (N_SIMS, nodes)
                - ``'v'``: conduction velocity. Shape: (N_SIMS,)

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

            \begin{gather}
            \tau^E\dot{E_i} = -E + (1 - E) S^E[c^{EE} E_i - c^{IE}I_i + G \sum_j{C_{ij} E_j} + P_i^E] + \sigma_i^E\epsilon_i^E \\
            \tau^I\dot{I_i} = -I + (1 - I) S^I[c^{EI} E_i - c^{II}I_i + P_i^I] + \sigma_i^I\epsilon_i^I, \text{where:} \\
            S^E(x) = \frac{1}{1 + e^{-a^E(x - \mu^E)}} \\
            S^I(x) = \frac{1}{1 + e^{-a^I(x - \mu^I)}} \\
            \end{gather}

        References
        ----------        
        * Wilson et al. 1972 Biophysical Journal (10.1016/S0006-3495(72)86068-5)
        * Papadopoulos et al. 2020 PLOS Computational Biology (10.1371/journal.pcbi.1008144)
        """
        # 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)

    @property
    def labels(self):
        """
        Labels of parameters and state variables
        to use in plots and reports
        """
        labels = super().labels
        labels.update({
            'G': r'G',
            'E': r'E',
            'I': r'I',
            'c_EE': r'$c_{EE}$',
            'c_IE': r'$c_{IE}$',
            'c_EI': r'$c_{EI}$',
            'c_II': r'$c_{II}$',
            'P_E': r'$P_E$',
            'P_I': r'$P_I$',
            'sigma_E': r'$\sigma_E$',
            'sigma_I': r'$\sigma_I$',
        })
        return labels
