*! Author: Hong Il Yoo (h.i.yoo@durham.ac.uk) 
*! HIY 1.1.1 07 November 2021

program define lclogitpr2, sortpreserve
	version 13.1	
	if ("`e(cmd)'" != "lclogitml2")&("`e(cmd)'" != "lclogit2") error 301
	
	syntax newvarname [if] [in] [, CLass(numlist >=1 <=`=e(nclasses)') pr0 pr up cp ] 
	
	** Check whether specified options are valid **
	if ("`class'" != "") & ("`pr0'" != "") {
		display as error "class() and pr0 cannot be specified at the same time."
		exit 184
	}	
	
	local check : word count `pr0' `pr' `up' `cp'
	if (`check' > 1) {
		display as error "only one of pr0, pr, up and cp can be specified at a time." 
		exit 184
	}
	
	** Check that group, id and other explanatory variables are numeric **
	foreach v of varlist `e(group)' `e(id)' `e(indepvars)' `e(indepvars2)' {
		capture confirm numeric variable `v'
		if _rc != 0 {
			display as error "variable `v' is not numeric."
			exit 498
		}
	}
	
	if ("`class'" == "") {
		forvalues c = 1/`=e(nclasses)' {
			local class `class' `=int(`c')'
		}
	}
		
	** Mark the prediction sample **
	marksample touse, novarlist
	
	** Define temporary objects **
	tempvar one
	qui gen `one' = 1 if `touse'
	
	** Define variables ** 
	forvalues c = 1/`=e(nclasses)' {
		tempvar pr_`c'
		// temporary variable to hold some kind of probability for class c
		qui gen double `pr_`c'' = .
		local pr_all `pr_all' `pr_`c''  
	}
	if ("`up'" == "" & "`cp'" == "") {
		tempvar pr 
		// temporary variable to hold average choice probability over classes
		qui gen double `pr' = . 
		local pr_all `pr' `pr_all'
	}
	if ("`up'" != "") local pr_type = 1
	else if ("`cp'" != "") local pr_type = 2
	else local pr_type = 3
	
	** Define macros **
	local id `e(id)'
	local group `e(group)'
	
	local nclasses = e(nclasses)
	local k_fix = e(k_fix)
	local k_rand = e(k_rand)
	local k_membership = e(k_share)
	
	local depvar `e(depvar)'
	local rand `e(indepvars_rand)'
	local membership `e(indepvars_share)'
	if (`k_fix' > 0) local fix `e(indepvars_fix)'
	
	// send basic information to Mata
	sort `id' `group' 
	mata: st_view(id=.,.,st_local("id"),st_local("touse"))  // subject id
	mata: st_view(group=.,.,st_local("group"),st_local("touse")) // choice set id
	mata: nclasses = strtoreal(st_local("nclasses")) // # of classes	
	mata: k_fix = strtoreal(st_local("k_fix")) // # of fixed preference coefs
	mata: k_rand = strtoreal(st_local("k_rand")) // # of random preference coefs	
	mata: k_membership = strtoreal(st_local("k_membership")) // # of class share coefs (excl. const)
	
	// send data to Mata
	mata: st_view(Y=.,.,"`depvar'",st_local("touse"))
	mata: st_view(X_rand=.,.,"`rand'",st_local("touse"))
	mata: st_view(X_share=.,.,"`membership' `one'",st_local("touse"))
	if ("`fix'" != "") mata: st_view(X_fix=.,.,"`fix'",st_local("touse"))			
	
	// compute predicted probabilities
	mata: lclogitml2pr(st_matrix("e(b)"),"`pr_all'", "`touse'", `=int(`pr_type')')
	
	// keep predicted probabilities 
	if ("`pr0'" == "") {
		foreach c of numlist `class' {
			gen `typlist' `varlist'`c' = `pr_`c'' if `touse'
			if ("`up'" != "") label variable `varlist'`c' "prior probability of being in class `c'"
			if ("`cp'" != "") label variable `varlist'`c' "posterior probability of being in class `c'"
			if ("`pr'" != "") label variable `varlist'`c' "probability of choice if in class `c'"
		}
	}
	if ("`up'" == "" & "`cp'" == "") {
		gen `typlist' `varlist' = `pr' if `touse'
		label variable `varlist' "probability of choice unconditional on class"
	}
end

