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;
}

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(soc_div);
    vector[N] z_int_div = z_scale(int_div);
    vector[N] z_res_soc_div = z_scale(res_soc_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=1> sigma;

    // vector[K] mu_x;
    // vector<lower=0>[K] sigma_x;
    // real mu_pop;
    // real<lower=0> sigma_pop;
}

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);

    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_research_area + mu);

    // m ~ beta(mu_pop, sigma_pop);
    // for (k in 1:N) {
    //    m[k] ~ normal(mu_x[primary_research_area[k]+1], sigma_x[primary_research_area[k]+1]);
    // }
    // mu_x ~ normal(0, 1);
    // mu_pop ~ normal(0, 1);
    // sigma_x ~ exponential(1);
    // sigma_pop ~ exponential(1);
}

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] 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_research_area + mu);

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