functions {
    real partial_sum_lpdf(array[] vector X,
        int start, int end,
        int R,
        int C,
        array[,] int indices,
        array[,] int NR,
        array[,] int NC,
        array[] vector cov,
        array[] vector expertise,
        matrix mu, vector L_sigma,
        array[] vector delta,
        array[] vector gamma,
        array[] vector beta_vec
        ) {
            vector[C] theta;
            matrix[C,R] beta;
            vector [C] tmp = rep_vector(0, C);

            real lpmf = 0;

            for (i in start:end) {
                for(r in 1:R) {
                    if (NR[i,r] == 0) {
                        beta[:,r] = rep_vector(0, C); 
                    }
                    else {
                        lpmf += normal_lpdf(beta_vec[indices[i,r]] | 0, 1);
                        tmp = mu[r,:]';//'
                        tmp[1:C-1] += beta_vec[indices[i,r]] .* L_sigma;
                        beta[:,r] = softmax(tmp + delta[r] .* cov[i] + gamma[r] .* expertise[i]); 
                    }
                }

                theta = beta*X[i-start+1,:];
                lpmf += multinomial_lpmf(NC[i,:] | theta);
            }
            return lpmf;
    }
}

data {
    int<lower=1> n_units;
    int<lower=2> R;
    int<lower=2> C;

    array[n_units,R] int NR;
    array[n_units,C] int NC;
    array[n_units] vector[C] cov;
    array[n_units] vector[C] expertise;
    matrix<lower=0,upper=1>[R,C] nu;

    int<lower=1> threads;
}

transformed data {
    array [n_units] int population;
    array[n_units] vector<lower=0,upper=1>[R] X;
    array[n_units,R] int indices;
    int largest_index = 1;
    
    for (i in 1:n_units) {
        population[i] = sum(NR[i,:]);
        for (r in 1:R) {
            X[i,r] = (NR[i,r]*1.0)/(population[i]*1.0);

            if (NR[i,r] > 0) {
                indices[i,r] = largest_index;
                largest_index += 1;
            }
            else {
                indices[i,r] = 0;
            }
        }
    }

    print("largest index:", largest_index);
    print("RxN: ", R*n_units);
}

parameters {
    matrix[R, C-1] mu_;
    array[R] vector[C] delta;
    array[R] vector[C] gamma;
    array [largest_index] vector[C-1] beta_vec;

    real delta_0;
    real delta_nu;
    real mu_nu;

    vector<lower=0>[C-1] L_sigma;
}

model {
    matrix[R, C] mu;
    for (r in 1:R) {
        mu[r,:C-1] = mu_[r];
        mu[r,C] = 0;
        mu[r,:] += mu_nu*nu[r,:];
    }

    L_sigma ~ exponential(1);
    
    for (r in 1:R) {
        mu_[r] ~ normal(0, 1);
        delta[r] ~ normal(delta_0+delta_nu*nu[r,:], 1);
        gamma[r] ~ normal(0, 1);
    }

    delta_0 ~ normal(0, 1);
    delta_nu ~ normal(0, 1);
    mu_nu ~ normal(0, 1);

    target += reduce_sum(
        partial_sum_lpdf, X, n_units%/%(threads*40),
        R, C, indices, NR, NC, cov, expertise,
        mu, L_sigma,
        delta, gamma,
        beta_vec
    );
}

generated quantities {
    array [n_units,R] vector[C] beta;
    array [R] vector[C] counts;

    for (r in 1:R) {
        counts[r] = rep_vector(0, C);
    }

    vector [C] tmp = rep_vector(0, C);

    matrix[R, C] mu;
    for (r in 1:R) {
        mu[r,:C-1] = mu_[r];
        mu[r,C] = 0;
        mu[r,:] += mu_nu*nu[r,:];
    }

    
    for (r in 1:R) {
        for (i in 1:n_units) {
            if (NR[i,r] == 0) {
                beta[i,r] = rep_vector(0, C);
            }
            else {
                tmp = mu[r,:]';//'
                tmp[1:C-1] += beta_vec[indices[i,r]] .* L_sigma;
                beta[i,r] = softmax(tmp + delta[r] .* cov[i] + gamma[r] .* expertise[i]);
                counts[r] += NR[i,r]*beta[i,r,:];
            }
            
        }
        counts[r] = counts[r]/sum(population);
    }
}