library('RSQLite')
library('ggplot2')
library('data.table')
library('reshape2')
library('plyr')
library('e1071') 
library('sm')
library('tikzDevice')
library('HiddenMarkov')
library('zoo')
library('ggthemes')
library('scales')
library('dplyr')

source('lib.R')


# This file is the last step in the HMM simulation analysis
# it relies on tables that are created by other scripts 
# exiting in the simulation database. 

dbname <- 'hmm_sim.db'
m <- dbDriver('SQLite')

con <- dbConnect(m, dbname = dbname)

# this table created by 'hmm-sim.R'
simulations <- as.data.table(dbReadTable(con,'simulations'))
dens <- as.data.table(dbReadTable(con,'simulation_pdfs'))

# this table is generated in 'hmm-input-data.R'
dt <- as.data.table(dbReadTable(con,'input_data'))

# these two table are created in 'hmm-sim-wavelet.R'
local_spec_med <- as.data.table(dbReadTable(con, 'wavelet_sim_median'))
spec_bounds <- as.data.table(dbReadTable(con, 'wavelet_sim_sawp_bounds'))

spec_obs <- wavelet_dt(dt$x,timelab=dt$date)
obs_scale <- wavelet_scale(dt$x)
obs_power <- wavelet_power(dt$x)
coi <- wavelet_default(dt$x)$coi[-(1:2)]
coi_time <- dt$date[-(1:2)]

plot_dir <- '../hmm-sim-paper/plots/'

options(warn=-1)
pars <- get.named.parlist(dt$x, 2, 'gamma', 'same.sd')
hmm <- dthmm(dt$x, Pi_init(2), delta_init(2), 'gamma', pars)

sink('/dev/null')
hmm <- BaumWelch(hmm)
sink()

#tikz(file.path(plot_dir,'sim-lengths-cum-mean.tex'),width=pc2in(21),height=pc2in(22),pointsize=10)
#cumsum.after <- 5
#max.plot <- 14
#runs <- rl.count(as.vector(x),ns)
#mean.counts <- rbind(
#    apply(ar.1.stats$counts,2,mean),
#    apply(hmm.2.stats$counts,2,mean)
#    #c(runs$count,rep(0,count.max-length(runs$count)))
#    )
#colnames(mean.counts) <- 1:count.max
#rownames(mean.counts) <- c('AR1','HM2G')#, 'Obs')#, 'NHM')
#mean.counts.pdf <- mean.counts[,(cumsum.after + 1):max.plot]
#mean.counts.cdf <- t(apply(mean.counts[,!(colnames(mean.counts) %in% 1:(cumsum.after))],1,cumsum))[,1:(max.plot-cumsum.after)]
#colnames(mean.counts.cdf) <- (cumsum.after + 1):max.plot

#layout(cbind(1:2))
#par(mar=c(3,4,1,1)+.1)
#barplot(mean.counts.pdf,beside=T,legend=T,xlab='Run Length',ylab='Mean Count',
#  ylim=c(0,max(mean.counts.pdf+.051)))
#barplot(mean.counts.cdf,beside=T,xlab='Run Length',ylab='Mean Count',
#  ylim=c(0,max(mean.counts.cdf+.051)))#c('AR1','HM2G','HM2N'),
#  #args.legend=list(x='bottomright',bg='white'))
#mtext('Run Length', 1,2)
#dev.off()

rl_obs <- rl.count(dt$x,2)
obs_rl <- data.frame(x=rl_obs$len[-1],y=rl_obs$count[-1])
rl = dplyr::summarise(group_by(simulations[kind=='HMC'],kind,sim),
        rl=rl.count(value,2)$count[-1],n=rl.count(value,2)$len[-1])
p_rl <- ggplot(rl[n<8])+
    geom_boxplot(aes(x=factor(n),y=rl),outlier.colour=NA)+
    theme_bw()+
    geom_point(data=obs_rl, aes(x=x-1,y=y), size=3)+
    geom_line(data=obs_rl, aes(x=x-1,y=y))+
    xlab('Run Length')+ylab('Number of Runs')+
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())
ggsave_tikz(file.path(plot_dir,'sim-run-length.tex'),width=pc2in(20),height=pc2in(15),pointsize=10)

