/* block for user defined functions */
functions {
  // The arguments to this function are:
  // - `time`: the time to evaluate the DAE system
  // - `state`: the state of the DAE system at the time specified
  // - `state_derivative`: the time derivatives of the state of the DAE system at the time specified
  // - `...`: sequence of arguments passed unmodified from the DAE solve function call.
  vector residual(real z, vector state, vector deriv, real lambda, real Omega_m, real Omega_r) {
    real E = state[1];
    real lhs = (E^2 - 2*lambda)*exp(lambda*E^-2);  // left hand side of the equation for E
    real rhs = Omega_m*(1+z)^3 + Omega_r*(1+z)^4;  // right hand side of the equation for E

    // compute the residuals
    vector[2] res;
    res[1] = rhs - lhs;       // E satisfies the algebraic relation, meaning that res[1] should be 0
    res[2] = 1/E - deriv[2];  // (? -> why is this of this form)

    return res;
  }
}

/* block for datasets used
- must match the columns for each input .csv file, and the number of events must be N1, N2, ...
- if two file have the same columns then they will stack as one big array.
- the order of the files in the CLI must be the same as the order on which we define each data variable here, unless we're stacking several measurements with the same column names, those can show up anywhere */
data {
  int N1;
  array[N1] real redshift;
  array[N1] real luminosity_distance;
  array[N1] real error;
}

/* block for transformed data
- only evaluated in the beginning of each chain */
transformed data {
}

/* block for model parameters */
parameters {
  real<lower=0> h;
  real<lower=0,upper=1> Omega_m;
  real<lower=0,upper=1> Omega_r;
}

/* block for transformed parameters
- allows derived quantities from parameters and data
- is evaluated on each leapfrog step */
transformed parameters {
  array[N1] real dGW;   // luminosity distance for gravitational waves
  real correction;      // correction from modification to gravity
  real E;               // E(z)
  real k = 2.9979;      // c/H0 = (k Gpc)/h

  // quantity derived from Omega_m and Omega_r
  real lambda;
  lambda = 0.5 + lambert_w0( -(Omega_m + Omega_r)/(2*exp(0.5)) );

  // compute dE/dz(z = 0), knowing that E(0) = 1
  real deriv_0;
  deriv_0 = 1/(2*exp(lambda)) * (3*Omega_r + 4*Omega_m) / (1 - lambda + 2*lambda^2);

  // call the differential algebraic equation solver
  // (? -> what are the ' following the arrays? is it because its a vector?)
  // `residual`: DAE residual function,
  // `initial_state`: initial state, type vector
  // `initial_state_derivative`: time derivative of the initial state, type vector
  // `initial_time`: initial time, type real
  // `times`: solution times, type array[] real
  // `...`: sequence of arguments that will be passed through unmodified to the DAE residual function.
  array[N1] vector[2] S;
  S = dae(residual, [1, 0.0]', [deriv_0, 1]', 0.0, redshift, lambda, Omega_m, Omega_r);

  // compute theoretical luminosity distance for gravitational waves for each measured redshift
  for (i in 1:N1) {
    E = S[i,1];
    correction = ((1-lambda)/(1-lambda/E^2))^0.5 * exp(lambda/2 * (E^2 - 1)/E^2);
    dGW[i] = correction * (1.0+redshift[i]) * (k/h) * S[i,2];
  }
}

/* block for model definition */
model {
  // priors
  h ~ normal(0.68, 10);
  Omega_m ~ normal(0.353, 0.5);
  Omega_r ~ normal(0.021, 0.5);

  // likelihood
  luminosity_distance ~ normal(dGW, error);
}

/* block for generated quantites
- allows derived quantities based on parameters, data, and rng
- is executed once per iteration */
generated quantities {
  // output the natural logarithm of the likelihood to use in model selection criteria
  vector[N1] log_lik;
  for (n in 1:N1) {
    log_lik[n] = normal_lpdf(luminosity_distance[n] | dGW[n], error[n]);
  }
}