
####################################################
#    Function definitions for Chapter 10


####################################################
# A NAIVE implementation of simulated annealing
NaiveSA <- function(fn,par,temp0,beta,maxit,sigma){
  xNow <- par    
  fNow <- fn(xNow)
  xBest <- xNow
  fBest <- fNow
  temp <- temp0
  dim <- length(x0)
  for (k in 1:maxit){
    # generate candidate
    xNew <- xNow+rnorm(dim,sd=sigma) 
    fNew <- fn(xNew)
    if (fNew<fNow){
      xNow <- xNew; fNow <- fNew  
      if (fNew < fBest){
        xBest <- xNew; fBest <- fNew
      }
    }
    else{
      if (runif(1) <= exp((fNow-fNew)/temp)){
        xNow <- xNew; fNow <- fNew  
      }
    }
    temp <- beta*temp
  }
  return(list(value=fBest,par=xBest,counts=maxit))
}


####################################################
# Optimal solution of a newsvendor problem (analytical)
newsOptimal <- function(pur,sell,mdown,mu,sigma){
  # define standard loss function
  lossFun <- function(z) dnorm(z) - z*(1-pnorm(z))
  # get economics
  m <- sell - pur   # profit margin
  cu <- pur - mdown  # cost of unsold items
  # get quantile and optimal q
  z <- qnorm(m/(m+cu))
  q <- mu + z*sigma
  # print solution and cost
  cat('optimal q ', q, '\n')
  cost <- cu*(q-mu)+(cu+m)*sigma*lossFun((q-mu)/sigma)
  cat('optimal cost ', cost, '\n')
}

####################################################
# Solving a newsvendor problem by scenario-based 
# stochastic linear programming
newsLP <- function(m,cu,demand,probs){
  library(linprog)
  numScen <- length(demand)
  Amat <- cbind(rep(1,numScen),diag(numScen),-diag(numScen))
  cvet <- c(0, m*probs, cu*probs)
  eqs <- rep("=",numScen)
  # solve and print solution
  result <- solveLP(cvet,demand,Amat,const.dir=eqs,
                    lpSolve=TRUE)
  return(list(cost=result$opt,q=result$solution[1]))
}


##########################################################
# State space discretization and splines to approximate
# the value functions for a deterministic budget allocation problem
Budgeting <- function(budget,funcList,numPoints){
  ### INITIALIZATION: set up a simple grid
  budgetGrid <- seq(0,budget,length.out=numPoints)
  # prepare matrix of function values for each time step
  numSteps <- length(funcList)
  valFunMatrix <- matrix(0,nrow=numPoints,ncol=numSteps)
  splinesList <- vector("list",numSteps)
  # start from last step (increasing contribution functions)
  valFunMatrix[,numSteps] <- funcList[[numSteps]](budgetGrid)
  # now step backward in time
  for (t in seq(from=numSteps-1,to=1,by=-1)) {
    # build function adding immediate contribution and 
    # value function t plus 1, at the next stage
    splinesList[[t+1]] <- splinefun(budgetGrid, 
                                    valFunMatrix[,t+1])
    for (k in 1:numPoints){
      b <- budgetGrid[k]
      objFun <- function(z) funcList[[t]](z)+
        splinesList[[t+1]](b-z)
      if (b>0){
        outOptim <- optimize(objFun,lower=0,upper=b,
                             maximum=TRUE)
        valFunMatrix[k,t] <- outOptim$objective}
      else  valFunMatrix[k,t] <- objFun(0)
    }
  }
  # Now optimize forward in time and return values
  residualB <- budget;  objVal <- 0
  x <- numeric(numSteps)
  for (t in 1:(numSteps-1)){
    objFun <- function(z) funcList[[t]](z)+
      splinesList[[t+1]](residualB-z)
    if (residualB > 0){
      outOptim <- optimize(objFun,lower=0,upper=residualB,
                           maximum=TRUE)
      x[t] <- outOptim$maximum}
    else  x[t] < 0
    objVal <- objVal + funcList[[t]](x[t])
    residualB <- residualB - x[t]
  }
  x[numSteps] <- residualB
  objVal <- objVal + funcList[[numSteps]](x[numSteps])
  # assemble output
  out <- list(par=x,value=objVal,valFuns=valFunMatrix)
}

