#  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 = "C:/Users/dbackus/Documents/Papers/BCH/model/steady_state"
setwd(dir) 

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

# Preferences
sigma <- 1

# Technology
alpha <- 0.33
delta <- .05

# Demographics
I0 <- 20
I  <- 101
IR <- 60

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


# 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 - (1/(beta*s[i-1]*R))^(-sigma)*(a[i-1]*R + w[i-1] + h - a[i])
  }
  return(a)
}

# 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[21: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)

# what is this??  
allbetas <-array(0, dim=c(6,6),dimnames=list(c("s60","s65","s70","s75","s80","s85"),c("c60", "c65","c70","c75","c80","c85")))

# big loop?
for(sctr in 1:6){
  for(cctr in 1:6){

    mortality <- alls[,sctr]
    cohorts <- allc[,cctr]
    
    h <- 0  # ?? why zero?
    #h<-0.05977687

    # bisection on beta:  pick to hit steady state K/Y 
    betalow <- .925
    betahigh <- 1.01

    KYlow <- NaN
    while(is.nan(KYlow)){
      betalow <- betalow + .005
      assets <- compute.asset.holdings(R,w,h,betalow,sigma,mortality,I0,I)
      N <- sum(cohorts[21:IR])
      K<-0
      for(i in 1 : 101){
        K <- K + assets[i]*cohorts[i]
      }
      Y <- K^alpha * N^(1-alpha)
      print(K/Y)
      KYlow <- K/Y - KYtarget
    }
    
    assets <- compute.asset.holdings(R,w,h,betahigh,sigma,mortality,I0,I)
    N <- sum(cohorts[21:61])
    K<-0
    for(i in 1 : 101){
      K <- K + assets[i]*cohorts[i]
    }
    Y <- K^alpha * N^(1-alpha)
    print(K/Y)
    KYhigh <- K/Y - KYtarget

    betamid <- (betahigh + betalow)/2
    assets <- compute.asset.holdings(R,w,h,betamid,sigma,mortality,I0,I)
    N <- sum(cohorts[21:61])
    K<-0
    for(i in 1 : 101){
      K <- K + assets[i]*cohorts[i]
    }
    Y <- K^alpha * N^(1-alpha)
    print(K/Y)
    KYmid <- K/Y - KYtarget

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

    diff <- 1
    iter <- 0
    while (diff > 1E-6) {
      if (KYmid*KYlow < 0) {
        betahigh <- betamid
        KYhigh <- KYmid
      } else {
        betalow <- betamid
        KYlow <- KYmid  
      }
      betamid <- (betahigh + betalow)/2
      assets <- compute.asset.holdings(R,w,h,betamid,sigma,mortality,I0,I)
      #N <- sum(cohorts[21:60])
      K<-0
      for(i in 1 : 101){
        K <- K + assets[i]*cohorts[i]
      }
      Y <- K^alpha * N^(1-alpha)
      KYmid <- K/Y - KYtarget

      diff <- abs(KYmid)
      iter <- iter + 1
      if(iter > 100){
        diff <- 1E-6
        print("iter > 100")
        print(K/Y)
      }
    }

    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 - assets[i+1]
      #c[i] <- assets[i]*R + w[i] + h - assets[i+1]
      wN <- wN + w[i]*cohorts[i]
      rK <- rK + (R-(1-delta))*assets[i]*cohorts[i]
      C <- C + c[i]*cohorts[i]
    }

    accbeq <- mat.or.vec(I+1,1)
    H <- 0
    for(i in 1 : I){
      if(i >= I0){
        accbeq[i] <- R*assets[i]*(1-mortality[i-1])*cohorts[i]
        H <- H + accbeq[i]
      }
    }

    print(K/Y)
    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)
    }

    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))
    axis(1)
    axis(2)
    print(allbetas)
  }
}