// Hashiguchi-Tanaka 2014 Agglomeration and Firm-level Productivity: A Bayesian Spatial Approach
// IV (Rossi type) model
// fn = Z*G + Dm*ga3 + ep1
// y = betaA*fn + X*B + Dm*alp + ep2
// ga3 = mu_ga + eta1
// alp = mu0*1 + mu1*UP + rho*Walp + eta2
// X = h ~ ExpDm ~ k ~ l ~ IndDm ~ OwnDm

#include <oxstd.h>
#include <oxprob.h>
#include <oxfloat.h>
#include <oxdraw.h>
#include "ReportMCMC.ox"

// ReportMCMC.ox: see Omori's website 
// http://www.omori.e.u-tokyo.ac.jp/Ox/Sample/MCMC/iwanami/index.html

/////////////////////////////////////////////////
// Generation of Wishart Distribution using
// Bartlett's (1933) decompositon
// See Johnson, M.E. (1986)
// "Multivariate Statistical Simulation"
// Wiley. (page 204)
/////////////////////////////////////////////////
fMyWishart(const n,const mS){
//Omori's Home page (http://www.omori.e.u-tokyo.ac.jp/Ox/Sample/MCMC/Wishart_ox.htm)
// X~W(n,S), X:(pxp) symmetrx matrix
// f(X|n,S)
// = const *|X|^{(n-p-1)/2}*exp-0.5*tr(S^{-1}X)
// n: degrees of freedom 
// S: (p x p) symmetric parameter matrix
// E(X_{ij})=n*s_{ij}, s_{ij}: (i,j) element of S.
decl ci,cj,cp,mA,mL,mT,mX;
cp=rows(mS);mT=zeros(cp,cp);
if(n<cp){print("Error! d.f. too small.");}
// generate A~W(n,I)
for(ci=0;ci<cp;++ci){
mT[ci][ci]=sqrt(ranchi(1,1,n-ci));
for(cj=0;cj<ci;++cj){mT[ci][cj]=rann(1,1);}
}
mA=mT*mT'; // A=TT' ~W(n,I)
mL=choleski(mS); // S=LL'(Choleski decomposition)
mX=mL*mA*mL'; // X=LAL'~W(n,S)
return(mX);
}

