#  betamatch.R
#  Computes steady state paths of consumption, assets, etc of OG model 
#  with annual time period and arbitrary inputs for mortality rates.  
#  K/Y given, hence interest rate, so the discount factor beta is the unknown.  
#  Variable definitions 
#    a = asset holdings by age 
#    c = consumption by age 
#    w = wage (scalar)
#    R = gross interest rate (scalar)
#    h = bequest income (scalar for now)
#    K = capital stock (= aggregate assets)
#    Y = GDP (= aggregate consumption + investment)
#  Demographics 
#    s = survival prob by age:  from age i to i+1 
#    I0 = initial age (20, nothing happens till I0+1)
#    I = terminal age (101)
#    IR = retirement age 
#  Parameters 
#    beta = discount factor 
#    sigma = curvature parameter in prefs:  u(c) = c^{1-sigma}/(1-sigma)
#    alpha = capital share (Cobb-Douglas production)
#    delta = depreciation rate (annual)
#  Written by:  Espen Henriksen
#  Comments added by Dave Backus, September 2012 
# ------------------------------------------------------------------------------
rm(list=ls())

# remember to change \ to /
#dir <- "/Users/espen/Dropbox/research/global_economy/model101/"
dir <- "C:/Users/dbackus/Documents/Papers/BCH/model/steady_state"
setwd(dir) 

# 1. Set basic parameters 
# ------------------------------------------------------------------------------

# Preferences
sigma <- 0.5

# Technology
alpha <- 0.33
delta <- .06

# Demographics
I0 <- 20
I  <- 100
IR <- 65

# K/Y ratio and interest rate (Cobb-Douglas net MPK)
KYtarget <- 3
R = 1 + alpha * 1/KYtarget - delta

# Closing the model / social security system
# sssystem <- 1  "black hole"
# sssystem <- 2  redistributed to survivors in same cohort
# sssystem <- 3  accidential bequests distributed evenly
sssystem <- 2


# 2. Functions 
# ------------------------------------------------------------------------------

# update net asset position a:  given starting points, compute whole path  
# based on second-order difference equation for net asset position 
law.of.motion <- function(a,R,w,h,beta,sigma,s,I0,I) {
  for(i in c((I0+2):I)){
    a[i+1] <- a[i]*R + w[i] + h[i] - (1/(beta*s[i-1]*R))^(-sigma)*(a[i-1]*R + w[i-1] + h[i-1] - a[i])
  }
  return(a)
}

# compute capital output ratios
KY <- function(alpha,cohorts,assets){
  N <- sum(cohorts[21:IR])
  K<-0
  for(i in 1 : I){K <- K + assets[i]*cohorts[i]}
  Y <- K^alpha * N^(1-alpha)
  return(K/Y)
}

# compute bequests in different social security systems
# sssystem <- 1  "black hole"
# sssystem <- 2  redistributed to survivors in same cohort
# sssystem <- 3  accidential bequests distributed evenly
bequests <- function(sssystem,I0,I,mortality,assets,R){
  if(sssystem == 1){
    h <- mat.or.vec(I+1,1)
  }else if(sssystem == 2){
    h <- mat.or.vec(I+1,1)
    for(i in I0 : I){
      h[i] <- (1-mortality[i-1])*assets[i-1]*R
    }
  }else if(sssystem == 3){
    h <- mat.or.vec(I+1,1)
    accbeq <- 0
    for(i in I0 : I){
      accbeq <- accbeq + (1-mortality[i-1])*assets[i-1]*R*cohorts[i-1]
    }
    for(i in I0+1 : I){
      h[i] <- accbeq/sum(cohorts[(I0+1):I])
    }
  }
  return(h)
}

