// file: proso_mixture.stan
data {
    // we split the dataset into zeros and not-zeros
    // we also allow two sets of covariates, one for presence-absence and one for nonzero counts
    int <lower = 0> n_zeros;
    int <lower = 0> n_counts;
    int <lower = 1> k_pa;
    int <lower = 1> k_pois;

    // four covariate matrices:
    //      observed zeros, presence-absence process
    //      observed nonzeros, presence-absence process
    //      observed zeros, poisson (count) process
    //      observed nonzeros, poisson process
    matrix [n_zeros, k_pa] X_zeros_pa; // binomial process, observed zeros
    matrix [n_counts, k_pa] X_count_pa; // binomial process, observed nonzeros
    matrix [n_zeros, k_pois] X_zeros_pois; // poisson process, observed zeros
    matrix [n_counts, k_pois] X_count_pois; // poisson process, observed zeros

    // the observed nonzero counts
    int <lower = 1> counts [n_counts];

    // prior hyperparams
    real a_pa_scale;
    real B_pa_scale;
    real a_pois_scale;
    real B_pois_scale;
}
parameters {
    // one set of linear parameters for determining the probability of presence
    real a_pa;
    vector [k_pa] B_pa;

    // a second set of parameters for determining the count if present
    real a_count;
    vector [k_pois] B_count;
}
transformed parameters {
    // first, we have a probability of presence and an expected count for each observed zero
    vector <lower = 0, upper = 1> [n_zeros] prob_pres_zeros;
    vector <lower = 0> [n_zeros] lam_zeros;

    // then we have the same for each observed (nonzero) count
    vector <lower = 0, upper = 1> [n_counts] prob_pres_counts;
    vector <lower = 0> [n_counts] lam_counts;
    
    prob_pres_zeros = inv_logit(a_pa + X_zeros_pa * B_pa);
    lam_zeros = exp(a_count + X_zeros_pois * B_count);

    prob_pres_counts = inv_logit(a_pa + X_count_pa * B_pa);
    lam_counts = exp(a_count + X_count_pois * B_count);
}
model {
    for(i in 1:n_zeros) {
        // on the probability scale, just to see
        // in the end we must work on the log scale, so it's a bit more complicated
        //      target *= (1 - prob_pres[i]) + prob_pres[i] * poisson_pmf(0 | lam_zeros[i]);

        // log_sum_exp performs the computation above, but keeping all values on the log scale
        // log_sum_exp(x1, x2) is equivalent to log(e^x1 + e^x2), but it never performs exponentiation
        // x1 and x2 are kept on the log scale, so we avoid numerical problems
        // see: https://mc-stan.org/docs/stan-users-guide/log-sum-of-exponentials.html
        target += log_sum_exp(
            // first term, the binomial term, now on the log scale
            log(1 - prob_pres_zeros[i]),
            // second term, the poisson term, on the log scale
            log(prob_pres_zeros[i]) + poisson_lpmf(0 | lam_zeros[i])
        );
    }

    // for the nonzero counts, we use a poisson likelihood as usual, with the added complication
    // that we must account for the probability of presence!
    for(i in 1:n_counts) {
        target += log(prob_pres_counts[i]) + poisson_lpmf(counts[i] | lam_counts[i]);
    }


    a_pa ~ normal(0, a_pa_scale);
    B_pa ~ normal(0, B_pa_scale);

    a_count ~ normal(0, a_pois_scale);
    B_count ~ normal(0, B_pois_scale);
}
generated quantities {
    // capture model deviance and lppd
    real deviance = 0;
    vector [n_zeros + n_counts] lppd;

    // simulate to get the PPD
    int ppd_counts [n_zeros + n_counts];

    // first simulate for all observed zeros
    for(i in 1:n_zeros) {
        // first term simulates the presence-absence part
        // then we multiply by a simulated poisson
        ppd_counts[i] = binomial_rng(1, prob_pres_zeros[i]) * poisson_rng(lam_zeros[i]);
        lppd[i] = log_sum_exp(log(1 - prob_pres_zeros[i]), 
            log(prob_pres_zeros[i]) + poisson_lpmf(0 | lam_zeros[i]));
        deviance += lppd[i];
    }

    // next simulate all observed nonzeros
    for(j in 1:n_counts) {
        ppd_counts[j + n_zeros] = binomial_rng(1, prob_pres_counts[j]) * poisson_rng(lam_counts[j]);
        lppd[j + n_zeros] = log(prob_pres_counts[j]) + poisson_lpmf(counts[j] | lam_counts[j]);
        deviance += lppd[j + n_zeros];
    }
    deviance *= -2;
}

data {
    int <lower = 1> n; // total number of districts, one data point per district
    int <lower = 1> k; // number of regression variables
    
    // spatial neighbourhood data
    // this is a sparse array
    // the district ID in column one is adjacent to column 2
    int <lower = 1> n_nb; // number of adjacencies
    int <lower = 1, upper = n> neighbours [n_nb, 2];

    // regression data
    int <lower = 0> deaths [n];
    vector <lower = 0> [n] exposure;
    matrix [n,k] X;
    
    // prior hyperparams
    real <lower=0> a_sig;
    real <lower=0> B_sig;
    
    // controls the strength of the spatial effect
    real <lower = 0> gamma_scale_sig;
}
transformed data {
    vector [n] nu = rep_vector(0, n); // number of neighbors per region
    for(i in 1:n_nb)
        nu[neighbours[i,1]] += 1;
    
}
parameters {
    // regression params
    real a;
    vector [k] B;

    // latent variable for spatial random effect
    real gamma_scale;
    vector [n] gamma;
}
transformed parameters {
    vector [n] gamma_expectation = rep_vector(0, n);
    vector <lower = 0> [n] lambda;
    
    for(i in 1:n_nb)
        gamma_expectation[neighbours[i,1]] += gamma[neighbours[i,2]];
    for(i in 1:n) {
        if(nu[i] > 0)
            gamma_expectation[i] = gamma_expectation[i] / nu[i];
    }
    lambda = exp(a + gamma + X*B);
}
model {
    deaths ~ poisson(exposure .* lambda);
    gamma ~ normal(gamma_expectation, gamma_scale);
    gamma_scale ~ normal(0, gamma_scale_sig);
    a ~ normal(0, a_sig);
    B ~ normal(0, B_sig);
}
generated quantities {
    int ppd [n];
    ppd = poisson_rng(exposure .* lambda);
}

library(sf)
library(rstan)
scotlip = st_read("../vu_advstats_students/data/scotlip.gpkg")
scotlip_nb = readRDS("../vu_advstats_students/data/scotlip_neighbours.rds")

## read neighbours
stan_cancer_car = stan_model("vu_advstats_students/stan/scotlip.stan")

X = matrix(scotlip$AFF, ncol = 1)

standata = list(    
    n = nrow(scotlip),
    k = ncol(X),
    
    n_nb = nrow(scotlip_nb),
    neighbours = scotlip_nb,
    deaths = scotlip$CANCER,
    exposure = scotlip$POP/1000,
    X = X,
    a_sig = 10,
    B_sig = 5,
    gamma_scale_sig = 10
)


# we need inits to keep the poisson function small
initfun = function() list(a = runif(1, -3,0), B = array(runif(ncol(X), -3, 0), dim = ncol(X)), gamma = runif(nrow(scotlip), -0.1, 0.1))

fit_car = sampling(cancer_car, data = standata, iter = 50000, chains = 4, control = list(max_treedepth = 20),
                   init = initfun, refresh = 0, open_progress = FALSE)