/* 
    Autogenerated using:
    - Template: 'src/ext/models/model.cu.mako'
    - Model definition: 'codegen/recipes/jr.yaml'
    Do not modify this autogenerated code. Instead modify template
    or model definition.
*/
#include "cubnm/includes.cuh"
#include "cubnm/defines.h"
#include "cubnm/models/jr.cuh"

__device__ __NOINLINE__ void JRModel::init(
    double* _state_vars, double* _intermediate_vars,
    double* _global_params, double* _regional_params,
    int* _ext_int, bool* _ext_bool,
    int* _ext_int_shared, bool* _ext_bool_shared
) {
    // y0 = 0.0
    _state_vars[0] = 0.0;
    // y1 = 0.0
    _state_vars[1] = 0.0;
    // y2 = 0.0
    _state_vars[2] = 0.0;
    // y3 = 0.0
    _state_vars[3] = 0.0;
    // y4 = 0.0
    _state_vars[4] = 0.0;
    // y5 = 0.0
    _state_vars[5] = 0.0;
    // y1_y2 = 0.0
    _state_vars[6] = 0.0;
    // s_y1_y2 = 0.0
    _state_vars[7] = 0.0;
}

__device__ __NOINLINE__ void JRModel::restart(
    double* _state_vars, double* _intermediate_vars, 
    double* _global_params, double* _regional_params,
    int* _ext_int, bool* _ext_bool,
    int* _ext_int_shared, bool* _ext_bool_shared
) {
    // y0 = 0.0
    _state_vars[0] = 0.0;
    // y1 = 0.0
    _state_vars[1] = 0.0;
    // y2 = 0.0
    _state_vars[2] = 0.0;
    // y3 = 0.0
    _state_vars[3] = 0.0;
    // y4 = 0.0
    _state_vars[4] = 0.0;
    // y5 = 0.0
    _state_vars[5] = 0.0;
    // y1_y2 = 0.0
    _state_vars[6] = 0.0;
    // s_y1_y2 = 0.0
    _state_vars[7] = 0.0;
}

__device__ void JRModel::step(
        double* _state_vars, double* _intermediate_vars,
        double* _global_params, double* _regional_params,
        double& tmp_globalinput,
        double* noise, long& noise_idx
        ) {
    // Sigmoid terms
    // sigm_y1_y2 = 2.0 * nu_max / (1.0 + exp(r * (v0 - (y1 - y2))))
    _intermediate_vars[0] = 2.0 * d_JRc.nu_max / (1.0 + exp(d_JRc.r * (d_JRc.v0 - (_state_vars[1] - _state_vars[2]))));
    // sigm_y0_1  = 2.0 * nu_max / (1.0 + exp(r * (v0 - (a_1 * J * y0))))
    _intermediate_vars[1] = 2.0 * d_JRc.nu_max / (1.0 + exp(d_JRc.r * (d_JRc.v0 - (_regional_params[1] * _regional_params[0] * _state_vars[0]))));
    // sigm_y0_3  = 2.0 * nu_max / (1.0 + exp(r * (v0 - (a_3 * J * y0))))
    _intermediate_vars[2] = 2.0 * d_JRc.nu_max / (1.0 + exp(d_JRc.r * (d_JRc.v0 - (_regional_params[3] * _regional_params[0] * _state_vars[0]))));

    // State variable derivatives
    // y0_dot = y3
    _intermediate_vars[3] = _state_vars[3];
    // y1_dot = y4
    _intermediate_vars[4] = _state_vars[4];
    // y2_dot = y5
    _intermediate_vars[5] = _state_vars[5];
    // y3_dot = A * a * sigm_y1_y2 - 2.0 * a * y3 - (pow(a, 2)) * y0
    _intermediate_vars[6] = d_JRc.A * d_JRc.a * _intermediate_vars[0] - 2.0 * d_JRc.a * _state_vars[3] - (pow(d_JRc.a, 2)) * _state_vars[0];
    // y4_dot = A * a * (mu + a_2 * J * sigm_y0_1 + G * globalinput) - 2.0 * a * y4 - (pow(a, 2)) * y1
    _intermediate_vars[7] = d_JRc.A * d_JRc.a * (d_JRc.mu + _regional_params[2] * _regional_params[0] * _intermediate_vars[1] + _global_params[0] * tmp_globalinput) - 2.0 * d_JRc.a * _state_vars[4] - (pow(d_JRc.a, 2)) * _state_vars[1];
    // y5_dot = B * b * (a_4 * J * sigm_y0_3) - 2.0 * b * y5 - (pow(b, 2)) * y2
    _intermediate_vars[8] = d_JRc.B * d_JRc.b * (_regional_params[4] * _regional_params[0] * _intermediate_vars[2]) - 2.0 * d_JRc.b * _state_vars[5] - (pow(d_JRc.b, 2)) * _state_vars[2];

    // Integration step
    // y0 += dt * y0_dot
    _state_vars[0] += d_JRc.dt * _intermediate_vars[3];
    // y1 += dt * y1_dot
    _state_vars[1] += d_JRc.dt * _intermediate_vars[4];
    // y2 += dt * y2_dot
    _state_vars[2] += d_JRc.dt * _intermediate_vars[5];
    // y3 += dt * y3_dot + sqrt_dt * noise * sigma
    _state_vars[3] += d_JRc.dt * _intermediate_vars[6] + d_JRc.sqrt_dt * noise[noise_idx] * _regional_params[5];
    // y4 += dt * y4_dot
    _state_vars[4] += d_JRc.dt * _intermediate_vars[7];
    // y5 += dt * y5_dot
    _state_vars[5] += d_JRc.dt * _intermediate_vars[8];

    // Calculate y1_y2 and s(y1_y2) which
    // are the non-integrated state variables
    // needed for BOLD and connectivity
    // y1_y2 = y1 - y2
    _state_vars[6] = _state_vars[1] - _state_vars[2];
    // s_y1_y2 is the sigmoid of y1_y2 which will be the
    // connectivity state variable, i.e., the output to
    // the other connected nodes. This is calculated based
    // on SigmoidalJansenRit.pre() method in TVB
    // s_y1_y2 = cmin + (cmax - cmin) / (1.0 + exp(r * (midpoint - (y1_y2))))
    _state_vars[7] = d_JRc.cmin + (d_JRc.cmax - d_JRc.cmin) / (1.0 + exp(d_JRc.r * (d_JRc.midpoint - (_state_vars[6]))));
}
