// Read the accompanying case study:
// https://mc-stan.org/learn-stan/case-studies/lotka-volterra-predator-prey.html
functions {
  vector dz_dt(real t, // time
               vector z, real alpha, real beta, real gamma, real delta) {
    real u = z[1];
    real v = z[2];
    
    real du_dt = (alpha - beta * v) * u;
    real dv_dt = (-gamma + delta * u) * v;
    
    return [du_dt, dv_dt]';
  }
}
data {
  int<lower=0> N; // number of measurement times
  array[N] real ts; // measurement times > 0
  array[2] real y_init; // initial measured populations
  array[N, 2] real<lower=0> y; // measured populations
}
transformed data {
  real rel_tol = 1e-5;
  real abs_tol = 1e-3;
  int max_num_steps = to_int(5e2);
}
parameters {
  real<lower=0> alpha, beta, gamma, delta;
  vector<lower=0>[2] z_init; // initial population
  array[2] real<lower=0> sigma; // measurement errors
}
transformed parameters {
  array[N] vector[2] z = ode_rk45_tol(dz_dt, z_init, 0.0, ts, rel_tol,
                                      abs_tol, max_num_steps, alpha, beta,
                                      gamma, delta);
}
model {
  alpha ~ normal(1, 0.5);
  gamma ~ normal(1, 0.5);
  beta ~ normal(0.05, 0.05);
  delta ~ normal(0.05, 0.05);
  sigma ~ lognormal(-1, 1);
  z_init ~ lognormal(log(10), 1);
  for (k in 1 : 2) {
    y_init[k] ~ lognormal(log(z_init[k]), sigma[k]);
    y[ : , k] ~ lognormal(log(z[ : , k]), sigma[k]);
  }
}
generated quantities {
  array[2] real y_init_rep;
  array[N, 2] real y_rep;
  for (k in 1 : 2) {
    y_init_rep[k] = lognormal_rng(log(z_init[k]), sigma[k]);
    for (n in 1 : N) {
      y_rep[n, k] = lognormal_rng(log(z[n, k]), sigma[k]);
    }
  }
}
