# Example of estimating an experimental QRE
# by Stephen R. Haptonstahl (srh@haptonstahl.org)
# 
#       D (-1,-3) Burn (1)
#    A   /
# 1 --- 2 --- (1,-2) Give up (2)
#  \       C
# ~A (0,0)
#        Stay home (3)
# 

Eu1.given.A <- function(lambda) (exp(lambda)-1) / (1+exp(lambda))

prob.B.given.lambda <- function(lambda) {
  1 / ((1+exp(lambda)) * (1 + exp(- lambda * Eu1.given.A(lambda))))
}
prob.G.given.lambda <- function(lambda) {
  exp(lambda) / ((1+exp(lambda)) * (1 + exp(- lambda * Eu1.given.A(lambda))))
}
prob.S.given.lambda <- function(lambda) {
  1 / (1 + exp(lambda * Eu1.given.A(lambda)))
}

# Choose "true" lambda values we hope to recover at the end
true.lambdas <- c(.5, 1, 2, 3, 5)

n <- 1000    # number of observations for each value of lambda

# Generate fake data as a list of binned observations
mydata <- lapply(true.lambdas, function(this.lambda) {
  p <- c(prob.B.given.lambda(this.lambda),
    prob.G.given.lambda(this.lambda),
    prob.S.given.lambda(this.lambda)
  )
  return( sample(c("B", "G", "S"), n, replace=T, prob=p) )
})
mydata[[1]]     # look at data from first time period

#######################
###  Fit using MLE  ###
#######################
llike <- function(lambda, outcomes) {
  logPrB <- log(prob.B.given.lambda(lambda))
  logPrG <- log(prob.G.given.lambda(lambda))
  logPrS <- log(prob.S.given.lambda(lambda))
  
  return( logPrB * sum(outcomes=="B") + logPrG * sum(outcomes=="G") + logPrS * sum(outcomes=="S") )
}
fit.mle.1 <- optim(
  par=1,                                    # start with lambda = 1
  fn=llike,                                 # function to optimize
  method="L-BFGS-B", lower=0, upper=Inf,    # search on non-negative reals
  control = list(fnscale=-1),               # maximize, not minimize
  hessian = TRUE,                           # so we can get standard errors
  outcomes=mydata[[1]]                      # other arguments for llike
)
# MLE point estimate
fit.mle.1$par
# standard error
mle.se.1 <- sqrt(- 1 / fit.mle.1$hessian)
mle.se.1

# Fit all using MLE
fit.mle <- list(fit.mle.1)
for(i in 2:length(mydata)) {
  next.fit <- optim(par=1, fn=llike, method="L-BFGS-B", lower=0, upper=Inf, 
    control = list(fnscale=-1), hessian = TRUE, 
    outcomes=mydata[[i]])
  fit.mle <- c(fit.mle, list(next.fit))
}
fit.mle[[1]]

# Grab estimated lambdas, standard errors, 95% conf intervals, and
# generate predicted probabilities given estimiated lambdas
mle.params <- data.frame(
  lambda=fit.mle[[1]]$par,
  se=sqrt(-1/fit.mle[[1]]$hessian),
  lower95=fit.mle[[1]]$par - 1.96 * sqrt(-1/fit.mle[[1]]$hessian),
  upper95=fit.mle[[1]]$par + 1.96 * sqrt(-1/fit.mle[[1]]$hessian)
)
for(i in 2:length(fit.mle)) {
  mle.params <- rbind(mle.params, data.frame(
    lambda=fit.mle[[i]]$par,
    se=sqrt(-1/fit.mle[[i]]$hessian),
    lower95=fit.mle[[i]]$par - 1.96 * sqrt(-1/fit.mle[[i]]$hessian),
    upper95=fit.mle[[i]]$par + 1.96 * sqrt(-1/fit.mle[[i]]$hessian)
  ))
}
row.names(mle.params) <- true.lambdas
mle.params

# Plot results
source("http://sasaki.wustl.edu/srhapton/software/software/public/roundNicely/roundNicely.R")
# initialize blank plot
plot(0, type='n', 
  xlim=c(0, roundNicely(max(mle.params$lambda))),
  ylim=c(0,1),
  main="Experimental QRE", 
  xlab=expression(paste("Learning parameter ",lambda)), 
  ylab="Predicted probability"
)
# plot QRE curves or predicted probabilities given lambda
curve(prob.B.given.lambda, add=T, from=0, col=rainbow(3)[1])
curve(prob.G.given.lambda, add=T, from=0, col=rainbow(3)[2])
curve(prob.S.given.lambda, add=T, from=0, col=rainbow(3)[3])
# plot observed probabilities
for(i in 1:nrow(mle.params)) {
  p.observed <- c(
    mean(mydata[[i]]=="B"),
    mean(mydata[[i]]=="G"),
    mean(mydata[[i]]=="S")
  )
  points(rep(mle.params$lambda[i],3), p.observed, col=rainbow(3), pch=19)
}