#fm_transition.R
#
#This file computes the baseline fully mobile transition in
#"Trapped Factors and China's Impact on Growth,"
#Bloom, Nicholas, Romer, Paul, Terry, Stephen J., and Van Reenen, John (2013).
#
#If you have questions about the code, please contact Stephen Terry at sterry@stanford.edu.
#
#September 27, 2013

rm(list=ls())
library(pso)
library(matlab)
set.seed(1)

assign('numyrs',10)
assign('growth.target',((1+0.02)^numyrs) - 1)
assign('IoverY.open.target',0.039)
assign('IoverY.fully.mobile.target',0.07)
assign('sigma',1); #flow utility u^(1-sigma) / (1-sigma)
assign('alpha',2/3); #share on human capital
assign('human.cap.ratio',2.96); # from Barro-Lee Data
assign('beta',(1/1.02)^numyrs);  #discount rate
assign('H',.0185);  #northern human capital endowment
assign('Hstar',H*human.cap.ratio);  #southern human capital endowment
assign('rho',0.5); #parameter of innovation production function
assign('gamma',1/rho); #parameter of innovation cost function
assign('phi',.1725);  #proportion of goods off of trade restriction, initialization only
shock = 0.075; # for initialization of phi.prime only 
assign('phi.prime',phi*shock); #post-shock trade parameter, initialization only
assign('eta',1); #eta=1 is fully internal marginal costs, eta=0 is external
assign('n',2); #number of firms (must be even for cross-sectional experiment to work properly)

#some useful constants, same name as in paper
assign('OMEGA',alpha* ( 1 - alpha ) ^ ( ( 2 - alpha ) / alpha ))
assign('PSI',( 1 - alpha ) ^ ( ( alpha - 1 ) / ( 2 - alpha ) ))   

#now, setup the transition path solution
assign('T',7); #of periods assumed for full transition
assign('A1',1); #assumed starting number of varieties
errortol=10^(-15)

