# Autogenerated using:
# - Template: 'src/cubnm/sim/simgroup.py.mako'
# - Model definition: 'codegen/recipes/jr.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 JRSimGroup(SimGroup):
    model_name = "JR"
    global_param_names = ['G']
    regional_param_names = ['J', 'a_1', 'a_2', 'a_3', 'a_4', 'sigma']
    state_names = ['y0', 'y1', 'y2', 'y3', 'y4', 'y5', 'y1_y2', 's_y1_y2']
    sel_state_var = "y1_y2"
    n_noise = 1
    default_params = {
        "G": None,  # required
        "J": 135.0,
        "a_1": 1.0,
        "a_2": 0.8,
        "a_3": 0.25,
        "a_4": 0.25,
        "sigma": 0.0001,
    }

    def __init__(self, *args, **kwargs):
        r"""
        Group of Jansen-Rit 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,)
                - ``'J'``: average number of synapses. Shape: (N_SIMS, nodes)
                - ``'a_1'``: average probability of synapses in feedback excitatory loop. Shape: (N_SIMS, nodes)
                - ``'a_2'``: average probability of synapses in slow feedback excitatory loop. Shape: (N_SIMS, nodes)
                - ``'a_3'``: average probability of synapses in feedback inhibitory loop. Shape: (N_SIMS, nodes)
                - ``'a_4'``: average probability of synapses in slow feedback inhibitory loop. Shape: (N_SIMS, nodes)
                - ``'sigma'``: local noise sigma. Shape: (N_SIMS, nodes)
                - ``'v'``: conduction velocity. Shape: (N_SIMS,)

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

            \begin{gather}
            \dot{y_{0_{i}}} = y_{3_{i}} \\
            \dot{y_{1_{i}}} = y_{4_{i}} \\
            \dot{y_{2_{i}}} = y_{5_{i}} \\
            \dot{y_{3_{i}}} = A a [\frac{2 \nu_{max}}{1+e^{r(v_0 - (y_{2_{i}} - y_{1_{i}}))}}]-2 a y_{3_{i}}-a^2 y_{0_{i}} + \sigma_i\epsilon_i \\
            \dot{y_{4_{i}}} = A a [\mu+a_{2_{i}} J_i \frac{2\nu_{max}}{1+e^{r(v_0 - (a_{1_{i}} J y_{1_{i}}))}}+ \sum_j{C_{ij}x_j} \\
            \dot{y_{5_{i}}} = B b (a_{4_{i}} J_i [\frac{2\nu_{max}}{1 + e^{r(v_0 - (a_{3_{i}} J y_{0_{i}})}}]) - 2by_{5_{i}} - b^2 y_{2_{i}} \\
            z_i = y_{1_{i}} - y_{2_{i}} \\
            x_i = c_{min} + \frac{c_{max}-c_{min}}{1 + e^{r(c_{mid} -z_i)}} \\
            \end{gather}

        References
        ----------        
        * Jansen and Rit 1995 Biological Cybernetics (10.1007/BF00199471)
        * Ritter et al. 2013 Brain Connectivity (10.1089/brain.2012.0120)
        """
        # 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)
