functions {
#include functions/primarycensored.stan
#include functions/convolve.stan
#include functions/pmfs.stan
#include functions/delays.stan
#include functions/observation_model.stan
#include functions/secondary.stan
#include functions/params.stan
}

data {
  int t;                             // time of observations
  int lt;                             // time of observations
  array[t] int<lower = 0> obs;             // observed secondary data
  array[lt] int obs_time;             // observed secondary data
  vector[t] primary;                 // observed primary data
  int burn_in;                       // time period to not use for fitting
  int any_accumulate; // Should any missing values be accumulated?
  array[t] int accumulate;  // Should missing values be accumulated (by time)
#include data/secondary.stan
#include data/delays.stan
#include data/observation_model.stan
#include data/params.stan
#include data/estimate_secondary_params.stan
}

transformed data{
  array[delay_types] int delay_type_max = get_delay_type_max(
    delay_types, delay_types_p, delay_types_id,
    delay_types_groups, delay_max, delay_np_pmf_groups
  );
}

parameters{
  // observation model
  vector<lower = delay_params_lower>[delay_params_length] delay_params;
  // raw gamma values for estimated nonparametric delay PMFs;
  // normalised within each ragged segment to give a Dirichlet draw
  vector<lower = 0>[delay_np_est_length] delay_np_est_raw;
  simplex[week_effect] day_of_week_simplex;  // day of week reporting effect
  vector<lower = params_lower, upper = params_upper>[n_params_variable] params;
}

transformed parameters {
  // combined fixed + estimated nonparametric delay PMF
  vector[delay_np_pmf_length] delay_np_pmf_use = combine_np_pmf(
    delay_np_pmf, delay_n_np_est, delay_np_est_groups,
    delay_np_est_pos, delay_np_est_raw
  );

  vector<lower=0>[t] secondary;
  // calculate secondary reports from primary

  {
    vector[t] scaled;
    vector[t] convolved = rep_vector(1e-5, t);

    // scaling of primary reports by fraction observed
    if (obs_scale) {
      real fraction_observed = get_param(
        param_id_fraction_observed, params_fixed_lookup, params_variable_lookup, params_value,
        params
      );
      scaled = scale_obs(primary, fraction_observed);
    } else {
      scaled = primary;
    }

    if (delay_id_reporting) {
      vector[delay_type_max[delay_id_reporting] + 1] reporting_rev_pmf =
        get_delay_rev_pmf(
          delay_id_reporting, delay_type_max[delay_id_reporting] + 1,
          delay_types_p, delay_types_id, delay_types_groups, delay_max,
          delay_np_pmf_use, delay_np_pmf_groups, delay_params, delay_params_groups,
          delay_dist, 0, 1, 0
        );
      convolved = convolved + convolve_to_report(scaled, reporting_rev_pmf, 0);
    } else {
      convolved = convolved + scaled;
    }

    secondary = calculate_secondary(
      scaled, convolved, obs, cumulative, historic, primary_hist_additive,
      current, primary_current_additive, t
    );
  }

  // weekly reporting effect
  if (week_effect > 1) {
    secondary = day_of_week_effect(secondary, day_of_week, day_of_week_simplex);
  }

  // truncate near time cases to observed reports
  if (delay_id_truncation) {
    vector[delay_type_max[delay_id_truncation] + 1] trunc_rev_cmf =
      get_delay_rev_pmf(
        delay_id_truncation, delay_type_max[delay_id_truncation] + 1,
        delay_types_p, delay_types_id, delay_types_groups, delay_max,
        delay_np_pmf_use, delay_np_pmf_groups, delay_params, delay_params_groups,
        delay_dist, 0, 1, 1
      );
    secondary = truncate_obs(secondary, trunc_rev_cmf, 0);
  }

  // accumulate reports
  if (any_accumulate) {
    profile("accumulate") {
      secondary = accumulate_reports(secondary, accumulate);
    }
  }
}

model {
  // penalised priors for delay distributions
  delays_lp(
    delay_params, delay_params_mean, delay_params_sd, delay_params_groups,
    delay_dist, delay_weight
  );
  delays_np_lp(delay_np_est_raw, delay_np_est_alpha);

  // parameter priors
  profile("param lp") {
    params_lp(
      params, prior_dist, prior_dist_params, params_lower, params_upper
    );
  }
  // observed secondary reports from mean of secondary reports (update likelihood)
  if (likelihood) {
    real reporting_overdispersion = get_param(
      param_id_reporting_overdispersion, params_fixed_lookup, params_variable_lookup, params_value,
      params
    );
    report_lp(
      obs[(burn_in + 1):t][obs_time], obs_time, secondary[(burn_in + 1):t],
      reporting_overdispersion, model_type, 1
    );
  }
}

generated quantities {
  array[t - burn_in] int sim_secondary;
  vector[return_likelihood > 1 ? t - burn_in : 0] log_lik;
  {
    real reporting_overdispersion = get_param(
      param_id_reporting_overdispersion, params_fixed_lookup, params_variable_lookup, params_value,
      params
    );
    // simulate secondary reports
    sim_secondary = report_rng(
      secondary[(burn_in + 1):t], reporting_overdispersion, model_type
    );
    // log likelihood of model
    if (return_likelihood) {
      log_lik = report_log_lik(
        obs[(burn_in + 1):t], secondary[(burn_in + 1):t],
        reporting_overdispersion, model_type, obs_weight
      );
    }
  }
}
