// mdcev model

functions {
#include common/mdcev_ll.stan
}


data {
// declares I J NPsi K dat_psi price_j quant_j income num_price M_factorial
// prior_psi_sd prior_gamma_sd prior_alpha_sd prior_scale_sd
// model_num fixed_scale1 trunc_data flat_priors weights
#include common/mdcev_data.stan
	int K; // number of mixtures

// additional LC data that can be set to 0 if not used
	int L; // number of predictors in membership equation
//	array[I] vector[L] data_class;   // predictors for component membership
	matrix[I, L] data_class;   // predictors for component membership
	int<lower=0, upper=1> single_scale; // indicator to estimate one scale for lc
	real prior_delta_sd;
}

transformed data {
#include common/mdcev_tdata.stan

if (single_scale == 0 && fixed_scale1 == 0)
	S = K;
}

parameters {
	array[K] vector[NPsi] psi;
	array[K] vector[NPhi] phi;
	array[K] vector<lower=0 >[Gamma] gamma;
	array[K] vector<lower=0, upper=1>[A] alpha;
	vector<lower=0>[S] scale;
  	matrix[K - 1, L] delta;  // mixture regression coeffs
}

transformed parameters {
	vector[I] log_like;
	matrix[K > 1 ? I : 0, K > 1 ? K - 1 : 0] delta_data;
	{
	array[K] vector[I] log_like_util;
	for (k in 1:K){
		matrix[I, J] gamma_full = gamma_ll(gamma[k], I, J, Gamma);
		matrix[I, J] lpsi;
		vector[I] alpha_1 = alpha_1_ll(alpha[k], I, model_num);
		real scale_full = (S == 1) ? scale[1] : ((S == K) ? scale[k] : 1.0);
		vector[NPsi] psi_k = psi[k];

		if (psi_ascs == 1){
			if (NPsi_ij > 0){
				lpsi = rep_matrix(append_row(0, head(psi_k, J-1))', I);
				lpsi += to_matrix(dat_psi * segment(psi_k, J, NPsi_ij), I, J, 0);
			} else {
				lpsi = rep_matrix(append_row(0, head(psi_k, J-1))', I);
			}
		} else if (psi_ascs == 0){
			if (NPsi_ij > 0)
    			lpsi = to_matrix(dat_psi * psi_k, I, J, 0);
			else
    			lpsi = rep_matrix(0, I, J);
		}

		if(model_num < 5){
			matrix[I, J] alpha_j = alpha_j_ll(alpha[k], I, J, model_num);

			if(K == 1) {
			log_like = mdcev_ll(quant_j, price_j, log_num, income, M, log_M_fact, // data
				lpsi, gamma_full, alpha_1, alpha_j, scale_full, 						// parameters
				I, J, nonzero, trunc_data);
			} else if (K > 1){
			log_like_util[k] = mdcev_ll(quant_j, price_j, log_num, income, M, log_M_fact, // data
				lpsi, gamma_full, alpha_1, alpha_j, scale_full, 						// parameters
				I, J, nonzero, trunc_data);
			}
		} else if (model_num == 5){
			matrix[I, J] phi_ij;
			if(NPhi > 0)
	    		phi_ij = exp(to_matrix(dat_phi * phi[k], I, J, 0));
			else if(NPhi == 0)
	    		phi_ij = rep_matrix(1, I, J);

			if(K == 1) {
			log_like = kt_ll(quant_j, price_j, log_num, income,
  				lpsi, phi_ij, gamma_full, alpha_1, scale_full,
  				I, J, nonzero, trunc_data, jacobian_analytical_grad);
			} else if (K > 1){
			log_like_util[k] = kt_ll(quant_j, price_j, log_num, income,
  				lpsi, phi_ij, gamma_full, alpha_1, scale_full,
  				I, J, nonzero, trunc_data, jacobian_analytical_grad);
			}
		}
	}

	if (K > 1){
		delta_data = data_class * delta';
		for(i in 1:I){
			vector[K] ltheta = log_softmax(append_row(0, delta_data[i]'));
			vector[K] lps = ltheta + to_vector(log_like_util[,i]);
			log_like[i] = log_sum_exp(lps);
		}
	}
	}
}

model {
  // priors on the parameters
  if(flat_priors == 0){ // no priors ensure optimizing returns MLE
	scale ~ normal(0, prior_scale_sd);

  	if(K == 1){
	  psi[1] ~ normal(0, prior_psi_sd);
	  phi[1] ~ normal(0, prior_phi_sd);
	  gamma[1] ~ normal(1, prior_gamma_sd);
	  alpha[1] ~ beta(prior_alpha_shape, prior_alpha_shape);

	} else if (K > 1){
		to_vector(delta) ~ normal(0, prior_delta_sd);
		for (k in 1:K){
			to_vector(psi[k]) ~ normal(0, prior_psi_sd);
			to_vector(phi[k]) ~ normal(0, prior_phi_sd);
			to_vector(gamma[k]) ~ normal(1, prior_gamma_sd);
			to_vector(alpha[k]) ~ beta(prior_alpha_shape, prior_alpha_shape);
		}
	}
  }

target += dot_product(log_like, weights);//objective to target
}

generated quantities{
	real sum_log_lik = dot_product(log_like, weights);
	matrix[K > 1 ? I : 0, K] theta;

	if (K > 1){
		for(i in 1:I)
  			theta[i] = softmax(append_row(0, delta_data[i]'))';
	}
}