#tikz(file.path(plot_dir,paste(dataset,dist,'sim-lengths-density.tex',sep='-')),width=pc2in(20),height=pc2in(20),pointsize=10)  
#    upper.hmm <- apply(hmm.2.stats$dens,2,quantile,.99)
#    lower.hmm <- apply(hmm.2.stats$dens,2,quantile,.01)
#    med.hmm <- apply(hmm.2.stats$dens,2,quantile,.5)
#
#    d.hmm.sim <- density(hmm.2.stats$runs[1,],bw=1,from=1,to=count.max,na.rm=T)
#    d.hist <- density(rl.count(as.vector(meko),3)$rl,bw=1,from=1,to=13)
#    hxp <- d.hmm.sim$x
#    
#    plot(hxp,upper.hmm,xlim=c(1,10),type='n',xlab='Run Lengths',ylab='Probability Density')
#    grid(lty=1,lwd=.5,col=gray(.9))
#    polygon(c(hxp,rev(d.hmm.sim$x)),c(upper.hmm,rev(lower.hmm)),col='lightgray',border=NA)
#    lines(d.hist,col='red')
#    legend('topright',c('HM2G','Historical'),lty=c(0,1,1),
#        col=c('lightgray','red'),pch=c(22,-1,-1),pt.bg='lightgray',bg='white')
#dev.off()

printf('Simulation Statistics...')
stats <- as.data.table(melt(ddply(simulations, .(sim,kind), summarise, 
    Mean=mean(value), 
    Minimum=min(value),
    Maximum=max(value), 
    `Standard Deviation`=sd(value), 
    Skewness=skewness(value), 
    `Lag 1 AC`=lag1(value)),
    id.vars=c('sim','kind')))

obs <- obs_stats(dt$x)

limits <- ddply(stats,.(kind,variable),summarise,
    value=quantile(value,c(.05,.25,.5,.75,.95)),
    limit=c('ymin','lower','middle','upper','ymax'))

limits <- as.data.table(dcast(limits,kind+variable~limit))


p_stats <- ggplot(limits[kind!='AR'&kind!='HMI'],aes(ymin=ymin,lower=lower,middle=middle,upper=upper,ymax=ymax))+
    geom_boxplot(aes(x=kind,group=kind),stat='identity',outlier.colour=NA,width=.8)+
    facet_wrap(~variable,scales='free_y')+
    geom_hline(data=obs,aes(yintercept=value,group=variable),linetype=2)+
    theme_bw()+
    xlab('')+ylab('')+
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())

ggsave_tikz(file.path(plot_dir,'sim-stats.tex'),width=pc2in(39),height=pc2in(24.4),pointsize=10)

######################################################################################################
######################################################################################################
printf('Simulation PDF...')
evalp <- seq(0,30,.5)
d <- sm.density(dt$x, eval.points=evalp,display='none')$estimate

limitsd <- ddply(dens[kind=='HMC'],.(x),summarise,
    value=fivenum(density),
    limit=c('ymin','lower','middle','upper','ymax'))

limitsd <- dcast(limitsd,x~limit)

p_pdf <- ggplot(limitsd,aes(ymin=ymin,lower=lower,middle=middle,upper=upper,ymax=ymax))+
    geom_boxplot(aes(x=x,group=x),stat='identity',outlier.colour=NA,width=.4)+
    geom_line(aes(x=evalp,y=d),col='#0072B2',size=1.5)+
    theme_bw()+
    xlab('Flow Volume [MAF]')+ylab('Probability Density') +
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())

#ggplot(dens[kind=='climate'])+
#    geom_boxplot(aes(x=x,y=density,group=x),outlier.colour=NA,width=.4)+
#    geom_line(aes(x=evalp,y=d),col='#0072B2',size=1.5)+
#    theme_bw()+
#    xlab('Flow Volume [MAF]')+ylab('Probability Density')