###################################################
# Numerical value iteration for a simple MDP
PW311_ValueIteration <- function(probs, payoff, discount, 
                                 eps=0.00001, maxIter=1000){
  # Initialization
  numStates <- length(payoff)
  oldV <- numeric(numStates)
  newV <- numeric(numStates)
  myeps <- eps*(1-discount)/2/discount
  stopit <- FALSE
  count <- 0
  # precompute invariant discounted probs
  dprobStay <- discount*probs
  dprobMove <- discount*(1-probs)
  # we have to decide only in interior states
  indexVet <- 2:(numStates-1)
  norm2 <- function(v) {return(sqrt(sum(v^2)))}
  # value iteration
  while (!stopit){
    count <- count + 1
    # update: first and last state are apart
    newV[1] <- dprobStay[1]*oldV[1]+dprobMove[1]*oldV[2]
    valueFlush <- payoff[indexVet] + discount*oldV[1]
    valueWait <- dprobStay[indexVet]*oldV[indexVet]+
      dprobMove[indexVet]*oldV[indexVet+1]
    newV[indexVet] <- pmax(valueWait,valueFlush )
    newV[numStates] <- payoff[numStates] + discount*oldV[1]
    if ((max(newV-oldV)<myeps) || (count > maxIter))
      stopit <- TRUE
    else
      oldV <- newV
  }
  # return state values and policy vector (1->flush is true)
  return(list(stateValues=newV, 
              policy=c(0,valueFlush>valueWait,1)))
}


###################################################
# Policy simulation for a simple MDP
PW311_Simulate <- function(probs,payoff,discount,startState, 
                           numRepl,numSteps,policy){
  numStates <- length(probs)
  sumRewards <- numeric(numRepl)
  for (k in 1:numRepl){
    state <- startState
    dRewards <- numeric(numSteps)
    discf <- 1
    for (t in 1:numSteps){
      if (state == 1){
        # state 1, no choice
        if (runif(1) > probs[state])
          state <- 2
      }
      else if (state == numStates){
        # last state, get reward and back to square 1
        state <- 1
        dRewards[t] <- discf*payoff[numStates]        
      }
      else {
        # apply policy
        if (policy[state] == 1){
          # get reward and go back to square 1
          dRewards[t] <- discf*payoff[state]        
          state <- 1
        }
        else {
          # random move 
          if (runif() > probs[state])
            state <- state + 1
        }
      } # end if on states
      discf <- discf * discount
    } # end single replication
    sumRewards[k] <- sum(dRewards)
  } # end replications
  # return point estimate and confidence interval (95%)
  aux <- t.test(sumRewards)
  value <- as.numeric(aux$estimate)
  ci <- as.numeric(aux$conf.int)
  return(list(estimate=value,conf.int=ci))
}

#################################################################
# Numerical DP for the consumption--saving problem
DPSaving <- function(W0,empl0,utilFun,gamma,numSteps,
                     numPoints,Wmin,Wmax,income,transMatrix,
                     retScenarios,probs,rf,splineOption="natural"){
  # initialization
  numScen <- length(probs)
  wealthValues <- seq(Wmin,Wmax,length.out=numPoints)
  valFunMatrix <- matrix(0,nrow=numPoints,ncol=3)
  valFunsList <- vector("list",numSteps)
  # initialize last time period  
  auxFunsList <- vector("list",3)
  for (i in 1:3) auxFunsList[[i]] <- utilFun
  valFunsList[[numSteps]] <- auxFunsList
  # backward recursion for each time period
  for (t in seq(from=numSteps-1,to=1,by=-1)) {
    # prepare nested list of splinefuns
    auxFunsList <- vector("list",3)
    valFunsList[[t]] <- auxFunsList
    # nested for loop on the grid
    for (i in 1:3){
      for (k in 1:numPoints){
        # build objective function
        objFun <- function(x){
          C <- x[1]; alpha <- x[2]; out <- 0
          for (j in 1:3){
            for (l in 1:numScen){
              futWealth <- income[j]+(wealthValues[k]-C)*
                (1+rf+alpha*(retScenarios[l]-rf))
              out <- out+probs[l]*transMatrix[i,j]*
                valFunsList[[t+1]][[j]](futWealth)}}
          out <- utilFun(C)+gamma*out
          return(-out)}
        # solve optimization problem
        x0 <- c(wealthValues[k]/2,0.5)
        outOptim <- nlminb(start=x0,objective=objFun,
                           lower=c(0,0),upper=c(wealthValues[k],1))
        valFunMatrix[k,i] <- -outOptim$objective
      }
      # use function values and build spline
      valFunsList[[t]][[i]] <- splinefun(wealthValues,
                                         valFunMatrix[,i],method=splineOption)
    } # end for on wealth values
  } # end for loop on employment states
  return(valFunsList)
}

