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<lower=0>[N] age;
    vector<lower=0>[N] m;
    matrix<lower=0,upper=1>[N,K] x;
    vector[N] stable;
    array [N] int<lower=0,upper=K-1> primary_research_area;
}

transformed data {
    vector[N] z_m = z_scale(m);
    vector[N] z_soc_cap = z_scale(soc_cap);
    vector[N] z_soc_div = z_scale(exp(soc_div));
    vector[N] z_int_div = z_scale(exp(int_div));
}

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

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

    real lambda_div;
    real mu_div;
    real<lower=0> sigma_div;

    vector<lower=0,upper=1>[K] mu_x;
    vector<lower=1>[K] eta;
    real<lower=0,upper=1> mu_pop;
    real<lower=1> eta_pop;
}

model {
    vector[N] beta_research_area;
    for (k in 1:N) {
        beta_research_area[k] = beta_x[primary_research_area[k]+1];
    }
    vector[N] res_soc_div = z_scale(soc_div - (lambda_div*z_int_div + mu_div));

    beta_soc_cap ~ normal(0, 1);
    beta_soc_div ~ normal(0, 1);
    beta_int_div ~ normal(0, 1);
    beta_x ~ normal(0, tau);
    beta_stable ~ normal(0, 1);

    z_soc_div ~ normal(lambda_div*z_int_div + mu_div, sigma_div);
    lambda_div ~ normal(0, 1);
    mu_div ~ normal(0, 1);
    sigma_div ~ exponential(1);

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

    z_m ~ normal(beta_soc_cap*z_soc_cap + beta_soc_div*res_soc_div + beta_int_div*z_int_div + beta_stable*stable + beta_research_area + mu, sigma);

    eta ~ pareto(1, 1.5);
    mu_x ~ uniform(0, 1);
    eta_pop ~ pareto(1, 1.5);
    mu_pop ~ uniform(0, 1);
    for (k in 1:N) {
       m[k] ~ beta_proportion(mu_x[primary_research_area[k]+1], eta[primary_research_area[k]+1]);
    }
    m ~ beta_proportion(mu_pop, eta_pop);
}

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];
        }
        vector[N] res_soc_div = z_scale(soc_div - (lambda_div*z_int_div + mu_div));
        vector[N] pred = beta_soc_cap*z_soc_cap + beta_soc_div*res_soc_div + beta_int_div*z_int_div + beta_stable*stable + beta_research_area + mu;

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