# Autogenerated using:
# - Template: 'src/cubnm/sim/simgroup.py.mako'
# - Model definition: 'codegen/recipes/kuramoto.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 KuramotoSimGroup(SimGroup):
    model_name = "Kuramoto"
    global_param_names = ['G']
    regional_param_names = ['init_theta', 'omega', 'sigma']
    state_names = ['theta']
    sel_state_var = "theta"
    n_noise = 1
    default_params = {
        "G": None,  # required
        "init_theta": None,  # required
        "omega": np.pi,
        "sigma": 0.17,
    }

    def __init__(
        self, 
        *args, 
        random_init_theta = True, 
        **kwargs
    ):
        r"""
        Group of Kuramoto simulations that are executed in parallel.

        Parameters
        ---------
        random_init_theta: :obj:`bool`
            set initial theta by randomly sampling from a uniform distribution
        *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,)
                - ``'init_theta'``: initial phase angle. Shape: (N_SIMS, nodes)
                - ``'omega'``: natural frequency. Shape: (N_SIMS, nodes)
                - ``'sigma'``: noise amplitude. Shape: (N_SIMS, nodes)
                - ``'v'``: conduction velocity. Shape: (N_SIMS,)

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

            \begin{gather}
            \dot{\theta_i} = \omega_i + G \sum_j{C_{ij} \sin(\theta_j - \theta_i)} + \sigma_i\epsilon_i \\
            \end{gather}

        References
        ----------        
        * Kuramoto et al. 1984 "Chemical Oscillations, Waves, and Turbulence"
        * Cabral et al. 2011 NeuroImage (10.1016/j.neuroimage.2011.04.010)

        Notes
        -----
        In Kuramoto model GPU and CPU results may differ under
        certain conditions (large G, presence of conduction delay, large
        sigma, etc.), especially when using larger dt.

        This discrepancy arises from subtle variations in mathematical
        function implementations between GPU and CPU hardware, which,
        depending on the system dynamics, can amplify through butterfly effects.

        As noted in NVIDIA's documentation
        (https://docs.nvidia.com/cuda/floating-point/index.html#considerations-for-a-heterogeneous-world),
        such differences do not indicate that either computation is incorrect.

        We therefore recommend to assess sensitivity of simulated
        findings in Kuramoto model (e.g. optimal simulation) to 
        CPU-based versus GPU-based simulations.
        """
        self.random_init_theta = random_init_theta
        # 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({
            'random_init_theta': self.random_init_theta,
        })
        return config
    
    @property
    def _model_config(self):
        """
        Internal model configuration used in the simulation
        """
        model_config = super()._model_config
        model_config.update({
            'random_init_theta': str(int(self.random_init_theta)),
        })
        return model_config

    def post_init(self):
        logger.warning(
            "In Kuramoto model GPU and CPU results may differ under "
            "certain conditions (large G, presence of conduction delay, large "
            "sigma, etc.), especially when using larger dt.\n"
            "This discrepancy arises from subtle variations in mathematical "
            "function implementations between GPU and CPU hardware, which, "
            "depending on the system dynamics, can amplify through butterfly effects.\n"
            "As noted in NVIDIA's documentation "
            "(https://docs.nvidia.com/cuda/floating-point/index.html#considerations-for-a-heterogeneous-world), "
            "such differences do not indicate that either computation is incorrect.\n"
            "We therefore recommend to assess sensitivity of simulated "
            "findings in Kuramoto model (e.g. optimal simulation) to "
            "CPU-based versus GPU-based simulations."
        )

    @SimGroup.N.setter
    def N(self, N):
        super(KuramotoSimGroup, KuramotoSimGroup).N.__set__(self, N)
        if self.random_init_theta:
            # sample from uniform distribution of [0, 2*pi] across nodes and
            # repeat it across simulations
            # use the same random seed as the simulation noise
            rng = np.random.default_rng(self.sim_seed)
            self.param_lists["init_theta"] = np.tile(rng.uniform(0, 2 * np.pi, self.nodes), (self._N, 1))
        else:
            self.param_lists["init_theta"] = np.zeros((self._N, self.nodes), dtype=float)
            logger.warning("init_theta is set to zero")