# compute asset path by "shooting method"
# start with initial value a[I0+1] = 0 and a guess of a[I0+2]
# adjust guess via bisection until we hit terminal value a[I+1] = 0
compute.asset.holdings <- function(R,w,h,beta,sigma,s,I0,I){
  alow  <- -20
  ahigh <- 15

  a <- numeric(length=I+1)
  a[I0+2] <- alow
  a <- law.of.motion(a,R,w,h,beta,sigma,s,I0,I)
  flow <- a[I+1]

  a <- numeric(length=I+1)
  a[I0+2] <- ahigh
  a <- law.of.motion(a,R,w,h,beta,sigma,s,I0,I)
  fhigh <- a[I+1]
  
  a <- numeric(length=I+1)
  amid <-  (ahigh + alow)/2                       
  a[I0+2] <- amid
  a <- law.of.motion(a,R,w,h,beta,sigma,s,I0,I)                     
  fmid <- a[I+1]

  diff <- 1
  iter <- 1
  
  if(flow*fhigh>0){
      print("here")
      print("sign flow = sign fhigh")
      diff = 1E-6
  }
  
  while (diff > 1E-6) {
    if (fmid*flow < 0) {
      ahigh <- amid
      fhigh <- fmid
    } else {
      alow <- amid
      flow <- fmid  
    }
    a <- numeric(length=21)
    amid <-  (ahigh + alow)/2                       
    a[I0+2] <- amid
    a <- law.of.motion(a,R,w,h,beta,sigma,s,I0,I)
    fmid <- a[I+1]
    
    diff = abs(fmid)

    iter <- iter + 1
    if(iter > 1000){
      diff <- 1E-6
      print("iter > 1000")
      print(a)
    }
  }
  return(a)
}


# compute survival probs??
survival.probabilities <- function(){
  source("MatchLE.R")

  s60<-SurvivalProbs(60)
  s65<-SurvivalProbs(65)
  s70<-SurvivalProbs(70)
  s75<-SurvivalProbs(75)
  s80<-SurvivalProbs(80)
  s85<-SurvivalProbs(85)

  alls <- cbind(s60,s65,s70,s75,s80,s85)
  return(alls)
}

cohort.sizes <- function(alls){
  c60 <- mat.or.vec(101,1)
  c60[1] = alls[1,1]
  c65 <- mat.or.vec(101,1)
  c65[1] = alls[1,2]
  c70 <- mat.or.vec(101,1)
  c70[1] = alls[1,3]
  c75 <- mat.or.vec(101,1)
  c75[1] = alls[1,4]
  c80 <- mat.or.vec(101,1)
  c80[1] = alls[1,5]
  c85 <- mat.or.vec(101,1)
  c85[1] = alls[1,6]
  for(i in 1:100){
    c60[i+1] = c60[i]*alls[i+1,1]
    c65[i+1] = c65[i]*alls[i+1,2]
    c70[i+1] = c70[i]*alls[i+1,3]
    c75[i+1] = c75[i]*alls[i+1,4]  
    c80[i+1] = c80[i]*alls[i+1,5]
    c85[i+1] = c85[i]*alls[i+1,6]
  }
  c60 <- c60/sum(c60)
  c65 <- c65/sum(c65)
  c70 <- c70/sum(c70)
  c75 <- c75/sum(c75)
  c80 <- c80/sum(c80)
  c85 <- c85/sum(c85)
  allc <- cbind(c60,c65,c70,c75,c80,c85)

  plot(alls[1:100,1], xlab="Age",ylab="Conditional survivial probability",lwd=2.0,type="l",col=c(2),axes=F)
  lines(alls[1:100,3],lwd=2.0,col=c(3))
  lines(alls[1:100,5],lwd=2.0,col=c(4))
  axis(1)
  axis(2)
  legend("bottomleft",legend = c("Life expectancy at birth: 60 years" ,"Life expectancy at birth: 70 years","Life expectancy at birth: 80 years"), cex=1.0,lwd=2.0,lty=c(1,1,1),col=c(2:4),bty="n")
  dev.print(device=pdf, "survivalprobs.pdf")

  plot(allc[1:101,1], xlab="Age",ylab="Stationary relative cohort sizes",lwd=2.0,type="l",col=c(2),axes=F)
  lines(allc[1:101,3],lwd=2.0,col=c(3))
  lines(allc[1:101,5],lwd=2.0,col=c(4))
  axis(1)
  axis(2)
  legend("bottomleft",legend = c("Life expectancy at birth: 60 years" ,"Life expectancy at birth: 70 years","Life expectancy at birth: 80 years"), cex=1.0,lwd=2.0,lty=c(1,1,1),col=c(2:4),bty="n")
  dev.print(device=pdf, "cohortsizes.pdf")

  return(allc)
}