calibration = function(IN,eval.dum){
H=IN[1]
phi=IN[2]
shock=IN[3]

assign('phi.prime',phi*shock)
assign('Hstar',human.cap.ratio*H)

##################
##################BGP growth rates
##################

g.open.fun = function(g) {
  

    LHS = g ^ ( gamma - 1) 
    RHS = ( beta ^ ( 1 / alpha ) ) * OMEGA 
    RHS = RHS * ( H + ( PSI ^ ( 1 / alpha ) ) * Hstar * ( ( ( phi * H ) / ( g * Hstar ) ) ^ ( 1 / ( 2 - alpha ) ) ) )
    
    RHS = RHS * ( 1 + g )^(- sigma / alpha )
  
    obj = (LHS - RHS)^2
  
  
}

#use g.open.fun to compute BGP growth rate for phi
g.open.guess=.02
g.open.list = optim(g.open.guess,g.open.fun,gr=NULL,method="Brent",lower=0,upper=1)
g.open = g.open.list$par
assign('g.open',g.open)
q.open = PSI * ( ( phi * H ) / ( g.open * Hstar ) ) ^ ( alpha / (2 - alpha ) ) 
assign('q.open',q.open)
intrate.open = (beta^-1) * (1+g.open)^sigma - 1
assign('intrate.open',intrate.open)
eqbm.check.phi = (q.open <=1)

#use g.open.fun to compute BGP growth rate for phi.prime
phi.orig=phi
assign('phi',phi.prime)
g.shock.mobile.list=optim(g.open,g.open.fun,gr=NULL,method="Brent",lower=0,upper=5)
g.shock.mobile=g.shock.mobile.list$par
assign('g.shock.mobile',g.shock.mobile)
q.shock.mobile = PSI * ( ( phi * H ) / ( g.shock.mobile * Hstar ) ) ^ ( alpha / (2 - alpha ) ) 
intrate.shock.mobile = (beta^-1) * (1+g.shock.mobile)^sigma - 1
assign('intrate.shock.mobile',intrate.shock.mobile)
assign('q.shock.mobile',q.shock.mobile)
assign('phi',phi.orig)
eqbm.check.phi.prime = (q.shock.mobile <=1)

##################
##################calibration moments for imports/gross output
##################

#pre-shock
Iphi = phi; Iphi.min1 = Iphi/(1+g.open); pIphi = q.open * (1+intrate.open); xIphi = (1-alpha)^(1/alpha) * H * (pIphi^(-1/alpha));
Rphi = (1-phi); Rphi.min1 = Rphi/(1+g.open); pRphi = (1+intrate.open); xRphi = (1-alpha)^(1/alpha) * H * (pRphi^(-1/alpha));
Mphi = g.open; Mphi.min1 = Mphi/(1+g.open); pMphi = (1+intrate.open)/(1-alpha); xMphi = (1-alpha)^(1/alpha) * H * (pMphi^(-1/alpha));
Yphi = (H^alpha) * (  Iphi.min1 * (xIphi^(1-alpha)) + Rphi.min1 * (xRphi^(1-alpha)) + Mphi.min1 * (xMphi^(1-alpha))         )
Ygphi = Yphi + Mphi.min1*pMphi*xMphi + Iphi.min1*pIphi*xIphi + Rphi.min1*pRphi*xRphi
IoverY.open = (Iphi.min1*pIphi*xIphi)/(Yphi)

#post-shock
Iphi.prime = phi.prime; Iphi.prime.min1 = Iphi.prime/(1+g.shock.mobile); pIphi.prime = q.shock.mobile * (1+intrate.shock.mobile); xIphi.prime = (1-alpha)^(1/alpha) * H * (pIphi.prime^(-1/alpha));
Rphi.prime = (1-phi.prime); Rphi.prime.min1 = Rphi.prime/(1+g.shock.mobile); pRphi.prime = (1+intrate.shock.mobile); xRphi.prime = (1-alpha)^(1/alpha) * H * (pRphi.prime^(-1/alpha));
Mphi.prime = g.shock.mobile; Mphi.prime.min1 = Mphi.prime/(1+g.shock.mobile); pMphi.prime = (1+intrate.shock.mobile)/(1-alpha); xMphi.prime = (1-alpha)^(1/alpha) * H * (pMphi.prime^(-1/alpha));
Yphi.prime = (H^alpha) * (  Iphi.prime.min1 * (xIphi.prime^(1-alpha)) + Rphi.prime.min1 * (xRphi.prime^(1-alpha)) + Mphi.prime.min1 * (xMphi.prime^(1-alpha))         )
Ygphi.prime = Yphi.prime + Mphi.prime.min1*pMphi.prime*xMphi.prime + Iphi.prime.min1*pIphi.prime*xIphi.prime + Rphi.prime.min1*pRphi.prime*xRphi.prime
IoverY.fully.mobile = (Iphi.prime.min1*pIphi.prime*xIphi.prime)/(Yphi.prime)

error1 = (growth.target-g.open)/growth.target
error2 = (IoverY.open.target-IoverY.open)/IoverY.open.target
error3 = (IoverY.fully.mobile.target-IoverY.fully.mobile)/IoverY.fully.mobile.target

error = error1^2 + error2^2 + error3^2

cat("H = ",H,"phi = ",phi,"phi.prime = ",phi.prime,"error = ",error,"\n")

#eqbm check for phi BGP
eqbm.check.phi = (q.open<=1)

#eqbm check for phi.prime BGP
eqbm.check.phi.prime = (q.shock.mobile<=1)

eqbm.check.BGP = eqbm.check.phi * eqbm.check.phi.prime

if (eval.dum==0){return(error)}

eval.vec = c(g.open,q.open,intrate.open,g.shock.mobile,q.shock.mobile,intrate.shock.mobile,eqbm.check.BGP)

if (eval.dum==1){return(eval.vec)}

}

#perform the calibration exercise, using the BGP
guess=c(.035,.0795,1.28); #H,phi,phi.prime/phi
lb=c(.001,.01,1); ub=c(5,.2,4)
BGPcalibration=psoptim(guess,calibration,lower=lb,upper=ub,eval.dum=0,control=list(abstol=1e-18,maxit=2000))

#extract the calibrated parameters
H = BGPcalibration$par[1]; assign('H',H); Hstar = H * human.cap.ratio; assign('Hstar',Hstar)
phi = BGPcalibration$par[2]; assign('phi',phi)
phi.prime = BGPcalibration$par[2]*BGPcalibration$par[3]; assign('phi.prime',phi.prime)