main()
{
	println("IV model (based on Rossi et al 2005) with No-cross products. logA = FD");
	
	//Sample size
	decl N = 97947;
	decl M = 223;
	decl k1 = 4;
	decl k2 = (29-1)+1;
	decl kalp = 2;
	decl kiv = 1;
	
	//data
	decl data = loadmat("data0.csv");
	
	decl Ni = loadmat( "data1.csv" );
	decl DmSan = loadmat("data2.csv" );//Industry dummy
	decl DmSho = loadmat("data3.csv" );//Ownership dummy
	decl DmChi = loadmat("data4.csv" );//Region dummy
	
	decl W = loadmat( "Wrin.csv" );//Spatial Weight Matrix
	decl lam, lam1, lam2, P;
	eigen(W, &lam, &P);
	lam1 = 1.0 / min(lam);
	lam2 = 1.0 / max(lam);
	
	decl Y = data[][0];
	decl K = data[][1];
	decl L = data[][2];
	decl h = 18*(data[][3] ./ L) + 16*(data[][4] ./ L)
		   + 14*(data[][5] ./ L) + 12*(data[][6] ./ L) + 9*(data[][7] ./ L);
	decl FN = data[][8] ./ data[][10];
	decl ED = data[][11];
	decl IS = data[][12] ./ (data[][12] + Y);
	decl UP = loadmat("UrbanPopDen.csv");
	decl IV = data[][13] ./  data[][10];
	decl MO = data[][14];
	
	decl y = log(Y);
	decl k = log(K);
	decl l = log(L);
	decl fn = log(FN);
	decl is = log(IS);
	decl up = log(UP);
	decl iv = log(IV + 1);
	decl mo = log(MO);
	
	//Variables X and q
	decl X = h ~ ED ~ k ~ l;
	decl q = iv; // N�~kiv
	
	//Variable G: 1 ~ up M * 2
	decl G = ones(M, 1) ~ up;
	
	//Variable D2: phi. Dummy
	decl D2 = DmSan[][1:] ~ DmSho[][2];
	
	//Variable D3: alp. Region Dummy (N * M)
	decl N1 = ones(N, 1);
	decl D3 =  new array[M];
	decl Nic = 0 | Ni[][2];
	for (decl i = 0; i < M; ++i){
		D3[i] = N1[Nic[i]:Nic[i+1]-1];
	}
	
	//Variable Xt: fn X D2, B_k �~ 1
	decl Xt = fn ~ X ~ D2;
	decl XD2 = X ~ D2;
	decl B_k = 1 + k1 + k2;
	
	//Variable Z: q X D2, G_k �~ 1
	decl Z = q ~ X ~ D2;
	decl G_k = kiv + k1 + k2;
	
	//Prior
	//Beta: (1+k1+k2) * 1, betaA beta phi
		decl pEBeta = zeros(B_k, 1);
		decl pVBeta = 100;
		decl I_pVBeta = invert(pVBeta);
	
	//Gamm: (kiv+k1+k2) * 1, ga0 ga1 ga2
		decl pEGamm = zeros(G_k, 1);
		decl pVGamm = 100*unit(G_k);
		decl I_pVGamm = invert(pVGamm);
	
	//Sigma: 2 * 2
		decl pESigma = 2;
		decl pVSigma = unit(2);
		decl I_pVSigma = invert(pVSigma);
	
	//mu_ga:
		decl pEmu_ga = 0;
		decl pVmu_ga = 100;
		decl I_pVmu_ga = invert(pVmu_ga);
		
	//tau2_ga:
		decl pnu_ga = 2*0.001;
		decl pom_ga = 2*0.001;
	
	//mu: kalp�~1
	decl pEmu = zeros(kalp, 1);
	decl pVmu = 100*unit(kalp);
	decl I_pVmu = invert(pVmu);
	
	//tau2
		decl pnu = 2*0.001;
		decl pom = 2*0.001;
		
	
	//MCMC
	decl cburn   = 50000, cn = 350000;
	decl caccept = 0;
	
	decl vEBeta = zeros(cn, B_k);//betaA ~ beta' ~ phi'
	decl vEGamm = zeros(cn, G_k);//ga0' ~ ga1' ~ ga2
	decl vESigma = zeros(cn, 4);
	decl vEalp = zeros(cn, M);
	decl vEga3 = zeros(cn, M);
	decl vEothers = zeros(cn, 6);//mu0, mu1, rho, tau2, mu_ga, tau2_ga
	decl vLL = zeros(cn, 1);
	
	//initial value
	decl Beta = 0.1 | zeros((B_k-1), 1);
	decl Gamm = zeros(G_k, 1);
	decl alp = zeros(M, 1);
	decl ga3 = zeros(M, 1);
	
	decl Sigma = unit(2);
	decl I_Sigma = invert(Sigma);
	decl MbetaA = <1, 0; 1, 0>;
	     MbetaA[1][1] = 1/Beta[0];
	decl ASA = MbetaA*Sigma*MbetaA';
	decl I_ASA = invert(ASA);
	decl U_ASA = choleski(ASA);
	decl sumI_ASA = I_ASA[0][0] + 2*I_ASA[0][1] + I_ASA[1][1];
	decl sig_a = Sigma[0][1] / Sigma[0][0];
	decl sig_b = Sigma[1][1] - Sigma[0][1]^2 / Sigma[0][0];
	decl sqsig_b = sqrt(sig_b);
	
	decl mu = zeros(kalp, 1);
	decl rho = 0;
	decl tau2 = 1;
	decl mu_ga = 0;
	decl tau2_ga = 1;
	
	decl UM = unit(M);
	decl ON = ones(N, 1);
	decl OM = ones(M, 1);
	decl WW = W'*W;
	decl XtXt = Xt'*Xt;
	decl D3D3v = zeros(M, 1);
	decl D3fn = zeros(M, 1);
	for (decl i = 0; i < M; ++i)
	{
		D3D3v[i] = D3[i]'*D3[i];
		D3fn[i] = D3[i]'*fn[Nic[i]:Nic[i+1]-1];
	}
	decl D3D3 = diag(D3D3v);
	decl GG = G'*G;
	decl ZZ = Z'Z;
	decl Zfn = Z'*fn;
	
	decl XtBe = Xt*Beta;
	decl ZGa = Z*Gamm;
	
	decl D3alp = zeros(N, 1);
	decl D3ga3 = zeros(N, 1);
	for (decl i = 0; i < M; ++i)
	{
		D3alp[Nic[i]:Nic[i+1]-1] = D3[i]*alp[i];
		D3ga3[Nic[i]:Nic[i+1]-1] = D3[i]*ga3[i];
	}
	
	decl XtBA =  XD2*Beta[1:] + D3alp;
	decl err1 = fn - D3ga3 - ZGa;
	decl err2 = y  - D3alp - XtBe;
	
	decl time0 = timer();
	println("\nStart time of MCMC: ", time());
	for (decl i = -cburn; i < cn; ++i)
	{	
		//print("%6.0f",i);
		//=== Beta ===
		decl yt1 =  (y - D3alp - sig_a*err1) / sqsig_b;
		decl Xt1 = Xt / sqsig_b;
		decl VBeta = invert( Xt1'Xt1 + I_pVBeta);
		decl EBeta = VBeta*( Xt1'*yt1 + I_pVBeta*pEBeta );
		Beta = EBeta + choleski(VBeta)*rann(B_k, 1);
		XtBe = Xt*Beta;
		XtBA = XD2*Beta[1:] + D3alp;
		MbetaA[1][1] = 1 / Beta[0];
		ASA = MbetaA*Sigma*MbetaA';
		I_ASA = invert(ASA);
		sumI_ASA = I_ASA[0][0] + 2*I_ASA[0][1] + I_ASA[1][1];
				
		//=== Gamm ===
		decl yt2 = (y - XtBA) / Beta[0] - D3ga3;
		decl Zfn1 = (I_ASA[0][0]+I_ASA[1][0])*Zfn;
		decl Zyt2 = (I_ASA[0][1]+I_ASA[1][1])*Z'*yt2;
		decl VGamm = invert( sumI_ASA*ZZ + I_pVGamm );
		decl EGamm = VGamm * ( Zfn1 + Zyt2 + I_pVGamm*pEGamm );
		Gamm = EGamm + choleski(VGamm)*rann(G_k, 1);
		ZGa = Z*Gamm;
		err1 = fn - D3ga3 - ZGa;
		
		//=== alp ===
		decl yt3 = (y - XtBe - sig_a*err1) / sqsig_b;
		decl Q = UM - rho*W;
		decl D3yt3 = zeros(M, 1);
		for (decl j = 0; j < M; ++j)
		{
			D3yt3[j] = (D3[j] / sqsig_b)' * yt3[Nic[j]:Nic[j+1]-1];
		}
		decl Valp = invert( D3D3 / sig_b + Q'*Q / tau2 );
		decl Ealp = Valp * ( D3yt3 + Q'*G*mu / tau2 );
		alp = Ealp + choleski(Valp)*rann(M, 1);
		for (decl i = 0; i < M; ++i)
		{
			D3alp[Nic[i]:Nic[i+1]-1] = D3[i]*alp[i];
		}
		XtBA = XD2*Beta[1:] + D3alp;
		
		//=== ga3 ===
		decl yt4 = ((y - XtBA) / Beta[0]) - ZGa;
		decl D3yt4_tmp = zeros(M, 1);
		for (decl i = 0; i < M; ++i)
		{
			D3yt4_tmp[i] = D3[i]'*yt4[Nic[i]:Nic[i+1]-1];
		}
		decl D3fn1 = (I_ASA[0][0]+I_ASA[1][0])*D3fn;
		decl D3yt4 = (I_ASA[0][1]+I_ASA[1][1])*D3yt4_tmp;
		decl Vga3 = invert( sumI_ASA*D3D3 + 1/tau2_ga );
		decl Ega3 = Vga3 * ( D3fn1 + D3yt4 + OM*mu_ga / tau2_ga );
		ga3 = Ega3 + choleski(Vga3)*rann(M, 1);
		for (decl i = 0; i < M; ++i)
		{
			D3ga3[Nic[i]:Nic[i+1]-1] = D3[i]*ga3[i];
		}
				
		//=== Sigma ===
		err1 = fn - D3ga3 - ZGa;
		err2 = y  - D3alp - XtBe;
		decl V = err1 ~ err2;
		decl VSigma = invert( V'*V + I_pVSigma );
		decl ESigma = N + pESigma;
		decl Sigma_p = fMyWishart(ESigma, VSigma);
		if (Sigma_p == 0) {
			println(i, "_Decomposition failed");
		}
		else {
			Sigma = invert(Sigma_p);
			I_Sigma = invert(Sigma);
			ASA = MbetaA*Sigma*MbetaA';
			I_ASA = invert(ASA);
			sumI_ASA = I_ASA[0][0] + 2*I_ASA[0][1] + I_ASA[1][1];
			sig_a = Sigma[0][1] / Sigma[0][0];
			sig_b = Sigma[1][1] - Sigma[0][1]^2 / Sigma[0][0];
			sqsig_b = sqrt(sig_b);
		}
		
		//=== mu ===
		decl Vmu = invert( GG / tau2 + I_pVmu );
		decl Emu = Vmu * ( G'*Q*alp / tau2 + I_pVmu*pEmu);
		mu = Emu + choleski(Vmu)*rann(2,1);
		
		//=== tau ===
		decl EEa = (Q*alp - G*mu)'*(Q*alp - G*mu);
		decl Enu = 0.5 * (M + pnu);
		decl Eom = 0.5 * (EEa + pom);
		tau2 = 1.0 / rangamma(1,1,Enu, Eom);
		
		//=== rho ===
		decl Vrho = invert( alp'*WW*alp / tau2 );
		decl Erho = Vrho * ( alp'*W'*(alp-G*mu) / tau2 );
		decl pRho;
		do {
			pRho  = Erho + sqrt(Vrho) * rann(1,1);
		} while (pRho <= lam1 || pRho >= lam2);
		decl IR   = UM - rho * W;
		decl pIR  = UM - pRho * W;
		decl num  = determinant( pIR );
		decl den  = determinant( IR  );
		decl ALP  = num / den;
		if (ranu(1,1) < ALP){
			rho = pRho;
			if(i >= 0){
				caccept += 1;
			}
		}
				
		//=== mu_ga ===
		decl Vmu_ga = invert( OM'*OM / tau2_ga + I_pVmu_ga );
		decl Emu_ga = Vmu_ga * ( OM'*ga3 / tau2_ga + I_pVmu_ga*pEmu_ga );
		mu_ga = Emu_ga + sqrt(Vmu_ga)*rann(1,1);
		
		//=== tau2_ga ===
		decl EEga = (ga3 - OM*mu_ga)'*(ga3 - OM*mu_ga);
		decl Enu_ga = 0.5 * (M + pnu_ga);
		decl Vnu_ga = 0.5 * (EEga + pom_ga);
		tau2_ga = 1.0 / rangamma(1,1,Enu_ga, Vnu_ga);
		
		if (i >= 0){
			vEGamm[i][] = Gamm';
			vEBeta[i][] = Beta';
			vEothers[i][] = mu'~rho~tau2~mu_ga~tau2_ga;
			vEalp[i][] = alp';
			vEga3[i][] = ga3';
			vESigma[i][] = Sigma[0][0] ~ Sigma[0][1] ~ Sigma[1][0] ~ Sigma[1][1];
			decl ESE = I_Sigma[0][0]*err1'*err1 + I_Sigma[0][1]*err1'*err2 + I_Sigma[1][1]*err2'*err2;
			decl LL = -0.5*N*log(2*M_PI) -0.5*N*log(determinant(Sigma)) - 0.5*ESE;
			vLL[i] = LL;
		}
	}	
	
	//=== ReportMCMC ===
	
	decl out1 = new ReportMCMC(vEBeta);
		 out1.SetOutfileName("Beta");
		 out1.Report();
	delete out1;
		
	decl out2 = new ReportMCMC(vEGamm);
		 out2.SetOutfileName("Gamm");
		 out2.Report();
	delete out2;
	
	decl out3 = new ReportMCMC(vEothers);
		 out3.SetOutfileName("Others");
		 out3.Report();
	delete out3;
	
	decl out5 = new ReportMCMC(vESigma);//Sigma[0][0] ~ Sigma[0][1] ~ Sigma[1][0] ~ Sigma[1][1]
		 out5.SetOutfileName("Sigma");
		 out5.Report();
	delete out5;
	
	println( "%12.5f", "%c", {"1/lam_min","1/lam_max","caccept"}, lam1~lam2~(caccept/cn));

	decl time1 = timer();
	println("\nTime: ", timespan(time1,time0));

}