# 3. Main program 
# ------------------------------------------------------------------------------

# set wage from marginal product 
w <- mat.or.vec(I,1)
w[(I0+1):IR] <- (1-alpha)*((R-(1-delta))/(alpha))^(alpha/(alpha-1))

# get survival probs and steady state cohort sizes 
alls <- survival.probabilities()
allc <- cohort.sizes(alls)

# define arrays that store steady-state betas, KYs, and Rs 
allbetas <- array(0, dim=c(6,6),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))
allKYs   <- array(0, dim=c(6,6),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))
allRs    <- array(0, dim=c(6,6),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))
allWKs   <- array(0, dim=c(6,6),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))

allassetsKY      <- array(0, dim=c(6,6,I+1),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))
allconsumptionKY <- array(0, dim=c(6,6,I+1),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))

allassetsbeta      <- array(0, dim=c(6,6,I+1),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))
allconsumptionbeta <- array(0, dim=c(6,6,I+1),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))

allassetsKYbeta      <- array(0, dim=c(6,6,I+1),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))
allconsumptionKYbeta <- array(0, dim=c(6,6,I+1),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))


# Solve model for beta / calibrate beta -- given targeted K/Y ratio 
# Loop over differences in life expectancy and associated cohort sizes
for(sctr in 1:6){
  for(cctr in 1:6){

    mortality <- alls[,sctr]
    cohorts <- allc[,cctr]

    # loop for different social security system / ways to (not to) close the model
    # ssystem set earlier
    ssdiff <- 1
    ssiter <-1
    laggedassets <-  mat.or.vec(I+1,1)
    while (ssdiff > 1E-6) {
      if(ssiter > 1){
        h <- (h + bequests(sssystem,I0,I,mortality,assets,R))/2}
      else if(ssiter==1){
        h <- mat.or.vec(I+1,1)
      }
      
      # bisection on beta:  pick to hit steady state K/Y
      betalow <- .905
      betahigh <- 1.01

      KYlow <- NaN
      while(is.nan(KYlow)){
        betalow <- betalow + .005
        assets <- compute.asset.holdings(R,w,h,betalow,sigma,mortality,I0,I)
        KYlow <- KY(alpha,cohorts,assets) - KYtarget
      }

      assets <- compute.asset.holdings(R,w,h,betahigh,sigma,mortality,I0,I)
      KYhigh <-  KY(alpha,cohorts,assets) - KYtarget

      if(KYlow*KYhigh>0){
        print("sign flow = sign fhigh")
        diff = 1E-6
      }

      diff <- 1
      iter <- 0
      while (diff > 1E-6) {
        betamid <- (betahigh + betalow)/2
        assets <- compute.asset.holdings(R,w,h,betamid,sigma,mortality,I0,I)
        KYmid <-  KY(alpha,cohorts,assets) - KYtarget
      
        if (KYmid*KYlow < 0) {
          betahigh <- betamid
          KYhigh <- KYmid
        } else {
          betalow <- betamid
          KYlow <- KYmid  
        }

        diff <- abs(KYmid)
        iter <- iter + 1
        if(iter > 100){
          diff <- 1E-6
          print("iter > 100")
        }
      }
      rm(diff,iter)
      rm(KYlow,KYmid,KYhigh,betalow,betahigh)
      
      ssdiff <- max(abs(assets-laggedassets))
      laggedassets <- assets
      ssiter <- ssiter + 1
    }
    rm(ssdiff,ssiter,laggedassets)
    
    c <- mat.or.vec(I+1,1)
    gs <- mat.or.vec(I+1,1)
    C <- 0
    wN <- 0
    rK <- 0
    GS <- 0
    for(i in 1 : I){
      c[i] <- assets[i]*R + w[i] + h[i] - assets[i+1]
      wN <- wN + w[i]*cohorts[i]
      rK <- rK + (R-(1-delta))*assets[i]*cohorts[i]
      C <- C + c[i]*cohorts[i]
    }
    rm(i)
    
    allbetas[sctr,cctr] <- betamid
    
    # Test of Euler condition
    eulerdiff <- mat.or.vec(I+1,1)
    for(i in I0+1 : I){
      eulerdiff[i] <- betamid*R*mortality[i]*(c[i+1]/c[i])^(-sigma)
    }
    rm(i)

    plot(assets, xlab="Age",ylab="Assets, consumption, wages",lwd=2.0,type="l",col=c(2),axes=F)
    lines(c[1:101],lwd=2.0,col=c(3))
    lines(w[1:101],lwd=2.0,col=c(4))
    lines(h[1:101],lwd=2.0,col=c(5))
    axis(1)
    axis(2)

    # Assign assets and consumption profiles to arrays
    allassetsKY[sctr,cctr,1:I+1] <- assets[1:I+1]
    allconsumptionKY[sctr,cctr,1:I+1] <- c[1:I+1]
    # WARNING (18Oct12): Check dimensions of assets and c
    
    print(allbetas)
    rm(betamid)
  }
}
rm(sctr,cctr)
rm(eulerdiff,rK,wN,GS,C)