#extract the growth, exchange, and interest rates
BGPvec = calibration(BGPcalibration$par,eval.dum=1)

assign('g.phi',BGPvec[1])
assign('q.phi',BGPvec[2])
assign('r.phi',BGPvec[3])
assign('g.phi.prime',BGPvec[4])
assign('q.phi.prime',BGPvec[5])
assign('r.phi.prime',BGPvec[6])

guessIN = zeros(3*(T-1),1)
guessIN[1:(T-1)] = (r.phi+r.phi.prime)/2
guessIN[(T-1)+1:(T-1)]=(r.phi+r.phi.prime)/2
guessIN[2*(T-1)+1:(T-1)]=(q.phi+q.phi.prime)/2
lb = guessIN*.8
ub = guessIN * 1.2

f.transition = function(IN){

#initialize vectors of allocations
Y = array(0,dim=c(T+1,1)); #northern output vector
Ystar = Y;  A = Y;  M = Y;  R = Y;  I = Y; u = Y; ustar = Y;
xM = Y; xMstar = Y; xR = Y; xRstar = Y; xI = Y; xIstar = Y; Z = Y
r = Y; rstar = Y; q = Y; g = Y;

#extract prices from input vector IN
r[2:T] = IN[1:(T-1)]; r[1]=r.phi; r[T+1]=r.phi.prime;
rstar[2:T] = IN[(T-1)+1:(T-1)]; rstar[1]=r.phi; rstar[T+1]=r.phi.prime;
q[2:T] = IN[2*(T-1)+1:(T-1)]; q[1] = q.phi; q[T+1]=q.phi.prime;

#feed in known values for period 1
A[1]=A1;
M[1]=g.phi/(1+g.phi) * A1;
R[1]=(1-phi)/(1+g.phi) * A1;
I[1]=phi/(1+g.phi)*A1;

#given prices, compute quantities for each period
xM = (1-alpha)^(2/alpha) * (1+r)^(-1/alpha) * H
xMstar = (1-alpha)^(2/alpha) * (1+r)^(-1/alpha) * q^(1/alpha) *  Hstar
xR = (1-alpha)^(1/alpha) * (1+r)^(-1/alpha) * H

xRstar = (1-alpha)^(1/alpha) * (1+rstar)^(-1/alpha) * Hstar
xI = (1-alpha)^(1/alpha) * (1+rstar)^(-1/alpha) * q^(-1/alpha) * H
xIstar = (1-alpha)^(1/alpha) * (1+rstar)^(-1/alpha) * Hstar

#get implied growth rates
g = ( OMEGA * (1 + r)^(-1/alpha) * (H + Hstar * q^(1/alpha)) )^(1/(gamma-1))

#get implied measures
for (t in 2:(T+1)) { 
A[t] = (1+g[t])*A[t-1] 
R[t] = (1-phi.prime) * A[t-1]
I[t] = phi.prime * A[t-1]
M[t] = g[t] * A[t-1]
}

#get implied innovation costs
Z[1:T] = g[2:(T+1)]^gamma * A[1:T] / (eta * (gamma-1) + 1)

#get implied output
Y =H^alpha * (M * xM^(1-alpha) + R * xR^(1-alpha) + I * xI^(1-alpha) )
Ystar = Hstar^alpha * (M * xMstar^(1-alpha) + R * xRstar^(1-alpha) + I * xIstar^(1-alpha) )

#get implied consumption in 1:T
u[1:T] = Y[1:T] - M[2:(T+1)] * (xM[2:(T+1)] + xMstar[2:(T+1)]) - R[2:(T+1)] * xR[2:(T+1)] - Z[1:T]
ustar[1:T] = Ystar[1:T] - R[2:(T+1)] * xRstar[2:(T+1)] - I[2:(T+1)] * (xI[2:(T+1)] + xIstar[2:(T+1)])

#Northern Eulers
LHSN = (u[2:T]/u[1:(T-1)])^sigma
RHSN = beta * (1+r[2:T])

#Southern Eulers
LHSS = (ustar[2:T]/ustar[1:(T-1)])^sigma
RHSS = beta * (1+rstar[2:T])

#BT
LHSBT = q[2:T]
RHSBT = ((phi.prime * H)/(g[2:T]*Hstar))^(alpha/(2-alpha)) * PSI * ((1+r[2:T])/(1+rstar[2:T]) )^((1-alpha)/(2-alpha))

error = sum(((LHSN - RHSN)/RHSN)^2) + sum(((LHSS - RHSS)/RHSS)^2) + sum(((LHSBT - RHSBT)/RHSBT)^2) 
cat(error,'\n')
return(error)
}