ggsave_tikz(file.path(plot_dir,'sim-density.tex'),p_pdf,width=pc2in(20),height=pc2in(15),pointsize=10)

######################################################################################################
######################################################################################################
printf('HMM Decoding...')
p1_decoding <- ggplot()+
    geom_bar(aes(x=dt$date,y=dt$state),width=.5,stat='identity')+
    xlab('Time')+ylab('State')+
    theme_bw()+
    coord_cartesian(ylim=c(.95,2.05)) +
    scale_y_continuous(breaks=seq(1,2,1)) +
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())

windows = data.frame(x1=c(1906,1930,1982,1987),x2=c(1929,1981,1986,2009))
ys = data.frame(y1=numeric(nrow(windows)),y2=numeric(nrow(windows)))
for(i in 1:nrow(windows))
    ys$y1[i] = ys$y2[i] = mean(dt[date>=windows[i,1] & date<windows[i,2]]$x)
means <- cbind(windows,ys)
p2_decoding <- ggplot()+
    geom_line(aes(x=dt$date,y=dt$x))+
    geom_segment(data=means,aes(x1,xend=x2,y1,yend=y2))+
    xlab('Time')+ylab('Flow Volume [MAF]')+
    theme_bw() +
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())

tikz(file.path(plot_dir,'decoding-period-means.tex'),width=pc2in(20),height=pc2in(20),pointsize=10)
multiplot(
    p2_decoding+annotate('text',2000,23,label='(a)'),
    p1_decoding+annotate('text',2000,1.8,label='(b)')
)
dev.off()


######################################################################################################
######################################################################################################
printf('Timeseries...')
tsdf <- data.frame(x=dt$date,y=dt$x-mean(dt$x))
smdf <- data.frame(x=dt$date,y=as.vector(rollapply(dt$x,5,mean,align='center',partial=T)-mean(dt$x)))
p_ts <- ggplot()+
    geom_bar(data=tsdf,aes(x=x,y=y),stat='identity',color='black',fill=NA,position='dodge')+
    geom_line(data=smdf,aes(x=x,y=y),size=1.3)+
    xlab('Time')+ylab('Flow [MAF]') +
    #geom_abline(slope=-0.03416,intercept=81.83134-mean(x))+
    theme_bw()+
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())
ggsave_tikz(file.path(plot_dir,'lees-ts.tex'),p_ts,width=pc2in(39),height=pc2in(14),pointsize=10)

######################################################################################################
######################################################################################################
printf('HMM Stationary Distribution...')
options(tikzMetricPackages=c(getOption('tikzMetricPackages'),'\\usepackage{amsmath}\n'))
p_statdist <- ggplot_stationary_hmm(hmm)+
    xlab('Flow Volume [MAF]')+ylab('density')+
    ggtitle(paste('$\\boldsymbol\\delta=[',paste(round(statdist(hmm$Pi),2),collapse=', '),']$',sep=''))+
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())
ggsave_tikz(file.path(plot_dir,'statdist.tex'),p_statdist,width=pc2in(20),height=pc2in(15),pointsize=10)


######################################################################################################
######################################################################################################
printf('HMM SAWP Spectrum...')
p <- ggplot(spec_bounds)+
    geom_line(aes(x=scale,y=value, color=kind, linetype=variable))+
    geom_line(data=data.frame(x=obs_scale,y=obs_power),aes(x=x,y=y))+
    scale_x_log10()+
    theme_bw()+
    scale_color_tableau() +
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())
ggsave_tikz(file.path(plot_dir,'wavelet-spectrum-ave.tex'),p,width=pc2in(39),height=pc2in(15),pointsize=10)


######################################################################################################
######################################################################################################
printf('HMM Imposed Spectrum...')
p_hmm_imposed <- ggplot(local_spec_med[kind=='HMI'])+
    geom_tile(aes(x=time,y=scale,fill=power))+
    geom_line(data=data.frame(x=coi_time,y=coi),aes(x=x,y=y),size=1)+
    theme_bw()+xlab('Year')+ylab('Scale')+
    scale_y_continuous(trans=log2_trans(),limits=c(1.99,65))+
    scale_fill_gradientn(colours=colorblind_pal()(8)[-1], guide=guide_colourbar(raster=FALSE, title='Power'))+
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())
ggsave_tikz(file.path(plot_dir,'wavelet-spectrum-imposed.tex'),p_hmm_imposed,width=pc2in(20),height=pc2in(15),pointsize=10)