######################################################################
# Solve model for K/Y ratio (and real interest rate) -- given beta
# Loop over differences in life expectancy and associated cohort sizes

for(sctr in 1:6){
  for(cctr in 1:6){

    mortality <- alls[,sctr]
    cohorts   <- allc[,cctr]
    beta      <- allbetas[1,1]
    
    #if(sctr >1 | cctr >1){
      # loop for different social security system / ways to (not to) close the model
      # ssystem set earlier
      ssdiff <- 1
      ssiter <-1
      laggedassets <-  mat.or.vec(I+1,1)
      while (ssdiff > 1E-6) {
        if(ssiter > 1){
          h <- (h + bequests(sssystem,I0,I,mortality,assets,R))/2}
        else if(ssiter==1){
          h <- mat.or.vec(I+1,1)
        }
      
        # bisection on steady state K/Y
        Rlow <-  .99
        Rhigh <- (1 + alpha/KYtarget - delta) + .005

        wlow <- mat.or.vec(I,1)
        wlow[(I0+1):IR] <- (1-alpha)*((Rlow-(1-delta))/(alpha))^(alpha/(alpha-1))      
        KYlow <- NaN
        while(is.nan(KYlow)){
          Rlow <- Rlow + .005
          assets <- compute.asset.holdings(Rlow,wlow,h,beta,sigma,mortality,I0,I)        
          KYlow <- KY(alpha,cohorts,assets)  - (alpha / (Rlow - (1-delta)))
          rm(assets)
        }
        rm(wlow)
      
        whigh <- mat.or.vec(I,1)
        whigh[(I0+1):IR] <- (1-alpha)*((Rhigh-(1-delta))/(alpha))^(alpha/(alpha-1))   
        assets <- compute.asset.holdings(Rhigh,whigh,h,beta,sigma,mortality,I0,I)
        KYhigh <- KY(alpha,cohorts,assets) -  (alpha / (Rhigh - (1-delta)))
        rm(whigh,assets)

        if(KYlow*KYhigh>0){
          print("here")
          print("sign flow = sign fhigh")
          diff = 1E-6
        }

        diff<-1
        iter<-0
        while (diff > 1E-6) {
          Rmid <- (Rhigh + Rlow)/2
          wmid <- mat.or.vec(I,1)
          wmid[(I0+1):IR] <- (1-alpha)*((Rmid-(1-delta))/(alpha))^(alpha/(alpha-1))   
          assets <- compute.asset.holdings(Rmid,wmid,h,beta,sigma,mortality,I0,I)
          KYratio <- KY(alpha,cohorts,assets)
          KYmid <- KYratio - (alpha / (Rmid - (1-delta)))

        
          if (KYmid*KYlow < 0) {
            Rhigh <- Rmid
            KYhigh <- KYmid
          } else {
            Rlow <- Rmid
            KYlow <- KYmid  
          }

          diff <- abs(KYmid)
          iter <- iter + 1
          if(iter > 100){
            diff <- 1E-6
            print("iter > 100")
          }
        }
        ssdiff <- max(abs(assets-laggedassets))
        laggedassets <- assets
        ssiter <- ssiter + 1
      }
    #}
    rm(KYlow,KYmid,KYhigh)
    
    #if(sctr >1 | cctr >1){
      allRs[sctr,cctr] <- 1 + alpha/KYratio - delta
      allKYs[sctr,cctr] <- KYratio

      c <- mat.or.vec(I+1,1)
      for(i in 1 : I){c[i] <- assets[i]*Rmid + wmid[i] + h[i] - assets[i+1]}
      
      plot(assets, xlab="Age",ylab="Assets, consumption, wages",lwd=2.0,type="l",col=c(2),axes=F)
      lines(c[1:101],lwd=2.0,col=c(3))
      lines(wmid[1:101],lwd=2.0,col=c(4))
      lines(h[1:101],lwd=2.0,col=c(5))
      axis(1)
      axis(2)
    #} else {
    #  allRs[sctr,cctr] <- 1 + alpha/KYtarget - delta
    #  allKYs[sctr,cctr] <- KYtarget
    #}
    print("----------------------------------------------------")
    print(allRs)
    print(allKYs)

    # Assign assets and consumption profiles to arrays
    allassetsbeta[sctr,cctr,1:I+1] <- assets[1:I+1]
    allconsumptionbeta[sctr,cctr,1:I+1] <- c[1:I+1]
    # WARNING (18Oct12): Check dimensions of assets and c
    rm(Rhigh,Rmid,Rlow,wmid)
  }
}
rm(sctr,cctr)

