library("HiddenMarkov")
library('ggplot2')
library("data.table")
library('ggthemes')

source('lib.R')

family <- "gamma"  # underlying distribution for hmm
discrete <- FALSE


aic1=c()

## Read May Lees Ferry Flows

test=read.table("http://civil.colorado.edu/~balajir/CVEN6834/R-sessions/session3/files-4HW3/LeesFerry-monflows-1906-2016.txt")
## select the May month flows and convert to MAF
x = test[,6]/10^6		

#LFann <- read.table("http://civil.colorado.edu/~balajir/CVEN6833/R-sessions/session3/files-4HW3/LF_1906-2005.txt")
#x = ts(LFann[,2])/10^6   #convert to MAF



## Fit HMM models of orders 2 through 6. Obtain the AIC  for each
## They are stored in arrays aic1 and bic1
## Spit it out and the best order is the one with the least value of AIC .
## For Colorado flows it will be order 2

for(imodel in 1:6){

m <- imodel      #model order to fit

stationary <- F   # use a stationary distribution of mixtures
    # different initial condition types when family == "norm"
ic <- "same.sd"#c("same.sd","same.both","both.diff")

fd.name <- ifelse(family == "norm", "normal", family)

  # T.P.M.
Pi <- Pi_init(m)
  # Other ways to specify Pi
#Pi <- diag(rep(1-.05*(m-1),m))
#Pi[Pi == 0] <- .05
  # random guess for the transition matrix
#Pi <- matrix(runif(m^2),m)
#Pi <- Pi/rowSums(Pi)
delta <- delta_init(m)

pars <- get.named.parlist(x, m, fd.name, lower=.0, ic)#,start=list(shape1=2,shape2=2))

  # set up the model
hmm <- dthmm(x, Pi=Pi, delta=delta, family, pars, nonstat=!stationary, discrete = discrete)

if(imodel < 2){
 hmm <- BaumWelch(hmm, bwcontrol(maxiter = 1000,posdiff=TRUE,converge = expression(diff > tol)))
 } else {
 
 hmm <- BaumWelch(hmm, bwcontrol(maxiter = 1000, tol = 1e-08))

}


  # get the hidden states from the fitted model 
  # Global decoding
  # To get the probability of being in a state: hmm$u
decoding <- Viterbi(hmm)


  # get AIC


aic <- AIC(hmm)

aic1=c(aic1,aic)

}

## Get the best order
bestorder = order(aic1)[1]

## Fit the model for this best order
###### For Colorado River Flows the parameters for bestorder model of 2
### obtained below will be same as what is reported in Bracken et al. (2015) paper

m <- bestorder      #model order to fit

stationary <- F   # use a stationary distribution of mixtures
    # different initial condition types when family == "norm"
ic <- "same.sd"#c("same.sd","same.both","both.diff")

fd.name <- ifelse(family == "norm", "normal", family)

  # T.P.M.
Pi <- Pi_init(m)
 
delta <- delta_init(m)

pars <- get.named.parlist(x, m, fd.name, lower=.0, ic)#,start=list(shape1=2,shape2=2))

  # set up the model
hmm <- dthmm(x, Pi=Pi, delta=delta, family, pars, nonstat=!stationary, discrete = discrete)
#sink(tempfile())

hmm <- BaumWelch(hmm, bwcontrol(maxiter = 1000, tol = 1e-08))

sink()

  # get the hidden states from the fittet model 
  # Global decoding
  # To get the probability of being in a state: hmm$u

  ##the vector decoding gives the best estimate of the hidden state
  ## for each observation
  
 decoding <- Viterbi(hmm)

# Gives all the model parameters
print(summary(hmm))

cat('Model order:',m,'\n')


p <- ggplot_stationary_hmm(hmm,.5)
print(p)
ggsave(paste0(tools::file_path_sans_ext(datafile),'.pdf'),p,width=8.5,height=4)

############################
############### Now simulate ###########

# First simulate a sequence of states from the TPM

## simulate from the transition probability
N = length(x)
nprob = length(decoding[decoding == 1])/N
delta1=c(nprob,1-nprob)   #stationary probability
zsim = mchain(NULL,hmm$Pi,delta=delta1)
zsim = simulate(zsim,nsim=N)

## now simulate the flows from the corresponding PDF

flowsim = c()
for(i in 1:N){

if(zsim$mc[i] == 1)xx=rgamma(1,shape=hmm$pm$shape[1],scale=1/hmm$pm$rate[1])
if(zsim$mc[i] == 2)xx=rgamma(1,shape=hmm$pm$shape[2],scale=1/hmm$pm$rate[2])

flowsim=c(flowsim,xx)
}