/* Code adapted from Virgile Baudrot
https://github.com/virgile-baudrot/gutsRstan */

#include /include/license.stan

functions {

  #include /include/common_stan_functions.stan

  vector TKTD_varSD( real t,
                     vector y,
                     array[] real theta,
                     array[] real x_r,
                     array[] int  x_i) {

    // - parameters
    real kd = theta[1];
    real zw = theta[2];
    real bw = theta[3];
    real hb = theta[4];

    // - new variables
    array[2] real max_zw; //
    vector[2] dy_dt; //

    // - latent variables
    int Nconc = x_i[1]; // Number of point measuring concentration

    vector[Nconc] tconc = to_vector(x_r[1:Nconc]);
    vector[Nconc] conc = to_vector(x_r[Nconc+1:2*Nconc]); // to take dose time which is in ts

    int pulse_index = find_interval_elem(t, tconc, 1); // pulse_index index

    // if_else statement
    real conc_linInterp = pulse_index != 0 ? linearInterp(t , tconc[pulse_index], tconc[pulse_index+1], conc[pulse_index], conc[pulse_index+1] ) : conc[1];
    //real conc_linInterp = linearInterp(t , tconc[pulse_index], tconc[pulse_index+1], conc[pulse_index], conc[pulse_index+1] );
    // - model
    dy_dt[1] =  kd * ( conc_linInterp - y[1]);
    max_zw[1] = 0;
    max_zw[2] = y[1] - zw;
    dy_dt[2] = bw * max(max_zw) + hb;

    return(dy_dt);
  }

  matrix solve_TKTD_varSD(vector y0, vector ts, array[] real theta, data vector tconc, data vector conc){
        // implementation of a modified Euler method
        array[1] int x_i;
        x_i[1] = size(tconc);
        vector[2] statevars = y0;
        vector[2] deriv = y0;
        vector[2] tempres = y0;
        array[size(ts)] vector[2] ode_res;
        // this is to avoid artificially inflating the Rhat value for the first point of the time vector
        ode_res[1] = y0 + 1e-9 * TKTD_varSD(1e-9,statevars,theta,to_array_1d(append_row(tconc, conc)),x_i);;
        real tforiter;
        for (i in 2:size(ts)){
          for (j in 1:101){
            tforiter = ts[i-1] + (ts[i] - ts[i-1])/100 * (j-1);
            deriv = TKTD_varSD(tforiter,statevars,theta,to_array_1d(append_row(tconc, conc)),x_i);
            statevars = statevars + deriv * 0.01;
          }
          ode_res[i] = statevars;
        }

        matrix[size(ts), 2] rtn;
        for(i in 1:size(ts)) {
          rtn[i] = transpose(ode_res[i]);
        }
        return rtn;
      }
    }

data {

  #include /include/data_guts.stan

  /* PRIORS */
  real bwMean_log10;
  real bwSD_log10;
  real zwMean_log10;
  real zwSD_log10;
}

transformed data{

  vector<lower=0>[2]  y0;
  vector[nData_Nsurv] tNsurv_ode; // time of Nbr survival to include in the ode !
  vector[nData_conc] tconc_ode; // time of Nbr survival to include in the ode !

  y0[1] = 0;
  y0[2] = 0;

  for(gr in 1:nGroup){
    tNsurv_ode[idS_lw[gr]:idS_up[gr]] = to_vector(tNsurv[idS_lw[gr]:idS_up[gr]]);
    tNsurv_ode[idS_lw[gr]] = tNsurv[idS_lw[gr]] + 1e-9 ; // to start ode integrator at 0
    tconc_ode[idC_lw[gr]:idC_up[gr]] = to_vector(tconc[idC_lw[gr]:idC_up[gr]]);
    tconc_ode[idC_lw[gr]] = tconc[idC_lw[gr]] + 1e-9 ; // to start ode integrator at 0
  }
}

parameters {

  array[3 + nDatasets] real sigma;

}

transformed parameters{

  real kd_log10 = kdMean_log10 + kdSD_log10 * sigma[1];
  real zw_log10 = zwMean_log10 + zwSD_log10 * sigma[2];
  real bw_log10 = bwMean_log10 + bwSD_log10 * sigma[3];
  array[nDatasets] real hb_log10;
  array[4] real<lower=0> param; //

  matrix[nData_Nsurv,2] y_hat;
  vector<lower=0, upper=1>[nData_Nsurv] Psurv_hat;
  vector<lower=0, upper=1>[nData_Nsurv] Conditional_Psurv_hat;
  param[1] = 10^kd_log10; // kd
  param[2] = 10^zw_log10; // zw
  param[3] = 10^bw_log10; // bw

  for(i in  1:nDatasets){
    hb_log10[i]  = hbMean_log10 + hbSD_log10 * sigma[3+i];
  }

  for(gr in 1:nGroup){
    param[4] = 10^hb_log10[groupDataset[gr]]; // hb

    /* initial time must be less than t0 = 0, so we use a very small close small number 1e-9 to at at time tNsurv and tconc */
    y_hat[idS_lw[gr]:idS_up[gr],1:2] = solve_TKTD_varSD(y0, tNsurv_ode[idS_lw[gr]:idS_up[gr]], param, tconc_ode[idC_lw[gr]:idC_up[gr]], to_vector(conc[idC_lw[gr]:idC_up[gr]]));

    Psurv_hat[idS_lw[gr]:idS_up[gr]] = exp( - y_hat[idS_lw[gr]:idS_up[gr], 2]);

    for(i in idS_lw[gr]:idS_up[gr]){
      Conditional_Psurv_hat[i] =  i == idS_lw[gr] ? Psurv_hat[i] : Psurv_hat[i] / Psurv_hat[i-1] ;
    }
  }

}

model {

  target += normal_lpdf(sigma | 0, 1);
  for(gr in 1:nGroup){
    target += binomial_lpmf(Nsurv[idS_lw[gr]:idS_up[gr]] | Nprec[idS_lw[gr]:idS_up[gr]], Conditional_Psurv_hat[idS_lw[gr]:idS_up[gr]]);
    // Nsurv[idS_lw[gr]:idS_up[gr]] ~ binomial( Nprec[idS_lw[gr]:idS_up[gr]], Conditional_Psurv_hat[idS_lw[gr]:idS_up[gr]]);
  }
}

generated quantities {

  #include /include/gen_quantities_guts.stan
}