######################################################################
# Solve model for wealth-to-capital ratio  -- given beta and real interest rate
# Loop over differences in life expectancy and associated cohort sizes

# beta equal to that calibrated to KY target for life expectancy 70
beta      <- allbetas[3,3]

w <- mat.or.vec(I,1)
w[(I0+1):IR] <- (1-alpha)*((R-(1-delta))/(alpha))^(alpha/(alpha-1))

for(sctr in 1:6){
  for(cctr in 1:6){

    mortality <- alls[,sctr]
    cohorts   <- allc[,cctr]
    
   
    # loop for different social security system / ways to (not to) close the model
    # ssystem set earlier
    ssdiff <- 1
    ssiter <-1
    laggedassets <-  mat.or.vec(I+1,1)
    while (ssdiff > 1E-6) {
      if(ssiter > 1){
        h <- (h + bequests(sssystem,I0,I,mortality,assets,R))/2}
      else if(ssiter==1){
        h <- mat.or.vec(I+1,1)
      }
      assets <- compute.asset.holdings(R,w,h,beta,sigma,mortality,I0,I)
      ssdiff <- max(abs(assets-laggedassets))
      laggedassets <- assets
      ssiter <- ssiter + 1
    }
    rm(ssdiff,ssiter,laggedassets)

    KYratio <- KY(alpha,cohorts,assets)
    WKratio <- (KYratio/KYtarget)^(1/(1-alpha))
    allWKs[sctr,cctr] <- WKratio
    print(allWKs)

    c <- mat.or.vec(I+1,1)
    for(i in 1 : I){c[i] <- assets[i]*R + w[i] + h[i] - assets[i+1]}

    allassetsKYbeta[sctr,cctr,1:I+1] <- assets[1:I+1]
    allconsumptionKYbeta[sctr,cctr,1:I+1] <- c[1:I+1]
    # WARNING (18Oct12): Check dimensions of assets and c
    rm(assets,c,WKratio)
  }
}
rm(sctr,cctr)

rm(age,b_0,b_1,bequests,cohort.sizes)
rm(compute.asset.holdings,diff,dir,gs)
rm(h,HundredandOne,i,InduceLE,KYratio)
rm(law.of.motion,o,Rhigh,Rlow,Rmid,ssdiff)
rm(survival.probabilities,SurvivalProbs,survprob,w,wmid)

save(list = ls(all=TRUE), file = "demographysteadystates.RData")