functions {
    vector z_scale(vector x) {
        return (x-mean(x))/sd(x);
    }
}

data {
    int<lower=1> N;
    int<lower=1> K;
    vector<lower=0>[N] soc_cap;
    vector<lower=0>[N] soc_div;
    vector<lower=0>[N] int_div;
    vector[N] res_soc_div;
    vector<lower=0>[N] age;
    array[N] int<lower=0,upper=1> m;
    matrix<lower=0,upper=1>[N,K] x;
    vector[N] stable;
    array [N] int<lower=0,upper=K-1> primary_research_area;
    vector<lower=0>[N] productivity;
    vector<lower=0>[N] productivity_solo;
}

transformed data {
    vector[N] z_m = z_scale(to_vector(m));
    vector[N] z_soc_cap = z_scale(soc_cap);
    vector[N] z_soc_div = z_scale(soc_div);
    vector[N] z_int_div = z_scale(int_div);
    vector[N] z_res_soc_div = z_scale(res_soc_div);
    vector[N] z_age = z_scale(age);
    vector[N] z_productivity = z_scale(productivity);
    vector[N] z_productivity_solo = z_scale(productivity_solo);
}

parameters {
    real beta_soc_cap;
    real beta_soc_div;
    real beta_int_div;
    real beta_stable;
    real beta_age;
    real beta_productivity;
    real beta_productivity_solo;
    vector[K] beta_x;

    real mu;
    real<lower=0> tau;
    real<lower=1> sigma;
}

model {
    vector[N] beta_research_area;
    for (k in 1:N) {
        beta_research_area[k] = tau*beta_x[primary_research_area[k]+1];
    }

    beta_soc_cap ~ normal(0, 1);
    beta_soc_div ~ normal(0, 1);
    beta_int_div ~ normal(0, 1);
    beta_x ~ double_exponential(0, 1);
    beta_stable ~ normal(0, 1);
    beta_age ~ normal(0, 1);
    beta_productivity ~ normal(0, 1);
    beta_productivity_solo ~ normal(0, 1);

    mu ~ normal(0, 1);
    tau ~ exponential(1);
    sigma ~ pareto(1,1.5);

    m ~ bernoulli_logit(beta_soc_cap*z_soc_cap + beta_soc_div*z_res_soc_div + beta_int_div*z_int_div + beta_stable*stable + beta_age*z_age + beta_productivity*z_productivity + beta_productivity_solo*z_productivity_solo + beta_research_area + mu);
}

generated quantities {
    real R2 = 0;
    {
        vector[N] beta_research_area;
        for (k in 1:N) {
            beta_research_area[k] = beta_x[primary_research_area[k]+1]*tau;
        }
        vector[N] pred = inv_logit(beta_soc_cap*z_soc_cap + beta_soc_div*z_res_soc_div + beta_int_div*z_int_div + beta_stable*stable + beta_age*z_age + beta_productivity*z_productivity + beta_productivity_solo*z_productivity_solo + beta_research_area + mu);

        R2 = mean(square(to_vector(m)-pred))/variance(m);
        R2 = 1-R2;
    }
}