# Run EM for HMM with splines or trigonometric functions in the state process (for diel variation)
# for one of the 8 elephants
# with 20 starting values each

## the R scripts contain the R code for reproducing all materials presented in the manuscript
## the code was tested in R (version 4.2.1) on Windows 10 with packages...
# ...dplyr (version 1.1.0)
# ...mgcv (version 1.8-41)
# ...circular (version 0.4-95)


library(dplyr)
library(mgcv)
library(circular) # for vonMises distribution
source("HMMFunctions.R")

# load data
load("elephant_data_prepped.RData")
elphID <- 5 # choose one elephant for this model
data <- data[data$animalID == elphID, ]

N <- 2 # number of states (although EM is not written for more than 2 states as for now)

# define the different starting values to find global maximum ----
S <- 20 # number of random starts

stepMean0s <- matrix(NA, nrow = S, ncol = N)
stepSD0s <- matrix(NA, nrow = S, ncol = N)
zeroMass0s <- matrix(NA, nrow = S, ncol = N)
angleCon0s <- matrix(NA, nrow = S, ncol = N)

set.seed(2510)

stepMean0s[, 1] <- runif(S, 0.001, 0.8)
stepMean0s[, 2] <- runif(S, 0.801, 3)
stepSD0s[, 1] <- runif(S, 0.001, 0.8)
stepSD0s[, 2] <- runif(S, 0.801, 3)

angleCon0s[, 1] <- runif(S, 0.0001, 0.9)
angleCon0s[, 2] <- runif(S, 0.9001, 5)

# for Gamma, the same starting values are used (which result in same t.p.m. across time)
Gamma0 <- diag(N)
Gamma0[!Gamma0] <- rep(0.1, (N * (N - 1)))
Gamma0 <- Gamma0 / rowSums(Gamma0)
Gamma0 <- array(Gamma0, dim = c(N, N, 24))


tod <- data$tod
step <- data$step
angle <- data$angle
trackID <- data$ID
ID <- data$animalID # all the variables that are needed

# run with splines -------
mods_splines <- list()
llks_splines <- durations_splines <- rep(NA, S)

for (s in 1:S) {
  ## sensible starting values for the estimation:
  stepMean0 <- stepMean0s[s, ] # step mean
  stepSD0 <- stepSD0s[s, ] # step SD
  angleMean0 <- rep(0, N) # angle mean
  angleCon0 <- angleCon0s[s, ] # angle concentration

  stepShape0 <- (stepMean0^2) / (stepSD0^2)
  stepRate0 <- stepMean0 / (stepSD0^2)

  # with splines:
  duration <- system.time(mods_splines[[s]] <- tryCatch(
    {
      EM(tod, step, angle, trackID, ID, N,
        stepShape = stepShape0, stepRate = stepRate0,
        anglemean = angleMean0, anglecon = angleCon0,
        Gamma = Gamma0,
        maxiter = 100, tol = 0.01
      )
    },
    error = function(err) {
      print(err)
    }
  ))
  durations_splines[s] <- duration["elapsed"]
  llks_splines[s] <- if (is.na(mods_splines[[s]][3])) {
    NA
  } else {
    if
    (is.null(mods_splines[[s]]$mllk)) {
      NA
    } else {
      mods_splines[[s]]$mllk
    }
  }
}


llks_splines
bestno <- which(llks_splines == max(llks_splines, na.rm = T))
mod_splines <- mods_splines[[bestno]]
duration <- durations_splines[bestno]

save(
  list = c("mod_splines", "duration"),
  file = paste0("mod_splines_ID", elphID, ".RData")
)


# run for K = 1, ..., 7 with trigonometric functions ----
for (K in 1:7) {
  mods_trigon <- list()
  llks_trigon <- durations_trigon <- rep(NA, S)

  for (s in 1:S) {
    ## sensible starting values for the estimation:
    stepMean0 <- stepMean0s[s, ] # step mean
    stepSD0 <- stepSD0s[s, ] # step SD
    angleMean0 <- rep(0, N) # angle mean
    angleCon0 <- angleCon0s[s, ] # angle concentration

    stepShape0 <- (stepMean0^2) / (stepSD0^2)
    stepRate0 <- stepMean0 / (stepSD0^2)

    # with splines:
    duration <- system.time(mods_trigon[[s]] <- tryCatch(
      {
        EM_trigon(tod, step, angle, trackID, ID, N,
          stepShape = stepShape0, stepRate = stepRate0,
          anglemean = angleMean0, anglecon = angleCon0,
          Gamma = Gamma0, TrigPairs = K,
          maxiter = 100, tol = 0.01
        )
      },
      error = function(err) {
        print(err)
      }
    ))
    durations_trigon[s] <- duration["elapsed"]
    llks_trigon[s] <- if (is.na(mods_trigon[[s]][3])) {
      NA
    } else {
      if
      (is.null(mods_trigon[[s]]$mllk)) {
        NA
      } else {
        mods_trigon[[s]]$mllk
      }
    }
  }

  bestno <- which(llks_trigon == max(llks_trigon, na.rm = T))
  assign(paste0("mod_trigon", K), mods_trigon[[bestno]])
  duration <- durations_trigon[bestno]

  save(
    list = c(paste0("mod_trigon", K), "duration"),
    file = paste0("mod_trigon", K, "_ID", elphID, "test.RData")
  )
}

# additionally run assuming homogeneous Markov chain ------

mods_trigon0 <- list()
llks_trigon0 <- durations_trigon0 <- rep(NA, S)

for (s in 1:S) {
  ## sensible starting values for the estimation:
  stepMean0 <- stepMean0s[s, ] # step mean
  stepSD0 <- stepSD0s[s, ] # step SD
  angleMean0 <- rep(0, N) # angle mean
  angleCon0 <- angleCon0s[s, ] # angle concentration

  stepShape0 <- (stepMean0^2) / (stepSD0^2)
  stepRate0 <- stepMean0 / (stepSD0^2)

  # with splines:
  duration <- system.time(mods_trigon0[[s]] <- tryCatch(
    {
      EM_trigon(tod, step, angle, trackID, ID, N,
        stepShape = stepShape0, stepRate = stepRate0,
        anglemean = angleMean0, anglecon = angleCon0,
        Gamma = Gamma0, TrigPairs = 0,
        maxiter = 100, tol = 0.01
      )
    },
    error = function(err) {
      print(err)
    }
  ))
  durations_trigon0[s] <- duration["elapsed"]
  llks_trigon0[s] <- if (is.na(mods_trigon0[[s]][3])) {
    NA
  } else {
    if
    (is.null(mods_trigon0[[s]]$mllk)) {
      NA
    } else {
      mods_trigon0[[s]]$mllk
    }
  }
}


llks_trigon0
bestno <- which(llks_trigon0 == max(llks_trigon0, na.rm = T))
mod0 <- mods_trigon0[[bestno]]
duration <- durations_trigon0[bestno]

save(
  list = c("mod0", "duration"),
  file = paste0("mod0_ID", elphID, ".RData")
)