version 13.1
mata:
function lclogitml2pr(real rowvector b, string rowvector pr_all, string scalar touse, real scalar pr_type) 
{
	//*******************************
	// Step 1. get things from Stata 
	//*******************************
	
	// initialise matrix of predicted probabilities 
	// [N x nclasses] for pr_type = 1 (unconditional membership prob) or 2 (conditional membership prob)
	// [N x (1+nclasses)] for pr_type = 3 (choice prob)
	st_view(PR=.,.,pr_all,touse)	
	
	// [N x 1] vectors
	external id // subject id 
	external group // choice set id

	// scalars 
	external nclasses // # of classes	
	external k_rand // # of random preference coefs
	external k_membership // # of class share coefs (excl. const)
	external k_fix // # of fixed preference coefs	

	k_coef = 0 // scalar to keep running total of # of coefficients

	// [N x 1] vector of dependent variable
	external Y
	
	// [N x # of coefs] matrices of regressors 
	external X_rand
	external X_share
	if (k_fix > 0) external X_fix
	
	// break down b into parameter blocks
	b_rand = b[1,1..k_rand*nclasses]
	k_coef = k_coef + k_rand*nclasses 	
	b_share = b[1,(k_coef+1)..(k_coef+(k_membership+1)*(nclasses-1))]
	k_coef = k_coef + (k_membership+1)*(nclasses-1)	
	if (k_fix > 0) {
		b_fix = b[1,(k_coef+1)..(k_coef+k_fix)]
		k_coef = k_coef + k_fix
	}	
	
	// [N x nclasses] matrices of class-specific linear indices
	// Random preference indices	
	Xb_rand = X_rand * b_rand[1,1..k_rand]' 
	for (c=2; c<=nclasses; c++) {
		Xb_rand = Xb_rand, X_rand * b_rand[1,k_rand*(c-1)+1..k_rand*c]'
	}	
	
	// Class share indices
	Xb_share = X_share * b_share[1,1..(k_membership+1)]' 
	if (nclasses > 2) {
		for (c=2; c<=nclasses-1; c++) {
			Xb_share = Xb_share, X_share * b_share[1,(k_membership+1)*(c-1)+1..(k_membership+1)*c]' 
		}	
	}
	Xb_share = Xb_share, J(rows(id),1,0) // index for last class's share is 0 (i.e. it's the base class) 
	
	// Fixed preference indices
	if (k_fix > 0) Xb_fix = (X_fix * b_fix'):* J(rows(id),nclasses,1)
	
	//**********************************
	// Step 2. transform linear indices 
	//**********************************
	// [N x nclasses] matrices of transformed indices
	// EXP : exp(sig*v) where sig is the scale function and v is the preference index 
	if (k_fix == 0) Xb_pref = Xb_rand
	else Xb_pref = Xb_rand + Xb_fix
	
	EXP = exp(Xb_pref)
	
	// Share: class shares 
	Share = exp(Xb_share) :/ quadrowsum(exp(Xb_share),1) 

	// fill in class shares in case option up (pr_type = 1) has been specified 
	if (pr_type == 1) PR[.,.] = Share
	
	//**************************************************************
	// Step 3. Compute conditional probabilities and log-likelihood
	//**************************************************************
	// set up panel information: input is [N x 1] vector "id", and output is [N_subject x 2] matrix "subject"
	subject = panelsetup(id,1)

	// # of panel units (subjects), identified by number of rows in "panel"
	N_subject = panelstats(subject)[1]
	
	// initialise [N_subject x 1] vector of each subject's log-likelihood 
	//lnfj = J(N_subject,1,.)
	
	// loop over subjects
	for(n=1; n<=N_subject; n++) {
		// read in data rows pertaining to subject n & store in a matrix suffixed _n
		Y_n = panelsubmatrix(Y,n,subject)
		EXP_n = panelsubmatrix(EXP,n,subject)
		Share_n = panelsubmatrix(Share,n,subject)
		group_n = panelsubmatrix(group,n,subject) 
	
		// set up panel information where each panel unit refers to a choice set 
		task_n = panelsetup(group_n,1)
		
		// # of choice sets for subject n
		N_task_n = panelstats(task_n)[1]
		
		// initialise [N_n x nclasses] matrix of choice probabilities where N_n is # of data rows for subject n
		Prob_n = J(rows(Y_n),nclasses,.)
		
		// loop over choice sets
		for(t=1; t<=N_task_n; t++) {
			// read in data rows pertaining to choice set t
			EXP_nt = panelsubmatrix(EXP_n,t,task_n)
			
			// fill in choice probabilities	
			Prob_n[task_n[t,1]..task_n[t,2],.] = EXP_nt :/ quadcolsum(EXP_nt,1)
		}
		// fill in subjet n's predicted choice probabilities in case option pr or pr0 (pr_type = 3) has been specified
		if (pr_type == 3) PR[subject[n,1]..subject[n,2],.] = quadrowsum(Prob_n :* Share_n,1), Prob_n
		
		// [1 x nclasses] vector of the likelihood of actual choice sequence
		ProbSeq_n = exp(quadcolsum(ln(Prob_n) :* Y_n,1))
	
		// compute subject n's log-likelihood
		//lnfj[n,1] = ln(ProbSeq_n * Share_n[1,.]')
		
		// fill in subject n's conditional membership probabilities in case option cp (pr_type = 2) has been specified
		if (pr_type == 2) PR[subject[n,1]..subject[n,2],.] = (ProbSeq_n :* Share_n[1,.])/(ProbSeq_n * Share_n[1,.]') :* J(rows(Y_n),nclasses,1)
	}
	
	// return [1 x 1] sample log-likelihood 
	//return(quadcolsum(lnfj,1)) 	
}
end