# Make scenarios for optimizing the consumption-saving 
# policy by Gaussian quadrature
MakeScenariosGQ <- function(mu,sigma,numScenarios){
  require(statmod)
  aux <- gauss.quad.prob(numScenarios,dist="normal",
                         mu=mu-sigma^2/2,sigma=sigma)
  probRets <- aux$weights
  riskyRets <- exp(aux$nodes)-1
  out <- list(retScenarios=riskyRets,probs=probRets)
}

# Generate scenarios to SIMULATE the consumption-saving
# policy
sampleEmployment <- function(empl0,cumProbs,numSteps){
  # cumProbs are cumulated probs for efficiency
  state <- c(empl0, numeric(numSteps))
  for (t in 1:numSteps){
    u <- runif(1)
    state[t+1] <- sum(u > cumProbs[state[t],])+1
  }
  return(state[-1])
}

MakeSamplePaths <- function(empl0,transMatrix,mu,sigma,
                            numSteps,numRepl){
  # cumulate transition probs once
  cumProbs <- t(apply(transMatrix,1,cumsum))
  emplPaths <- matrix(0,nrow=numRepl,ncol=numSteps)
  retPaths <- matrix(0,nrow=numRepl,ncol=numSteps)
  for (k in 1:numRepl){
    # sample employment scenario
    emplPaths[k,] <- 
      sampleEmployment(empl0,cumProbs,numSteps)
    # sample return scenario
    retPaths[k,] <- 
      exp(rnorm(numSteps,mu-sigma^2/2,sigma))-1
  }
  return(list(emplPaths=emplPaths,retPaths=retPaths))
}


### Simulate a FIXED consumption-saving policy
SimulateFixedSaving <- function(W0,empl0,utilFun,gamma,
                                income,transMatrix,mu,sigma,rf,
                                scenarios,alpha,fractCons){
  # allocate data structures
  dims <- dim(scenarios$retPaths)
  numRepl <- dims[1]
  numSteps <- dims[2]
  alphaPaths <- matrix(alpha,nrow=numRepl,ncol=numSteps)
  CPaths <- matrix(0,nrow=numRepl,ncol=numSteps+1)
  wealthPaths <- matrix(0,nrow=numRepl,ncol=numSteps+1)
  utilVals <- numeric(numRepl)
  # start replications
  for (k in 1:numRepl){
    # get employment and return scenario
    emplPath <- scenarios$emplPaths[k,]
    riskyRets <- scenarios$retPaths[k,]
    wealthNow <- W0
    stateNow <- empl0
    for (t in 0:(numSteps-1)){
      # collect path info
      C <- fractCons*wealthNow
      CPaths[k,t+1] <- C
      wealthPaths[k,t+1] <- wealthNow
      utilVals[k] <- utilVals[k] + gamma^(t)*utilFun(C)
      # update wealth according to actual scenario
      wealthNow <- income[emplPath[t+1]]+(wealthNow-C)*
        (1+rf+alpha*(riskyRets[t+1]-rf))      
    }
    # at last step, just consume available wealth
    wealthPaths[k,numSteps+1] <- wealthNow
    CPaths[k,numSteps+1] <- wealthNow
    utilVals[k] <- utilVals[k] + 
      gamma^(numSteps)*utilFun(wealthNow)
  }
  out <- list(alphaPaths=alphaPaths, CPaths=CPaths,
              wealthPaths=wealthPaths, utilVals=utilVals)
  return(out)
}