######################################################################################################
######################################################################################################
printf('HMM Climate Spectrum...')
p_hmm_climate <- ggplot(local_spec_med[kind=='HMC'])+
    geom_tile(aes(x=time,y=scale,fill=power))+
    geom_line(data=data.frame(x=coi_time,y=coi),aes(x=x,y=y),size=1)+
    theme_bw()+xlab('Year')+ylab('Scale')+
    scale_y_continuous(trans=log2_trans(),limits=c(1.99,65))+
    scale_fill_gradientn(colours=colorblind_pal()(8)[-1], guide=guide_colourbar(raster=FALSE, title='Power'))+
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())
ggsave_tikz(file.path(plot_dir,'wavelet-spectrum-climate.tex'),p_hmm_climate,width=pc2in(20),height=pc2in(15),pointsize=10)

######################################################################################################
######################################################################################################
printf('HMM Default Spectrum...')
p_hmm_default <- ggplot(local_spec_med[kind=='HMG'])+
    geom_tile(aes(x=time,y=scale,fill=power))+
    geom_line(data=data.frame(x=coi_time,y=coi),aes(x=x,y=y),size=1)+
    theme_bw()+xlab('Year')+ylab('Scale')+
    scale_y_continuous(trans=log2_trans(),limits=c(1.99,65))+
    scale_fill_gradientn(colours=colorblind_pal()(8)[-1], guide=guide_colourbar(raster=FALSE, title='Power'))+
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())
ggsave_tikz(file.path(plot_dir,'wavelet-spectrum-default.tex'),p_hmm_default,width=pc2in(20),height=pc2in(15),pointsize=10)

######################################################################################################
######################################################################################################
printf('AR Spectrum...')
p_hmm_ar <- ggplot(local_spec_med[kind=='AR'])+
    geom_tile(aes(x=time,y=scale,fill=power))+
    geom_line(data=data.frame(x=coi_time,y=coi),aes(x=x,y=y),size=1)+
    theme_bw()+xlab('Year')+ylab('Scale')+
    scale_y_continuous(trans=log2_trans(),limits=c(1.99,65))+
    scale_fill_gradientn(colours=colorblind_pal()(8)[-1], guide=guide_colourbar(raster=FALSE, title='Power'))+
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())
ggsave_tikz(file.path(plot_dir,'wavelet-spectrum-ar.tex'),p_hmm_ar,width=pc2in(20),height=pc2in(15),pointsize=10)

######################################################################################################
######################################################################################################
printf('Observed Spectrum...')
p4 <- ggplot(spec_obs)+
    geom_tile(aes(x=time,y=scale,fill=power))+
    geom_line(data=data.frame(x=coi_time,y=coi),aes(x=x,y=y),size=1)+
    theme_bw()+xlab('Year')+ylab('Scale')+
    scale_y_continuous(trans=log2_trans(),limits=c(1.99,65))+
    scale_fill_gradientn(colours=colorblind_pal()(8)[-1], guide=guide_colourbar(raster=FALSE, title='Power'))+
    theme(panel.grid.minor=element_blank(), panel.grid.major=element_blank())
ggsave_tikz(file.path(plot_dir,'wavelet-spectrum-obs.tex'),p4,width=pc2in(20),height=pc2in(15),pointsize=10)

tikz(file.path(plot_dir,'wavelet-spectrum-all.tex'), width=pc2in(39), height=pc2in(30), pointsize=10)
multiplot(
    p4+ggtitle('(a)'), 
    p_hmm_imposed+ggtitle('(c)'),
    p_hmm_default+ggtitle('(b)'),
    p_hmm_climate+ggtitle('(d)'),
    cols=2)
dev.off()