#' "Flexible modelling of diel and other periodic variation in hidden Markov models"
#' HMMs with different ways of modelling the diel variation in the state switching dynamics
#' plots for manuscript

rm(list = ls())
# load libraries, source code, ... ------
library(dplyr) # version 1.1.0
library(mgcv) # version 1.8-41
library(circular) # version 0.4-95

source("HMMFunctions.R")
source("plotFunctions.R")

cbPalette <- c("#999999", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7")
cbPalette_alpha <- rgb(col2rgb(cbPalette)[1, ],
  col2rgb(cbPalette)[2, ],
  col2rgb(cbPalette)[3, ],
  alpha = 50, maxColorValue = 255
)

# load data ----------------------------------------------

load("elephant_data_prepped.RData")
data <- data[data$animalID == 5, ]

N <- 2
cycle <- 24
no_tp <- 240
var_x <- seq(cycle / no_tp, cycle, length = no_tp)


# load and name models
load(file = "mod_splines_ID5.RData")
for (i in 1:7) {
  load(file = paste0("mod_trigon", i, "_ID5.RData"))
}
load(file = "mod0_ID5.RData")

# get Viterbi decoded state sequences -------------------------

vit_states_splines <- getVit(mod = mod_splines, data = data, cycle = cycle)
for (i in 1:7) {
  assign(paste0("vit_states_trigon", i), getVit(mod = get(paste0("mod_trigon", i)), data = data, cycle = cycle))
}
vit_states_0 <- getVit0(mod = mod0, data = data, cycle = cycle)


# proportions of time spent in each state
props_splines <- table(vit_states_splines) / length(vit_states_splines)
for (i in 1:7) {
  assign(paste0("props_trigon", i), table(get(paste0("vit_states_trigon", i))) / length(get(paste0("vit_states_trigon", i))))
}
props_0 <- table(vit_states_0) / length(vit_states_0)


# for each hour
vit_props <- group_by(data.frame(
  states_splines = vit_states_splines,
  states_trigon1 = vit_states_trigon1,
  states_trigon2 = vit_states_trigon2,
  states_trigon3 = vit_states_trigon3,
  states_trigon4 = vit_states_trigon4,
  states_trigon5 = vit_states_trigon5,
  states_trigon6 = vit_states_trigon6,
  states_trigon7 = vit_states_trigon7,
  states_0 = vit_states_0,
  tod = data$tod
), tod) %>%
  summarize(
    prop2_splines = mean(states_splines, na.rm = TRUE) - 1,
    prop2_trigon1 = mean(states_trigon1, na.rm = TRUE) - 1,
    prop2_trigon2 = mean(states_trigon2, na.rm = TRUE) - 1,
    prop2_trigon3 = mean(states_trigon3, na.rm = TRUE) - 1,
    prop2_trigon4 = mean(states_trigon4, na.rm = TRUE) - 1,
    prop2_trigon5 = mean(states_trigon5, na.rm = TRUE) - 1,
    prop2_trigon6 = mean(states_trigon6, na.rm = TRUE) - 1,
    prop2_trigon7 = mean(states_trigon7, na.rm = TRUE) - 1,
    prop2_0 = mean(states_0, na.rm = TRUE) - 1
  )

# get t.p.m.s ---------
Gamma_splines <- getGamma(mod = mod_splines, cycle = cycle)
for (i in 1:7) {
  assign(paste0("Gamma_trigon", i), getGamma(mod = get(paste0("mod_trigon", i)), cycle = cycle))
}


# confidence intervals for t.p.m.s
probs12Conf_splines <- getConfint(gammod = mod_splines$mods$gammod1, cycle = cycle)
probs21Conf_splines <- getConfint(gammod = mod_splines$mods$gammod2, cycle = cycle)

# probs12Conf_trigon1 <- getConfint(gammod = mod_trigon1$mods$gammod1, cycle = cycle)
# probs21Conf_trigon1 <- getConfint(gammod = mod_trigon1$mods$gammod2, cycle = cycle)
#
# probs12Conf_trigon2 <- getConfint(gammod = mod_trigon2$mods$gammod1, cycle = cycle)
# probs21Conf_trigon2 <- getConfint(gammod = mod_trigon2$mods$gammod2, cycle = cycle)
#
# probs12Conf_trigon3 <- getConfint(gammod = mod_trigon3$mods$gammod1, cycle = cycle)
# probs21Conf_trigon3 <- getConfint(gammod = mod_trigon3$mods$gammod2, cycle = cycle)

# plot time spent in states -------------------------------------------
pdf(file = "timespent.pdf", width = 9, height = 5)
# setEPS()
# postscript("timespent.eps")

par(mfrow = c(1, 1), mar = c(5.1, 4.1, 4.1, 4))
plot(
  x = seq(1, cycle), y = vit_props$prop2_splines,
  ylim = c(0, 1), type = "l", lwd = 1, col = cbPalette_alpha[2],
  xlab = "time of day", ylab = "prop. of time in state 2", bty = "n", xaxt = "n", cex.lab = 1.2
)
lines(x = seq(1, cycle), y = vit_props$prop2_trigon1, col = cbPalette_alpha[6], lwd = 1)
lines(x = seq(1, cycle), y = vit_props$prop2_trigon2, col = cbPalette_alpha[8], lwd = 1)
lines(x = seq(1, cycle), y = vit_props$prop2_trigon3, col = cbPalette_alpha[4], lwd = 1)
lines(x = seq(1, cycle), y = vit_props$prop2_0, col = cbPalette_alpha[7], lwd = 1)


points(x = seq(1, cycle), y = vit_props$prop2_splines, col = cbPalette[2], pch = 20)
points(x = seq(1, cycle), y = vit_props$prop2_trigon1, col = cbPalette[6], pch = 20)
points(x = seq(1, cycle), y = vit_props$prop2_trigon2, col = cbPalette[8], pch = 20)
points(x = seq(1, cycle), y = vit_props$prop2_trigon3, col = cbPalette[4], pch = 20)
points(x = seq(1, cycle), y = vit_props$prop2_0, col = cbPalette[7], pch = 20)


axis(1, seq(0, cycle, by = 4))
par(xpd = TRUE)
legend("topright", c("cyclic P-splines", "trigonometric (K=1)", "trigonometric (K=2)", "trigonometric (K=3)", "no diel variation modelled"),
  col = c(cbPalette[2], cbPalette[6], cbPalette[8], cbPalette[4], cbPalette[7]), lty = c(1, 1), bty = "n", inset = c(-0.1, -0.05)
)

dev.off()

# plot t.p.m.s -------------------------------------------

pdf(file = "tpms.pdf", width = 9, height = 5)
# setEPS()
# postscript("tpms.eps")

par(mfrow = c(1, 2), mar = c(5.1, 4.7, 2.1, 2))
plot(
  x = var_x, y = Gamma_splines[1, 2, ],
  ylim = c(0, 0.55), type = "l", lwd = 1, col = cbPalette[2],
  xlab = "time of day t", ylab = expression(gamma[12]^(t)), bty = "n", xaxt = "n", cex.lab = 1.2
)
axis(1, seq(0, cycle, by = 4))

lines(x = var_x, y = Gamma_trigon1[1, 2, ], col = cbPalette[6], lwd = 1)
lines(x = var_x, y = Gamma_trigon2[1, 2, ], col = cbPalette[8], lwd = 1)
lines(x = var_x, y = Gamma_trigon3[1, 2, ], col = cbPalette[4], lwd = 1)

polygon(c(var_x, rev(var_x)), c(probs12Conf_splines[, 1], rev(probs12Conf_splines[, 3])),
  col = cbPalette_alpha[2], border = NA
)

par(xpd = TRUE)
legend("topright", c("cyclic P-splines", "trigonometric (K=1)", "trigonometric (K=2)", "trigonometric (K=3)"),
  col = c(cbPalette[2], cbPalette[6], cbPalette[8], cbPalette[4]), lty = c(1, 1), bty = "n", inset = c(-0.1, 0)
)

par(mar = c(5.1, 4.7, 2.1, 2))
plot(
  x = var_x, y = Gamma_splines[2, 1, ],
  ylim = c(0, 0.55), type = "l", lwd = 1, col = cbPalette[2],
  xlab = "time of day t", ylab = expression(gamma[21]^(t)), bty = "n", xaxt = "n", cex.lab = 1.2
)
axis(1, seq(0, cycle, by = 4))

lines(x = var_x, y = Gamma_trigon1[2, 1, ], col = cbPalette[6], lwd = 1)
lines(x = var_x, y = Gamma_trigon2[2, 1, ], col = cbPalette[8], lwd = 1)
lines(x = var_x, y = Gamma_trigon3[2, 1, ], col = cbPalette[4], lwd = 1)

polygon(c(var_x, rev(var_x)), c(probs21Conf_splines[, 1], rev(probs21Conf_splines[, 3])),
  col = cbPalette_alpha[2], border = NA
)


dev.off()

# plot state dependent distributions  ----------
pdf(file = "statedepdists_splines.pdf", width = 11, height = 5)

par(mfrow = c(1, 2))
hist(data$step,
  breaks = 30, freq = FALSE, xlim = c(0, 3.5),
  xlab = "step lengths in km", ylab = "density", main = "", bor = "white", cex.lab = 1.2
)
z <- seq(min(data$step, na.rm = T), max(data$step, na.rm = T), length = 200)
lines(z, dgamma(z, shape = mod_splines$stepShape[1], rate = mod_splines$stepRate[1]) * props_splines[1], col = cbPalette[3], lwd = 3)
# points(0, zeroMass[1], col="orange",lwd=0.5)
lines(z, dgamma(z, shape = mod_splines$stepShape[2], rate = mod_splines$stepRate[2]) * props_splines[2], col = cbPalette[7], lwd = 3)
# points(0, zeroMass[2], col="mediumseagreen",lwd=0.5)
lines(z, (dgamma(z, shape = mod_splines$stepShape[1], rate = mod_splines$stepRate[1]) * props_splines[1] +
  dgamma(z, shape = mod_splines$stepShape[2], rate = mod_splines$stepRate[2]) * props_splines[2]), col = "black", lwd = 1, lty = 2)
# lines(density(data$step, na.rm=T),lwd=3)
legend("topright", c("state 1 (encamped)", "state 2 (exploratory)", "estimated mixture"),
  col = c(cbPalette[3], cbPalette[7], "black"),
  lwd = c(3, 3, 1), lty = c(1, 1, 2), bty = "n"
)


hist(data$angle,
  breaks = 20, freq = FALSE, xlim = c(-pi, pi),
  xlab = "turning angles in radians", ylab = "density", main = "", bor = "white", cex.lab = 1.2
)
z <- seq(min(data$angle, na.rm = T), max(data$angle, na.rm = T), length = 200)
lines(z, dvonmises(z, mu = mod_splines$anglemean[1], kappa = mod_splines$anglecon[1]) * props_splines[1], col = cbPalette[3], lwd = 3)
lines(z, dvonmises(z, mu = mod_splines$anglemean[2], kappa = mod_splines$anglecon[2]) * props_splines[2], col = cbPalette[7], lwd = 3)
lines(z, (dvonmises(z, mu = mod_splines$anglemean[1], kappa = mod_splines$anglecon[1]) * props_splines[1] +
  dvonmises(z, mu = mod_splines$anglemean[2], kappa = mod_splines$anglecon[2]) * props_splines[2]), col = "black", lwd = 1, lty = 2)


dev.off()





# information criteria ----

# no. of parameters in state-dep process : N * 4 #(stepShape, stepRate, anglemu, anglecon)
# no. of parameters in state process : N*(N-1) * (1+K*2) #(intercept, K*2 betas)
InfoCriteria <- data.frame(AIC = rep(NA, 9), BIC = NA, edf = NA, loglik = NA)
rownames(InfoCriteria) <- c("trigon1", "trigon2", "trigon3", "trigon4", "trigon5", "trigon6", "trigon7", "splines", "homogeneous")


for (K in 1:7) {
  InfoCriteria$edf[K] <- N * 4 + N * (N - 1) * (1 + 2 * K)
  InfoCriteria$loglik[K] <- get(paste0("mod_trigon", K))$mllk
}

InfoCriteria$loglik[2] <- mod_trigon2$mllk
InfoCriteria$loglik[3] <- mod_trigon3$mllk
InfoCriteria$loglik[4] <- mod_trigon4$mllk
InfoCriteria$loglik[5] <- mod_trigon5$mllk
InfoCriteria$loglik[6] <- mod_trigon6$mllk
InfoCriteria$loglik[7] <- mod_trigon7$mllk

InfoCriteria["homogeneous", ]$loglik <- mod0$mllk
InfoCriteria["homogeneous", ]$edf <- N * 4 + N * (N - 1)

InfoCriteria["splines", ]$edf <- N * 4 + sum(mod_splines$mods$gammod1$edf) + sum(mod_splines$mods$gammod2$edf)
InfoCriteria["splines", ]$loglik <- mod_splines$mllk

InfoCriteria$AIC <- (-2) * InfoCriteria$loglik + 2 * InfoCriteria$edf
InfoCriteria$BIC <- (-2) * InfoCriteria$loglik + log(nrow(data)) * InfoCriteria$edf

InfoCriteria