### Simulate the optimal (DP) consumption-saving policy
SimulateDPSaving <- function(W0,empl0,utilFun,gamma,
                             valFunsList,income,transMatrix,inRetScenarios,
                             probs,rf,outScenarios){
  dims <- dim(outScenarios$retPaths)
  numRepl <- dims[1]; numSteps <- dims[2]
  alphaPaths <- matrix(0,nrow=numRepl,ncol=numSteps)
  CPaths <- matrix(0,nrow=numRepl,ncol=numSteps+1)
  wealthPaths <- matrix(0,nrow=numRepl,ncol=numSteps+1)
  utilVals<-numeric(numRepl);numInScenarios<-length(probs)
  for (k in 1:numRepl){  # start replications
    emplOutScenario <- outScenarios$emplPaths[k,]
    retOutScenario <- outScenarios$retPaths[k,]
    wealthNow <- W0  # initialize state
    stateNow <- empl0
    for (t in 0:(numSteps-1)){
      objFun <- function(x){ # build objective function
        C <- x[1]; alpha <- x[2]; out <- 0
        for (j in 1:3){
          for (l in 1:numInScenarios){
            futWealth <- income[j]+(wealthNow-C)*
              (1+rf+alpha*(retScenarios[l]-rf))
            out <- out+probs[l]*transMatrix[stateNow,j]*
              valFunsList[[t+1]][[j]](futWealth)}}
        out <- utilFun(C)+gamma*out
        return(-out)}
      x0 <- c(wealthNow/2,0.5) # solve optim. problem
      outOptim <- nlminb(start=x0,objective=objFun,
                         lower=c(0,0),upper=c(wealthNow,1))
      # collect path info
      C <- outOptim$par[1]; alpha <- outOptim$par[2]
      CPaths[k,t+1] <- C; alphaPaths[k,t+1] <- alpha
      wealthPaths[k,t+1] <- wealthNow
      utilVals[k] <- utilVals[k] + gamma^(t)*utilFun(C)
      wealthNow<-income[emplOutScenario[t+1]]+(wealthNow-C)*
        (1+rf+alpha*(retOutScenario[t+1]-rf))      
    }
    wealthPaths[k,numSteps+1] <- wealthNow
    CPaths[k,numSteps+1] <- wealthNow
    utilVals[k] <- utilVals[k] + 
      gamma^(numSteps)*utilFun(wealthNow)
  }
  out <- list(alphaPaths=alphaPaths, CPaths=CPaths,
              wealthPaths=wealthPaths, utilVals=utilVals)
  return(out)
}


### Solving a simple MDP by Q-learning
PW311_QLearning <- function(probs,payoff,discount,numSteps,
                            startState,startQFactors,alpha){
  # NOTE: action indexes in R are incremented by 1
  # when accessing vectors and matrices
  QFactors <- startQFactors
  state <- startState
  numStates <- length(probs)  
  for (k in 1:numSteps){
    if (state == 1){ # state 1, no choice
      newState <- ifelse(runif(1)>probs[state],2,1)
      action <- 0
      reward <- 0
    }
    else if (state == numStates){
      # last state, get reward and go back to square 1
      newState <- 1
      action <- 1
      reward <- payoff[state]
    }
    else { # Choose optimal action
      action <- which.max(QFactors[state,]) - 1
      if (action == 1){# get reward and go back to square 1
        newState <- 1
        reward <- payoff[state]
      }
      else { # random transition 
        if (runif(1) > probs[state]){
          newState <- newState + 1
          reward <- 0
        }
      }
    } # end if on states
    qhat <- reward+discount*max(QFactors[newState,])
    QFactors[state,action+1] <- 
      alpha*qhat+(1-alpha)*QFactors[state,action+1]
    state <- newState
  } # end for
  # find optimal policy (for interior states)
  indexVet <- 2:(numStates-1)
  policy <- c(0,QFactors[indexVet,2]>QFactors[indexVet,1],1)
  out <- list(QFactors=QFactors,policy=policy)
  return(out)
}

### Solving a simple MDP by Q-learning and EPS Greedy strategy
PW311_QepsLearning <- function(probs,payoff,discount,numSteps,startState,
                               startQFactors,alpha,eps){
  # NOTE: action indexes in R are incremented by 1
  # when accessing vectors and matrices
  QFactors <- startQFactors
  state <- startState
  numStates <- length(probs)  
  for (k in 1:numSteps){
    if (state == 1){
      # state 1, no choice
      newState <- ifelse(runif(1)>probs[state],2,1)
      action <- 0
      reward <- 0
    }
    else if (state == numStates){
      # last state, get reward and go back to square 1
      newState <- 1
      action <- 1
      reward <- payoff[state]
    }
    else {
      if (runif(1) > eps) # Choose optimal action
        action <- which.max(QFactors[state,]) - 1
      else # Choose random action
        action <- ifelse(runif(1)<0.5,0,1)
      if (action == 1){
        # get reward and go back to square 1
        newState <- 1
        reward <- payoff[state]
      }
      else {
        # random transition 
        if (runif(1) > probs[state]){
          newState <- newState + 1
          reward <- 0
        }
      }
    } # end if on states
    qhat <- reward+discount*max(QFactors[newState,])
    QFactors[state,action+1] <- 
      alpha*qhat+(1-alpha)*QFactors[state,action+1]
    state <- newState
  } # end for
  # find optimal policy (for interior states)
  indexVet <- 2:(numStates-1)
  policy <- c(0,QFactors[indexVet,2]>QFactors[indexVet,1],1)
  out <- list(QFactors=QFactors,policy=policy)
  return(out)
}