transition=psoptim(guessIN,f.transition,lower=lb,upper=ub,control=list(abstol=errortol,maxit=100000))

r = zeros(T+1,1)
rstar = zeros(T+1,1)
q = zeros(T+1,1)
r[2:T] = transition$par[1:(T-1)]; r[1]=r.phi; r[T+1]=r.phi.prime;
rstar[2:T] = transition$par[(T-1)+1:(T-1)]; rstar[1]=r.phi; rstar[T+1]=r.phi.prime;
q[2:T] = transition$par[2*(T-1)+1:(T-1)]; q[1] = q.phi; q[T+1]=q.phi.prime;

Y = array(0,dim=c(T+1,1)); #northern output vector
Ystar = Y;  A = Y;  M = Y;  R = Y;  I = Y; u = Y; ustar = Y;
xM = Y; xMstar = Y; xR = Y; xRstar = Y; xI = Y; xIstar = Y; Z = Y

#feed in known values for period 1
A[1]=A1;
M[1]=g.phi/(1+g.phi) * A1;
R[1]=(1-phi)/(1+g.phi) * A1;
I[1]=phi/(1+g.phi)*A1;

#given prices, compute quantities for each period
xM = (1-alpha)^(2/alpha) * (1+r)^(-1/alpha) * H
xMstar = (1-alpha)^(2/alpha) * (1+r)^(-1/alpha) * q^(1/alpha) *  Hstar
xR = (1-alpha)^(1/alpha) * (1+r)^(-1/alpha) * H

xRstar = (1-alpha)^(1/alpha) * (1+rstar)^(-1/alpha) * Hstar
xI = (1-alpha)^(1/alpha) * (1+rstar)^(-1/alpha) * q^(-1/alpha) * H
xIstar = (1-alpha)^(1/alpha) * (1+rstar)^(-1/alpha) * Hstar

#get implied growth rates
g = ( OMEGA * (1 + r)^(-1/alpha) * (H + Hstar * q^(1/alpha)) )^(1/(gamma-1))

#get implied measures
for (t in 2:(T+1)) { 
  A[t] = (1+g[t])*A[t-1] 
  R[t] = (1-phi.prime) * A[t-1]
  I[t] = phi.prime * A[t-1]
  M[t] = g[t] * A[t-1]
}

#get implied innovation costs
Z = g[2:(T+1)]^gamma * A[1:T] / (eta * (gamma-1) + 1)

#get implied output
Y =H^alpha * (M * xM^(1-alpha) + R * xR^(1-alpha) + I * xI^(1-alpha) )
Ystar = Hstar^alpha * (M * xMstar^(1-alpha) + R * xRstar^(1-alpha) + I * xIstar^(1-alpha) )

#get implied consumption in 1:T
u[1:T] = Y[1:T] - M[2:(T+1)] * (xM[2:(T+1)] + xMstar[2:(T+1)]) - R[2:(T+1)] * xR[2:(T+1)] - Z[1:T]
ustar[1:T] = Ystar[1:T] - R[2:(T+1)] * xRstar[2:(T+1)] - I[2:(T+1)] * (xI[2:(T+1)] + xIstar[2:(T+1)])

#create a series of the lower bound prices for the I goods for Northern firms

pI_North = 1+r
pI_South = q*(1+rstar)

EQBM_CHECK = sum((pI_North < pI_South)==TRUE)==0
cat('IT IS AN I GOODS EQBM = ',EQBM_CHECK,'\n')

#summary spreadsheet - for trapped factors allocations
M1vec = 0.5*M; 
M2vec = 0.5*M; 

R1vec = 0.5*R; 
R2vec = 0.5*R; 

mu1vec = ones(T+1,1); 
mu2vec = ones(T+1,1); 

u[T+1]=(1+g[T+1])*u[T];
ustar[T+1]=(1+g[T+1])*ustar[T];

xM1vec = xM; 
xM2vec = xM; 

xMstar1vec= xMstar; 
xMstar2vec= xMstar; 

xR1vec = xR; 
xR2vec = xR; 

phivec=ones(T+1,1); phivec[1]=phi; phivec[2:(T+1)]=phi.prime

numvars = 28
allocations = zeros(T+1,numvars);
list=allocations.labels=list()
allocations[,1]=r; allocations.labels = c(allocations.labels,'r') 
allocations[,2]=rstar; allocations.labels = c(allocations.labels,'rstar')
allocations[,3]=q; allocations.labels = c(allocations.labels,'q')
allocations[,4]=g; allocations.labels = c(allocations.labels,'g')
allocations[,5]=mu1vec; allocations.labels = c(allocations.labels,'mu1')
allocations[,6]=mu2vec; allocations.labels = c(allocations.labels,'mu2')
allocations[,7]=u; allocations.labels = c(allocations.labels,'u')
allocations[,8]=ustar; allocations.labels = c(allocations.labels,'ustar')
allocations[,9]=Y; allocations.labels = c(allocations.labels,'Y')
allocations[,10]=Ystar; allocations.labels = c(allocations.labels,'Ystar')
allocations[,11]=A; allocations.labels = c(allocations.labels,'A')
allocations[,12]=M; allocations.labels = c(allocations.labels,'M')
allocations[,13]=I; allocations.labels = c(allocations.labels,'I')
allocations[,14]=R; allocations.labels = c(allocations.labels,'R')
allocations[,15]=M1vec; allocations.labels = c(allocations.labels,'M1')
allocations[,16]=M2vec; allocations.labels = c(allocations.labels,'M2')
allocations[,17]=R1vec; allocations.labels = c(allocations.labels,'R1')
allocations[,18]=R2vec; allocations.labels = c(allocations.labels,'R2')
allocations[,19]=xM1vec; allocations.labels = c(allocations.labels,'xM1')
allocations[,20]=xM2vec; allocations.labels = c(allocations.labels,'xM2')
allocations[,21]=xMstar1vec; allocations.labels = c(allocations.labels,'xMstar1')
allocations[,22]=xMstar2vec; allocations.labels = c(allocations.labels,'xMstar2')
allocations[,23]=xR1vec; allocations.labels = c(allocations.labels,'xR1')
allocations[,24]=xR2vec; allocations.labels = c(allocations.labels,'xR2')
allocations[,25]=xRstar; allocations.labels = c(allocations.labels,'xRstar')
allocations[,26]=xI; allocations.labels = c(allocations.labels,'xI')
allocations[,27]=xIstar; allocations.labels = c(allocations.labels,'xIstar')
allocations[,28]=phivec; allocations.labels = c(allocations.labels,'phivec')

allocations.frame=data.frame(allocations)
colnames(allocations.frame)=allocations.labels

parametersvec = c(alpha,beta,sigma,rho,eta,H,Hstar,n,phi,phi.prime,IoverY.open.target,IoverY.fully.mobile.target)
parameterslabels =  c('alpha','beta','sigma','rho','eta','H','Hstar','n','phi','phi.prime','IoverY.open.target','IoverY.fully.mobile.target')
parametersvec = data.frame(t(parametersvec))
colnames(parametersvec)=parameterslabels

write.csv(allocations.frame,file="allocations.csv")
write.csv(parametersvec,file="parameters.csv")

numgraphvars = 13
graphout.fm = zeros(T+1,numgraphvars)
graphlabels = c('r','rstar','q','g','phi','u','ustar','Y','Ystar','A','M','M1','M2')

graphout.fm[,1] = r
graphout.fm[,2] = rstar
graphout.fm[,3] = q
graphout.fm[,4] = g
graphout.fm[,5] = phivec
graphout.fm[,6] = u
graphout.fm[,7] = ustar
graphout.fm[,8] = Y
graphout.fm[,9] = Ystar
graphout.fm[,10] = A
graphout.fm[,11] = M
graphout.fm[,12] = M1vec
graphout.fm[,13] = M2vec

graphout.fm = data.frame(graphout.fm)
colnames(graphout.fm) = graphlabels

save(file="fm_transition.rda",list=ls())
save(file="graphout_fm.rda",graphout.fm)