rm(list = ls()); gc()

suppressPackageStartupMessages({
  library(dplyr)
  library(tidyr)
  library(ggplot2)
  library(igraph)
  library(amt)
  library(hmmSSF) #Installation link: devtools::install_github("NJKlappstein/hmmSSF
  library(patchwork)
})

# ----------------------------------------------------------------------------
# OUTPUT
# ----------------------------------------------------------------------------
set.seed(123)
output_dir <- "hierarchical_hmmssf_results"
if(!dir.exists(output_dir)) dir.create(output_dir, recursive = TRUE)

.save_png <- function(p, filename, w = 8, h = 5, dpi = 200){
  ggsave(file.path(output_dir, filename), p, width = w, height = h, dpi = dpi)
}

scale_color_viridis_discrete_safe <- function(..., option = "D", name = waiver()){
  if(requireNamespace("viridis", quietly = TRUE)){
    viridis::scale_color_viridis(discrete = TRUE, option = option, name = name, ...)
  } else {
    scale_color_discrete(name = name, ...)
  }
}
scale_fill_viridis_safe <- function(..., option = "magma", name = waiver()){
  if(requireNamespace("viridis", quietly = TRUE)){
    viridis::scale_fill_viridis(option = option, name = name, ...)
  } else {
    scale_fill_continuous(name = name, ...)
  }
}

.get_state_palette <- function(){
  if(requireNamespace("viridis", quietly = TRUE)){
    cols <- viridis::viridis(3, option = "D", begin = 0.10, end = 0.90)
  } else if(requireNamespace("viridisLite", quietly = TRUE)){
    cols <- viridisLite::viridis(3, begin = 0.10, end = 0.90)
  } else {
    cols <- c("#1b9e77", "#d95f02", "#7570b3")
  }
  names(cols) <- c("WITHIN_SITE","BETWEEN_SITE","BETWEEN_CAMP")
  cols
}

# ----------------------------------------------------------------------------
# ANGLE HELPERS
# ----------------------------------------------------------------------------
wrap_pi <- function(a) (a + pi) %% (2*pi) - pi
blend_bearing <- function(prev, target, w){
  x <- (1-w)*cos(prev) + w*cos(target)
  y <- (1-w)*sin(prev) + w*sin(target)
  atan2(y, x)
}

# ----------------------------------------------------------------------------
# RANDOM TARGET IN SITE + PERSISTENT MOVEMENT
# ----------------------------------------------------------------------------
sample_point_in_SITE <- function(SITE_id, SITE_data){
  cx <- SITE_data$SITE_x[SITE_id]
  cy <- SITE_data$SITE_y[SITE_id]
  r  <- sqrt(runif(1)) * SITE_data$SITE_size[SITE_id]
  a  <- runif(1, 0, 2*pi)
  c(cx + r*cos(a), cy + r*sin(a))
}

move_toward_point_persistent <- function(current_x, current_y, prev_bearing,
                                         target_x, target_y,
                                         step_length, w_blend, sd_angle,
                                         landscape_extent){
  target_bearing <- atan2(target_y - current_y, target_x - current_x)
  base_angle <- blend_bearing(prev_bearing, target_bearing, w = w_blend)
  theta <- wrap_pi(base_angle + rnorm(1, 0, sd_angle))
  new_x <- current_x + step_length * cos(theta)
  new_y <- current_y + step_length * sin(theta)
  new_x <- pmin(pmax(new_x, landscape_extent[1]), landscape_extent[2])
  new_y <- pmin(pmax(new_y, landscape_extent[3]), landscape_extent[4])
  list(theta = theta, new_x = new_x, new_y = new_y)
}

# ----------------------------------------------------------------------------
# AUTOCORRELATED (0,1) GENERATOR
# ----------------------------------------------------------------------------
generate_autocorrelated_01 <- function(n = 100, phi = 0.95, mu = 0.5, sigma = 0.4) {
  logit    <- function(p) log(p / (1 - p))
  invlogit <- function(x) 1 / (1 + exp(-x))
  mu_logit <- logit(mu)
  x <- numeric(n)
  x[1] <- rnorm(1, mu_logit, sigma)
  for(i in 2:n){
    x[i] <- mu_logit + phi * (x[i-1] - mu_logit) + rnorm(1, 0, sigma)
  }
  invlogit(x)
}

# ----------------------------------------------------------------------------
# 0) NETWORK
# ----------------------------------------------------------------------------
create_simplified_hierarchical_network <- function(
    n_camps = 3,
    SITEes_per_camp = c(6, 6, 6),
    camp_xy = data.frame(
      camp_id = 1:3,
      x = c(20, 140, 260),
      y = c(20, 160, 20)
    )
){
  stopifnot(length(SITEes_per_camp) == n_camps)
  camps  <- camp_xy
  SITEes <- data.frame()
  SITE_id_counter <- 1L
  
  for(i in 1:n_camps){
    m    <- SITEes_per_camp[i]
    ang  <- seq(0, 2*pi, length.out = m + 1)[-1]
    radii <- runif(m, 10, 18)
    camp_SITEes <- data.frame(
      SITE_id   = SITE_id_counter:(SITE_id_counter + m - 1L),
      camp_id   = i,
      SITE_x    = camps$x[i] + radii * cos(ang),
      SITE_y    = camps$y[i] + radii * sin(ang),
      SITE_size = runif(m, 4.5, 6.5)
    )
    SITEes <- rbind(SITEes, camp_SITEes)
    SITE_id_counter <- SITE_id_counter + m
  }
  
  n_SITEes <- nrow(SITEes)
  adj <- matrix(0, n_SITEes, n_SITEes)
  
  for(i in 1:(n_SITEes-1)){
    for(j in (i+1):n_SITEes){
      if(SITEes$camp_id[i] == SITEes$camp_id[j]){
        d <- sqrt((SITEes$SITE_x[i] - SITEes$SITE_x[j])^2 +
                    (SITEes$SITE_y[i] - SITEes$SITE_y[j])^2)
        if(d < 25) adj[i, j] <- adj[j, i] <- 1
      }
    }
  }
  
  p1 <- which(SITEes$camp_id == 1)
  p2 <- which(SITEes$camp_id == 2)
  p3 <- which(SITEes$camp_id == 3)
  
  if(length(p1) > 0 && length(p2) > 0){
    a <- p1[which.max(SITEes$SITE_x[p1])]
    b <- p2[which.min(SITEes$SITE_x[p2])]
    adj[a, b] <- adj[b, a] <- 1
  }
  if(length(p2) > 0 && length(p3) > 0){
    a <- p2[which.max(SITEes$SITE_x[p2])]
    b <- p3[which.min(SITEes$SITE_x[p3])]
    adj[a, b] <- adj[b, a] <- 1
  }
  
  g <- igraph::graph_from_adjacency_matrix(adj, mode = "undirected", diag = FALSE)
  
  coords <- as.matrix(SITEes[, c("SITE_x","SITE_y")])
  el <- igraph::as_edgelist(g)
  if(nrow(el) > 0){
    w <- apply(el, 1, function(eij){
      i <- as.integer(eij[1]); j <- as.integer(eij[2])
      sqrt((coords[i,1]-coords[j,1])^2 + (coords[i,2]-coords[j,2])^2)
    })
    E(g)$weight <- w
  } else {
    E(g)$weight <- numeric(0)
  }
  
  if(ecount(g) > 0){
    el2   <- igraph::as_edgelist(g)
    etype <- apply(el2, 1, function(eij){
      i <- as.integer(eij[1]); j <- as.integer(eij[2])
      if(SITEes$camp_id[i] == SITEes$camp_id[j]) "within_camp" else "between_camps"
    })
    E(g)$edge_type <- etype
  } else {
    E(g)$edge_type <- character(0)
  }
  
  list(camps = camps, SITEes = SITEes, g = g)
}

net        <- create_simplified_hierarchical_network()
camps      <- net$camps
SITE_data  <- net$SITEes
g          <- net$g
total_SITEes <- nrow(SITE_data)

neighbors_list    <- lapply(1:total_SITEes, function(i) as.integer(neighbors(g, i)))
names(neighbors_list) <- as.character(1:total_SITEes)
camp_SITE_indices <- lapply(1:nrow(camps), function(cc) which(SITE_data$camp_id == cc))

xmin <- min(SITE_data$SITE_x, camps$x) - 25
xmax <- max(SITE_data$SITE_x, camps$x) + 25
ymin <- min(SITE_data$SITE_y, camps$y) - 25
ymax <- max(SITE_data$SITE_y, camps$y) + 25
landscape_extent <- c(xmin, xmax, ymin, ymax)

cat("Network created: camps =", nrow(camps), " SITEes =", total_SITEes,
    " edges =", ecount(g), "\n")

edge_df <- as_data_frame(g, what = "edges") %>%
  mutate(edge_type = E(g)$edge_type)
vert_df <- SITE_data %>% mutate(name = as.character(SITE_id))
p_net <- ggplot() +
  geom_segment(data = edge_df,
               aes(x    = SITE_data$SITE_x[from], y    = SITE_data$SITE_y[from],
                   xend = SITE_data$SITE_x[to],   yend = SITE_data$SITE_y[to],
                   linetype = edge_type),
               linewidth = 0.7, alpha = 0.7) +
  geom_point(data = vert_df,
             aes(x = SITE_x, y = SITE_y, fill = factor(camp_id)),
             shape = 21, size = 3.2, color = "grey20") +
  coord_equal() + theme_minimal() +
  labs(title = "SITE network (3 camps)", x = "x", y = "y", linetype = "edge")
.save_png(p_net, "00_network.png", 7, 5)

# ----------------------------------------------------------------------------
# SIMULATION SETTINGS
# ----------------------------------------------------------------------------
n_steps     <- 500
state_names <- c("WITHIN_SITE","BETWEEN_SITE","BETWEEN_CAMP")
n_states    <- 3

# ----------------------------------------------------------------------------
# 1) mu FIELD
# ----------------------------------------------------------------------------
grid_n  <- 180
x_range <- seq(landscape_extent[1], landscape_extent[2], length.out = grid_n)
y_range <- seq(landscape_extent[3], landscape_extent[4], length.out = grid_n)

X <- matrix(rep(x_range, times = grid_n), nrow = grid_n, ncol = grid_n)
Y <- t(matrix(rep(y_range, times = grid_n), nrow = grid_n, ncol = grid_n))

SITE_id_raster <- matrix(0L,  nrow = grid_n, ncol = grid_n)
best_d2        <- matrix(Inf, nrow = grid_n, ncol = grid_n)

for(k in seq_len(nrow(SITE_data))){
  dx     <- X - SITE_data$SITE_x[k]
  dy     <- Y - SITE_data$SITE_y[k]
  d2     <- dx*dx + dy*dy
  inside <- d2 <= (SITE_data$SITE_size[k]^2)
  better <- inside & (d2 < best_d2)
  SITE_id_raster[better] <- k
  best_d2[better]        <- d2[better]
}
rm(X, Y, best_d2); gc()

KSITE <- nrow(SITE_data)
A_k   <- as.numeric(tabulate(SITE_id_raster[SITE_id_raster > 0L], nbins = KSITE))
if(sum(A_k) == 0) stop("No raster cells fall inside SITE disks.")

phi_mu   <- 0.95
mu_mean  <- 0.55
sigma_mu <- 0.35

u_SITE <- sapply(1:KSITE, function(k){
  generate_autocorrelated_01(n = n_steps, phi = phi_mu, mu = mu_mean, sigma = sigma_mu)
})
if(is.null(dim(u_SITE))) u_SITE <- matrix(u_SITE, nrow = n_steps, ncol = KSITE)
if(nrow(u_SITE) != n_steps) u_SITE <- t(u_SITE)

eps    <- 1e-8
u_SITE <- pmin(pmax(u_SITE, eps), 1 - eps)

mu_raw  <- u_SITE / (1 - u_SITE)
mu_SITE <- mu_raw

Ncells_in <- sum(A_k)
for(t in 1:n_steps){
  denom_in <- sum(A_k * mu_raw[t, ])
  s <- if(!is.finite(denom_in) || denom_in <= 0) 1 else (Ncells_in / denom_in)
  mu_SITE[t, ] <- s * mu_raw[t, ]
}
mu_SITE[!is.finite(mu_SITE)] <- 1

.get_raster_idx <- function(x, y){
  xi <- round((x - landscape_extent[1]) /
                (landscape_extent[2] - landscape_extent[1]) * (grid_n-1)) + 1L
  yi <- round((y - landscape_extent[3]) /
                (landscape_extent[4] - landscape_extent[3]) * (grid_n-1)) + 1L
  xi[!is.finite(xi) | is.na(xi)] <- 1L
  yi[!is.finite(yi) | is.na(yi)] <- 1L
  xi <- pmin(pmax(as.integer(xi), 1L), grid_n)
  yi <- pmin(pmax(as.integer(yi), 1L), grid_n)
  cbind(xi, yi)
}

get_mu_time <- function(x, y, t){
  if(!is.finite(x) || !is.finite(y)) return(1)
  t   <- as.integer(t); t <- pmin(pmax(t, 1L), n_steps)
  idx <- .get_raster_idx(x, y)
  pid <- SITE_id_raster[idx[1], idx[2]]
  if(is.na(pid) || pid <= 0L) return(1)
  mu_SITE[t, as.integer(pid)]
}

get_mu_time_vec <- function(x, y, t){
  x <- as.numeric(x); y <- as.numeric(y)
  n  <- length(x)
  out <- rep(1, n)
  ok  <- is.finite(x) & is.finite(y)
  if(!any(ok)) return(out)
  tt <- if(length(t) == 1L) rep(as.integer(t), n) else as.integer(t)
  tt[!is.finite(tt) | is.na(tt)] <- 1L
  tt <- pmin(pmax(tt, 1L), n_steps)
  idx  <- .get_raster_idx(x[ok], y[ok])
  pid  <- SITE_id_raster[cbind(idx[,1], idx[,2])]
  good <- is.finite(pid) & !is.na(pid) & (pid > 0L)
  if(any(good)){
    pid2     <- pmin(pmax(as.integer(pid[good]), 1L), KSITE)
    idx_good <- which(ok)[good]
    out[idx_good] <- mu_SITE[cbind(tt[idx_good], pid2)]
  }
  out
}

get_SITE_id_vec <- function(x, y){
  x <- as.numeric(x); y <- as.numeric(y)
  n   <- length(x)
  out <- rep(0L, n)
  ok  <- is.finite(x) & is.finite(y)
  if(!any(ok)) return(out)
  idx <- .get_raster_idx(x[ok], y[ok])
  pid <- SITE_id_raster[cbind(idx[,1], idx[,2])]
  pid[!is.finite(pid) | is.na(pid)] <- 0L
  out[which(ok)] <- as.integer(pid)
  out
}
get_SITE_id <- function(x, y){
  as.integer(get_SITE_id_vec(x, y)[1])
}

# ----------------------------------------------------------------------------
# 2) RESOURCE DYNAMICS
# ----------------------------------------------------------------------------
rv       <- 0.15
Kcap     <- 100
a_attack <- 0.3
h_handle <- 0.03
dt       <- 1

g_of_V <- function(V){
  V <- pmax(V, 0)
  (a_attack * V) / (1 + a_attack * h_handle * V)
}
V_SITE <- runif(KSITE, 0.6*Kcap, 1.0*Kcap)
V_hist <- matrix(NA_real_, nrow = n_steps, ncol = KSITE)

# ----------------------------------------------------------------------------
# 3) LEAVING RULE PARAMETERS
# ----------------------------------------------------------------------------
T_const <- 1

resource_rel_thresh  <- 0.85
camp_mu_thresh       <- 1.10
local_mu_thresh      <- 1.10
p_camp_switch_min    <- 0.15
p_camp_switch_max    <- 0.70
p_spont_between_SITE <- 0.10

# ----------------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------------
nearest_camp_id <- function(curr_camp, camps){
  cx    <- camps$x[curr_camp]; cy <- camps$y[curr_camp]
  other <- setdiff(camps$camp_id, curr_camp)
  d2    <- (camps$x[other] - cx)^2 + (camps$y[other] - cy)^2
  other[which.min(d2)]
}

next_SITE_on_route <- function(g, from_SITE, to_SITE){
  if(from_SITE == to_SITE) return(from_SITE)
  sp <- igraph::shortest_paths(g, from = from_SITE, to = to_SITE,
                               weights = igraph::E(g)$weight)$vpath[[1]]
  sp <- as.integer(sp)
  if(length(sp) <= 1) return(from_SITE)
  sp[2]
}
safe_next_hop <- function(g, from_SITE, to_SITE, neighbors_list){
  hop <- try(next_SITE_on_route(g, from_SITE, to_SITE), silent = TRUE)
  if(inherits(hop, "try-error") || is.na(hop) || length(hop) == 0){
    neigh <- neighbors_list[[as.character(from_SITE)]]
    if(length(neigh) > 0) return(sample(neigh, 1))
    return(from_SITE)
  }
  hop
}

pick_destination_neighbor_within_camp <- function(curr_SITE, neighbors_list,
                                                  SITE_data, camp_SITE_indices){
  curr_camp  <- SITE_data$camp_id[curr_SITE]
  neigh      <- neighbors_list[[as.character(curr_SITE)]]
  neigh_same <- neigh[SITE_data$camp_id[neigh] == curr_camp]
  neigh_same <- neigh_same[neigh_same != curr_SITE]
  if(length(neigh_same) > 0) return(sample(neigh_same, 1))
  cand <- camp_SITE_indices[[curr_camp]]
  cand <- cand[cand != curr_SITE]
  if(length(cand) == 0) curr_SITE else sample(cand, 1)
}

pick_destination_other_camp_random <- function(curr_SITE, camps, SITE_data,
                                               camp_SITE_indices){
  curr_camp <- SITE_data$camp_id[curr_SITE]
  dest_camp <- nearest_camp_id(curr_camp, camps)
  cand      <- camp_SITE_indices[[dest_camp]]
  if(length(cand) == 0) curr_SITE else sample(cand, 1)
}

rtruncnorm <- function(n, mean, sd, lo, hi){
  out <- numeric(n)
  for(i in 1:n){
    x <- rnorm(1, mean, sd)
    while(x < lo || x > hi) x <- rnorm(1, mean, sd)
    out[i] <- x
  }
  out
}

sl_within  <- function() pmin(rgamma(1, shape = 2.2, scale = 0.12), 0.7)
sl_between <- function() rtruncnorm(1, mean = 6.0,  sd = 1.1, lo = 3.5,  hi = 9.5)
sl_camp    <- function() rtruncnorm(1, mean = 17.0, sd = 1.7, lo = 11.5, hi = 23.0)

# ----------------------------------------------------------------------------
# SIMULATION LOOP
# ----------------------------------------------------------------------------
cat("Simulating movement...\n")

current_SITE  <- sample.int(total_SITEes, 1)
cx <- SITE_data$SITE_x[current_SITE]
cy <- SITE_data$SITE_y[current_SITE]
current_x     <- cx + runif(1, -SITE_data$SITE_size[current_SITE]/3,
                            SITE_data$SITE_size[current_SITE]/3)
current_y     <- cy + runif(1, -SITE_data$SITE_size[current_SITE]/3,
                            SITE_data$SITE_size[current_SITE]/3)
current_state <- 1L
dest_SITE     <- NA_integer_
dest_camp_trip <- NA_integer_
prev_bearing  <- runif(1, -pi, pi)

x1 <- y1 <- x2 <- y2 <- step_len_out <- bearing_out <- turn_out <- numeric(n_steps)
mu_end_out    <- numeric(n_steps)
true_state_out <- character(n_steps)
SITE_id_out   <- camp_id_out <- integer(n_steps)

for(step_i in 1:n_steps){
  
  current_camp <- SITE_data$camp_id[current_SITE]
  
  g_cur        <- g_of_V(V_SITE[current_SITE])
  gbar_camp    <- mean(g_of_V(V_SITE[camp_SITE_indices[[current_camp]]]))
  gbar_global  <- mean(g_of_V(V_SITE))
  
  local_mu     <- get_mu_time(current_x, current_y, step_i)
  camp_mu_mean <- mean(mu_SITE[step_i, camp_SITE_indices[[current_camp]]], na.rm = TRUE)
  
  camp_resource_bad  <- (gbar_camp < resource_rel_thresh * gbar_global)
  camp_predation_bad <- (camp_mu_mean > camp_mu_thresh)
  local_predation_bad <- (local_mu > local_mu_thresh)
  
  trigger_score <- 1.0 * as.numeric(camp_resource_bad) +
    1.0 * as.numeric(camp_predation_bad) +
    0.8 * as.numeric(local_predation_bad)
  
  p_camp_switch <- if(trigger_score <= 0) 0 else
    pmin(p_camp_switch_max, p_camp_switch_min + 0.22 * trigger_score)
  
  leave_SITE_rule <- (g_cur < (gbar_camp * pmax(local_mu, 1e-6)) / (1 + T_const))
  
  state_t <- current_state
  
  if(current_state == 1L){
    if(runif(1) < p_camp_switch){
      state_t       <- 3L
      dest_SITE     <- pick_destination_other_camp_random(current_SITE, camps, SITE_data, camp_SITE_indices)
      dest_camp_trip <- SITE_data$camp_id[dest_SITE]
    } else if(leave_SITE_rule || runif(1) < p_spont_between_SITE){
      state_t       <- 2L
      dest_SITE     <- pick_destination_neighbor_within_camp(current_SITE, neighbors_list, SITE_data, camp_SITE_indices)
      dest_camp_trip <- NA_integer_
    } else {
      state_t       <- 1L
      dest_SITE     <- NA_integer_
      dest_camp_trip <- NA_integer_
    }
  } else if(current_state == 2L){
    if(runif(1) < (0.25 * p_camp_switch)){
      state_t       <- 3L
      dest_SITE     <- pick_destination_other_camp_random(current_SITE, camps, SITE_data, camp_SITE_indices)
      dest_camp_trip <- SITE_data$camp_id[dest_SITE]
    } else {
      state_t <- 2L
      if(is.na(dest_SITE) || dest_SITE == current_SITE || runif(1) < 0.25)
        dest_SITE <- pick_destination_neighbor_within_camp(current_SITE, neighbors_list, SITE_data, camp_SITE_indices)
      dest_camp_trip <- NA_integer_
    }
  } else {
    state_t <- 3L
    if(is.na(dest_SITE) || SITE_data$camp_id[dest_SITE] == current_camp){
      dest_SITE     <- pick_destination_other_camp_random(current_SITE, camps, SITE_data, camp_SITE_indices)
      dest_camp_trip <- SITE_data$camp_id[dest_SITE]
    }
  }
  
  state_next <- state_t
  
  if(state_t == 1L){
    step_length <- sl_within()
    theta <- wrap_pi(prev_bearing + runif(1, -pi, pi))
    new_x <- current_x + step_length * cos(theta)
    new_y <- current_y + step_length * sin(theta)
    cx2 <- SITE_data$SITE_x[current_SITE]
    cy2 <- SITE_data$SITE_y[current_SITE]
    if(sqrt((new_x - cx2)^2 + (new_y - cy2)^2) > SITE_data$SITE_size[current_SITE]){
      ang_to_center <- atan2(cy2 - current_y, cx2 - current_x)
      theta <- wrap_pi(ang_to_center + runif(1, -pi/6, pi/6))
      new_x <- current_x + step_length * cos(theta)
      new_y <- current_y + step_length * sin(theta)
    }
    
  } else if(state_t == 2L){
    step_length <- sl_between()
    tgt  <- sample_point_in_SITE(dest_SITE, SITE_data)
    cand <- move_toward_point_persistent(
      current_x, current_y, prev_bearing,
      tgt[1], tgt[2], step_length, 0.60, 0.35, landscape_extent)
    theta <- cand$theta; new_x <- cand$new_x; new_y <- cand$new_y
    
    pid_new <- get_SITE_id(new_x, new_y)
    if(pid_new > 0L){
      current_SITE <- pid_new
      if(current_SITE == dest_SITE){
        state_next <- 1L; dest_SITE <- NA_integer_
      } else {
        state_next <- 2L
        dest_SITE  <- pick_destination_neighbor_within_camp(current_SITE, neighbors_list, SITE_data, camp_SITE_indices)
      }
    }
    
  } else {
    step_length <- sl_camp()
    hop_SITE <- safe_next_hop(g, current_SITE, dest_SITE, neighbors_list)
    tgt  <- sample_point_in_SITE(hop_SITE, SITE_data)
    cand <- move_toward_point_persistent(
      current_x, current_y, prev_bearing,
      tgt[1], tgt[2], step_length, 0.88, 0.10, landscape_extent)
    theta <- cand$theta; new_x <- cand$new_x; new_y <- cand$new_y
    
    pid_new <- get_SITE_id(new_x, new_y)
    if(pid_new > 0L){
      current_SITE <- pid_new
      if(!is.na(dest_camp_trip) && SITE_data$camp_id[current_SITE] == dest_camp_trip){
        state_next    <- 1L
        dest_SITE     <- NA_integer_
        dest_camp_trip <- NA_integer_
      }
    }
  }
  
  if(!is.finite(new_x) || !is.finite(new_y)){
    new_x <- current_x; new_y <- current_y; theta <- prev_bearing
  }
  new_x <- pmin(pmax(new_x, landscape_extent[1]), landscape_extent[2])
  new_y <- pmin(pmax(new_y, landscape_extent[3]), landscape_extent[4])
  
  growth <- rv * V_SITE * (1 - V_SITE / Kcap)
  cons   <- numeric(KSITE)
  if(state_t == 1L) cons[current_SITE] <- g_of_V(V_SITE[current_SITE])
  V_SITE <- pmin(pmax(V_SITE + dt * (growth - cons), 0), Kcap)
  V_hist[step_i, ] <- V_SITE
  
  x1[step_i] <- current_x; y1[step_i] <- current_y
  x2[step_i] <- new_x;     y2[step_i] <- new_y
  step_len_out[step_i]  <- sqrt((new_x - current_x)^2 + (new_y - current_y)^2)
  bearing_out[step_i]   <- atan2(new_y - current_y, new_x - current_x)
  turn_out[step_i]      <- if(step_i == 1) 0 else
    wrap_pi(bearing_out[step_i] - bearing_out[step_i-1])
  mu_end_out[step_i]    <- get_mu_time(new_x, new_y, step_i)
  true_state_out[step_i] <- state_names[state_t]
  SITE_id_out[step_i]   <- current_SITE
  camp_id_out[step_i]   <- SITE_data$camp_id[current_SITE]
  
  prev_bearing  <- bearing_out[step_i]
  current_x     <- new_x
  current_y     <- new_y
  current_state <- state_next
}

movement_data <- data.frame(
  step_id          = 1:n_steps,
  x1 = x1, y1 = y1, x2 = x2, y2 = y2,
  step_length      = step_len_out,
  bearing          = bearing_out,
  turning_angle    = turn_out,
  predation_stress = mu_end_out,
  true_state       = true_state_out,
  SITE_id          = SITE_id_out,
  camp_id          = camp_id_out,
  stringsAsFactors = FALSE
) %>%
  mutate(
    ID   = "ind1",
    time = as.POSIXct("2020-01-01 00:00:00", tz = "UTC") + (step_id - 1) * 3600
  )

cat("Movement generated. True-state counts:\n")
print(table(movement_data$true_state))

# ----------------------------------------------------------------------------
# CASE-CONTROL
# ----------------------------------------------------------------------------
cat("Generating controls...\n")
n_controls <- 20

trk <- movement_data %>%
  transmute(
    ID   = as.character(ID),
    time = as.POSIXct(time, tz = "UTC"),
    x    = as.numeric(x2),
    y    = as.numeric(y2)
  ) %>%
  arrange(ID, time) %>%
  distinct(ID, time, .keep_all = TRUE) %>%
  filter(is.finite(x), is.finite(y), !is.na(time), !is.na(ID))

controls <- hmmSSF::get_controls(
  obs        = trk,
  n_controls = n_controls,
  distr      = c("gamma", "vm")
)

t0    <- min(trk$time)
t_idx <- as.integer(round(as.numeric(difftime(controls$time, t0, units = "hours")))) + 1L
t_idx <- pmin(pmax(t_idx, 1L), n_steps)

mu_cc <- numeric(nrow(controls))
for(tt in sort(unique(t_idx))){
  ii        <- which(t_idx == tt)
  mu_cc[ii] <- get_mu_time_vec(as.numeric(controls$x[ii]),
                               as.numeric(controls$y[ii]), tt)
}
controls$predation_stress <- mu_cc

pid_end  <- get_SITE_id_vec(as.numeric(controls$x), as.numeric(controls$y))
camp_end <- rep(NA_integer_, length(pid_end))
inside   <- pid_end > 0L
camp_end[inside] <- SITE_data$camp_id[pid_end[inside]]

g_end        <- numeric(nrow(controls))
gbar_camp_end <- numeric(nrow(controls))

for(tt in sort(unique(t_idx))){
  ii       <- which(t_idx == tt)
  g_all    <- g_of_V(V_hist[tt, ])
  g_global <- mean(g_all, na.rm = TRUE)
  pid_i    <- pid_end[ii]
  camp_i   <- camp_end[ii]
  g_end[ii] <- ifelse(pid_i > 0L, g_all[pmax(1L, pid_i)], g_global)
  gbar_camp_end[ii] <- vapply(seq_along(ii), function(jj){
    cc <- camp_i[jj]
    if(is.na(cc)) return(g_global)
    mean(g_all[camp_SITE_indices[[cc]]], na.rm = TRUE)
  }, numeric(1))
}

controls$delta_g_end <- g_end -
  (gbar_camp_end * pmax(controls$predation_stress, 1e-8)) / (1 + T_const)

camp_mat <- as.matrix(camps[, c("x","y")])
xy_end   <- cbind(as.numeric(controls$x), as.numeric(controls$y))
d2       <- matrix(0, nrow = nrow(xy_end), ncol = nrow(camp_mat))
for(j in 1:nrow(camp_mat)){
  dx <- xy_end[,1] - camp_mat[j,1]; dy <- xy_end[,2] - camp_mat[j,2]
  d2[,j] <- dx*dx + dy*dy
}
controls$dist_nearest_camp <- sqrt(apply(d2, 1, min))

# ----------------------------------------------------------------------------
# FIT HMM-SSF
# ----------------------------------------------------------------------------
cat("Fitting HMM-SSF...\n")

data        <- controls
data$case_  <- data$obs
data$step   <- pmax(as.numeric(data$step), 1e-8)
data$log_step <- log(data$step)
data$angle  <- as.numeric(data$angle)
data$angle[!is.finite(data$angle) | is.na(data$angle)] <- 0
data$delta_g_end <- as.numeric(data$delta_g_end)
data$delta_g_end[!is.finite(data$delta_g_end)] <- NA_real_

need_vars <- c("case_","step","log_step","angle","delta_g_end","stratum")
data <- data[stats::complete.cases(data[, need_vars, drop = FALSE]), , drop = FALSE]

ssf_formula <- case_ ~ step + log_step + cos(angle)
tpm_formula <- ~ delta_g_end

init_SSF_from_single <- function(data, ssf_formula, N, n_states,
                                 variation_scale = c(0.2, 0.5), seed = NULL){
  if(!is.null(seed)) set.seed(seed)
  terms_all <- attr(stats::terms(ssf_formula), "term.labels")
  terms_all <- terms_all[!grepl("^strata\\(", terms_all)]
  fml_issf  <- as.formula(paste("case_ ~",
                                paste(c(terms_all, "strata(stratum)"), collapse = " + ")))
  m1        <- amt::fit_issf(fml_issf, data = data, model = TRUE)
  coef_0    <- stats::coef(m1$model)
  coef_names <- names(coef_0)
  p         <- length(coef_0)
  se_0      <- rep(variation_scale[2], p)
  if("step"       %in% coef_names) se_0[coef_names == "step"]       <- variation_scale[1]
  if("log_step"   %in% coef_names) se_0[coef_names == "log_step"]   <- variation_scale[1]
  if("cos(angle)" %in% coef_names) se_0[coef_names == "cos(angle)"] <- variation_scale[1]
  state_names0 <- paste0("S", seq_len(n_states))
  lapply(seq_len(N), function(ii){
    mat <- sapply(seq_len(n_states), function(s)
      as.numeric(coef_0) + stats::rnorm(p, mean = 0, sd = se_0))
    mat <- matrix(mat, nrow = p, ncol = n_states)
    rownames(mat) <- coef_names; colnames(mat) <- state_names0
    mat
  })
}

.use_optim_opts <- "optim_opts" %in% names(formals(hmmSSF::hmmSSF))

.hmmSSF_call_compat <- function(data, ssf_formula, tpm_formula,
                                n_states, ssf_par0, maxit){
  if(.use_optim_opts){
    hmmSSF::hmmSSF(data = data, ssf_formula = ssf_formula,
                   tpm_formula = tpm_formula, n_states = n_states,
                   ssf_par0 = ssf_par0, optim_opts = list(trace = 0, maxit = maxit))
  } else {
    hmmSSF::hmmSSF(data = data, ssf_formula = ssf_formula,
                   tpm_formula = tpm_formula, n_states = n_states,
                   ssf_par0 = ssf_par0, control = list(trace = 0, maxit = maxit))
  }
}

.get_fit_info <- function(res){
  if(inherits(res, "try-error") || is.null(res))
    return(list(ok = FALSE, value = Inf, convergence = NA_integer_))
  val  <- Inf
  if(!is.null(res$fit$value) && is.finite(res$fit$value)) val <- res$fit$value
  conv <- if(!is.null(res$fit$convergence)) res$fit$convergence else NA_integer_
  list(ok = is.finite(val) && (is.na(conv) || conv == 0),
       value = val, convergence = conv)
}

short_run_long_run_SSF <- function(initial_params_list, data, ssf_formula,
                                   tpm_formula, n_states,
                                   n_iter_short, n_iter_long,
                                   use_parallel = TRUE){
  if(length(initial_params_list) == 0) stop("initial_params_list is empty.")
  if(!("obs" %in% names(data))) data$obs <- data$case_
  infos <- vector("list", length(initial_params_list))
  for(i in seq_along(initial_params_list)){
    fit_i <- try(.hmmSSF_call_compat(data, ssf_formula, tpm_formula,
                                     n_states, initial_params_list[[i]],
                                     n_iter_short), silent = TRUE)
    infos[[i]] <- .get_fit_info(fit_i)
  }
  ok_idx <- which(vapply(infos, function(x) isTRUE(x$ok), logical(1)))
  if(length(ok_idx) == 0){
    finite_idx <- which(vapply(infos, function(x) is.finite(x$value), logical(1)))
    if(length(finite_idx) == 0) stop("All short-run fits failed.")
    ok_idx <- finite_idx
  }
  vals      <- vapply(infos[ok_idx], function(x) x$value, numeric(1))
  best_local <- ok_idx[which.min(vals)]
  message("Short-run: kept ", length(ok_idx), "/", length(initial_params_list),
          " fits. Best index = ", best_local)
  .hmmSSF_call_compat(data, ssf_formula, tpm_formula, n_states,
                      initial_params_list[[best_local]], n_iter_long)
}

N            <- 24
n_iter_short <- 250
n_iter_long  <- 5000

initial_params_list <- init_SSF_from_single(
  data = data, ssf_formula = ssf_formula, N = N,
  n_states = n_states, variation_scale = c(0.2, 0.6), seed = 123)

mod <- short_run_long_run_SSF(
  initial_params_list = initial_params_list,
  data = data, ssf_formula = ssf_formula, tpm_formula = tpm_formula,
  n_states = n_states, n_iter_short = n_iter_short,
  n_iter_long = n_iter_long, use_parallel = FALSE)

cat("Done. Saving model object...\n")
saveRDS(mod, file.path(output_dir, "hmmssf_fit.rds"))

# ----------------------------------------------------------------------------
# Viterbi decoding
# ----------------------------------------------------------------------------
cat("Decoding states...\n")
decoded_states_all  <- hmmSSF::viterbi_decoding(mod = mod)
decoded_case        <- decoded_states_all

trk_obs <- trk %>% arrange(ID, time)
trk_cases_aligned <- trk_obs %>%
  group_by(ID) %>%
  slice((n() - length(decoded_case) + 1):n()) %>%
  ungroup() %>%
  mutate(viterbi_state = decoded_case)

pal_state       <- .get_state_palette()

# ----------------------------------------------------------------------------
# Join TRUE state with Viterbi
# ----------------------------------------------------------------------------
trk_joined <- trk_cases_aligned %>%
  left_join(
    movement_data %>%
      transmute(
        ID            = as.character(ID),
        time          = as.POSIXct(time, tz = "UTC"),
        true_state    = true_state,
        step_length   = step_length,
        turning_angle = turning_angle
      ),
    by = c("ID", "time")
  )

# ----------------------------------------------------------------------------
# State labelling
# ----------------------------------------------------------------------------
state_summ <- trk_joined %>%
  filter(!is.na(viterbi_state), is.finite(step_length)) %>%
  group_by(viterbi_state) %>%
  summarise(mean_step = mean(step_length), n = dplyr::n(), .groups = "drop")

ord       <- state_summ$viterbi_state[order(state_summ$mean_step)]
label_map <- data.frame(
  viterbi_state = ord,
  label         = c("WITHIN_SITE","BETWEEN_SITE","BETWEEN_CAMP"),
  stringsAsFactors = FALSE
)

pal_viterbi_matched <- setNames(
  pal_state[label_map$label],
  paste0("S", label_map$viterbi_state)
)
trk_joined <- trk_joined %>%
  left_join(label_map, by = "viterbi_state") %>%
  mutate(viterbi_label = factor(label,
                                levels = c("WITHIN_SITE","BETWEEN_SITE","BETWEEN_CAMP")))

trk_joined <- trk_joined %>%
  arrange(ID, time) %>%
  group_by(ID) %>%
  mutate(step_id = row_number()) %>%
  ungroup() %>%
  mutate(
    viterbi_label_chr = as.character(viterbi_label),
    match_state = !is.na(viterbi_label_chr) & (true_state == viterbi_label_chr)
  )

# ----------------------------------------------------------------------------
# Trajectory plots
# ----------------------------------------------------------------------------
p_map_overlap <- ggplot(trk_joined) +
  geom_path(aes(x = x, y = y), color = "grey35", linewidth = 0.35, alpha = 0.55) +
  geom_point(aes(x = x, y = y, color = true_state, fill = viterbi_label),
             shape = 21, size = 2.0, stroke = 0.85, alpha = 0.95) +
  scale_color_manual(values = pal_state, drop = FALSE, name = "True state") +
  scale_fill_manual(values = pal_state,  drop = FALSE, name = "Viterbi state") +
  coord_equal() + theme_minimal() + theme(legend.position = "bottom") +
  labs(title = "Overlap vs mismatch on the trajectory", x = "x", y = "y")
.save_png(p_map_overlap, "traj_overlap_true_border_viterbi_fill.png",
          w = 9.2, h = 6.4, dpi = 220)

p_traj_viterbi_labeled <- ggplot(trk_joined) +
  geom_path(aes(x = x, y = y), color = "grey35", linewidth = 0.35, alpha = 0.55) +
  geom_point(aes(x = x, y = y, color = viterbi_label), size = 1.9, alpha = 0.95) +
  scale_color_manual(values = pal_state, drop = FALSE, name = "Label") +
  coord_equal() + theme_minimal() +
  labs(title = "Trajectory with Viterbi states (auto-labelled)", x = "x", y = "y")

p_traj_true <- ggplot(movement_data) +
  geom_path(aes(x = x2, y = y2), color = "grey35", linewidth = 0.35, alpha = 0.55) +
  geom_point(aes(x = x2, y = y2, color = true_state), size = 1.7, alpha = 0.95) +
  scale_color_manual(values = pal_state, drop = FALSE, name = "True state") +
  coord_equal() + theme_minimal() +
  labs(title = "Simulated trajectory (true states)", x = "x", y = "y")
.save_png(p_traj_true, "02_trajectory_true_states.png", 7.5, 5)

.get_legend <- function(p){
  g   <- ggplot2::ggplotGrob(p)
  idx <- which(vapply(g$grobs, function(x) x$name, character(1)) == "guide-box")
  if(length(idx) == 0) return(NULL)
  g$grobs[[idx[1]]]
}

leg <- .get_legend(p_traj_viterbi_labeled)
p1  <- p_traj_true            + theme(legend.position = "none")
p2  <- p_traj_viterbi_labeled + theme(legend.position = "none")
top <- gridExtra::arrangeGrob(p1, p2, ncol = 2)
p_traj_and_time <- gridExtra::arrangeGrob(top, leg, ncol = 2, widths = c(1, 0.3))
.save_png(p_traj_and_time, "traj_true_viterbi.png", 9, 5.0, dpi = 220)

# ----------------------------------------------------------------------------
# Density plots — step length and turning angle
# ----------------------------------------------------------------------------

# matched palette: each Viterbi S-number gets the colour of its decoded label
pal_viterbi_matched <- setNames(
  pal_state[label_map$label],
  paste0("S", label_map$viterbi_state)
)

# ---- step-length densities
df_step_true <- trk_joined %>%
  transmute(state = true_state, value = step_length) %>%
  filter(is.finite(value), !is.na(state))

df_step_vit <- trk_joined %>%
  transmute(state = paste0("S", viterbi_state), value = step_length) %>%
  filter(is.finite(value), !is.na(state))

p_step_dens_True <- ggplot(df_step_true, aes(x = value, color = state, fill = state)) +
  geom_density(alpha = 0.25) +
  scale_color_manual(values = pal_state, drop = FALSE) +
  scale_fill_manual( values = pal_state, drop = FALSE) +
  theme_minimal() +
  labs(title = "Step-length (True state)", x = "step length", y = "density") +
  theme(legend.position = "bottom")

p_step_dens_Viterbi <- ggplot(df_step_vit, aes(x = value, color = state, fill = state)) +
  geom_density(alpha = 0.25) +
  scale_color_manual(values = pal_viterbi_matched, drop = FALSE) +
  scale_fill_manual( values = pal_viterbi_matched, drop = FALSE) +
  theme_minimal() +
  labs(title = "Step-length (Viterbi decoding)", x = "step length", y = "density") +
  theme(legend.position = "bottom")

# ---- turning-angle densities
df_ang_true <- trk_joined %>%
  transmute(state = true_state, value = turning_angle) %>%
  filter(is.finite(value), !is.na(state))

df_ang_vit <- trk_joined %>%
  transmute(state = paste0("S", viterbi_state), value = turning_angle) %>%
  filter(is.finite(value), !is.na(state))

p_ang_dens_True <- ggplot(df_ang_true, aes(x = value, color = state, fill = state)) +
  geom_density(alpha = 0.25) +
  scale_color_manual(values = pal_state, drop = FALSE) +
  scale_fill_manual( values = pal_state, drop = FALSE) +
  theme_minimal() +
  labs(title = "Turning-angle (True state)", x = "turning angle (rad)", y = "density") +
  theme(legend.position = "bottom")

p_ang_dens_Viterbi <- ggplot(df_ang_vit, aes(x = value, color = state, fill = state)) +
  geom_density(alpha = 0.25) +
  scale_color_manual(values = pal_viterbi_matched, drop = FALSE) +
  scale_fill_manual( values = pal_viterbi_matched, drop = FALSE) +
  theme_minimal() +
  labs(title = "Turning-angle (Viterbi decoding)", x = "turning angle (rad)", y = "density") +
  theme(legend.position = "bottom")

# ---- combine
p_dens_both <- (
  (p_step_dens_True | p_step_dens_Viterbi) /
    (p_ang_dens_True  | p_ang_dens_Viterbi)
)
.save_png(p_dens_both, "densities_step_and_angle_true_vs_viterbi.png",
          9.2, 6.6, dpi = 220)

# ----------------------------------------------------------------------------
# SSF parameter plot
# ----------------------------------------------------------------------------
ssf_par <- as.matrix(mod$par$ssf)
rownames(ssf_par) <- rownames(mod$par$ssf)

ci <- try(confint(mod, range = 0.95, pretty = TRUE), silent = TRUE)

if(!inherits(ci, "try-error") && !is.null(ci$ssf)){
  ssf_df <- data.frame(
    state     = rep(colnames(ssf_par), each = nrow(ssf_par)),
    covariate = rep(rownames(ssf_par), times = ncol(ssf_par)),
    estimate  = as.numeric(ci$ssf[,1]),
    lower     = as.numeric(ci$ssf[,2]),
    upper     = as.numeric(ci$ssf[,3])
  )
} else {
  ssf_df <- data.frame(
    state     = rep(colnames(ssf_par), each = nrow(ssf_par)),
    covariate = rep(rownames(ssf_par), times = ncol(ssf_par)),
    estimate  = as.numeric(ssf_par),
    lower = NA_real_, upper = NA_real_
  )
}

pd <- position_dodge(0.25)
p_ssf <- ggplot(ssf_df, aes(x = covariate, y = estimate, color = state, group = state)) +
  scale_color_manual(values = pal_viterbi_matched, drop = FALSE) +
  geom_point(position = pd, size = 2) +
  geom_errorbar(aes(ymin = lower, ymax = upper), position = pd,
                width = 0.15, linewidth = 0.6, na.rm = TRUE) +
  theme_minimal(base_size = 15) +
  labs(title = "SSF parameter estimates", x = "", y = expression(beta), color = "State") +
  theme(
    plot.title   = element_text(size = 18, face = "bold"),
    axis.title   = element_text(size = 16),
    axis.text    = element_text(size = 14),
    axis.text.x  = element_text(angle = 35, hjust = 1, size = 13),
    legend.title = element_text(size = 14),
    legend.text  = element_text(size = 13),
    legend.position = "bottom"
  )

# ----------------------------------------------------------------------------
# TPM
# ----------------------------------------------------------------------------
delta_grid <- seq(
  quantile(data$delta_g_end, 0.05, na.rm = TRUE),
  quantile(data$delta_g_end, 0.95, na.rm = TRUE),
  length.out = 80
)
new_data <- data.frame(delta_g_end = delta_grid)

tpm_pred <- hmmSSF::predict_tpm(mod = mod, new_data = new_data, return_CI = TRUE)

df_mle <- as.data.frame.table(tpm_pred$mle)
colnames(df_mle) <- c("from","to","idx","value")
df_mle$delta_g_end <- rep(delta_grid, each = n_states * n_states)
df_mle$low <- as.data.frame.table(tpm_pred$lower)[,4]
df_mle$upp <- as.data.frame.table(tpm_pred$upper)[,4]


# Extract the raw state numbers from label_map
s_within  <- label_map$viterbi_state[label_map$label == "WITHIN_SITE"]
s_between <- label_map$viterbi_state[label_map$label == "BETWEEN_SITE"]
s_camp    <- label_map$viterbi_state[label_map$label == "BETWEEN_CAMP"]

# Build the transition labels dynamically
df_tpm_both <- df_mle %>%
  dplyr::filter(
    (from == paste0("S", s_within)  & to == paste0("S", s_between)) |
      (from == paste0("S", s_between) & to == paste0("S", s_camp))
  ) %>%
  dplyr::mutate(
    transition = dplyr::case_when(
      from == paste0("S", s_within)  & to == paste0("S", s_between) ~ 
        "WITHIN_SITE \u2192 BETWEEN_SITE",
      from == paste0("S", s_between) & to == paste0("S", s_camp) ~ 
        "BETWEEN_SITE \u2192 BETWEEN_CAMP",
      TRUE ~ paste0(from, "\u2192", to)
    )
  )

df_tpm_both$transition <- factor(
  df_tpm_both$transition,
  levels = c("WITHIN_SITE \u2192 BETWEEN_SITE",
             "BETWEEN_SITE \u2192 BETWEEN_CAMP")
)

pal_tpm <- c(
  "WITHIN_SITE \u2192 BETWEEN_SITE"  = "#1f77b4",
  "BETWEEN_SITE \u2192 BETWEEN_CAMP" = "#d62728"
)

p_tpm_both <- ggplot(df_tpm_both,
                     aes(x = delta_g_end, y = value,
                         color = transition, fill = transition)) +
  geom_ribbon(aes(ymin = low, ymax = upp), alpha = 0.22, linewidth = 0) +
  geom_line(linewidth = 1.1) +
  scale_color_manual(values = pal_tpm, drop = FALSE, name = "Transition") +
  scale_fill_manual( values = pal_tpm, drop = FALSE, name = "Transition") +
  theme_minimal(base_size = 15) +
  theme(
    plot.title   = element_text(size = 18, face = "bold"),
    axis.title   = element_text(size = 16),
    axis.text    = element_text(size = 15),
    legend.title = element_text(size = 14),
    legend.text  = element_text(size = 13),
    legend.position = "bottom"
  ) +
  labs(x = expression(Delta*g(t)), y = "Transition probability",
       title = "Transition probabilities")

p_ssf2 <- p_ssf +
  guides(color = guide_legend(nrow = 2, byrow = TRUE)) +
  theme(legend.position = "bottom",
        legend.box.margin = margin(t = 2, r = 2, b = 6, l = 2))

p_tpm2 <- p_tpm_both +
  guides(color = guide_legend(nrow = 2, byrow = TRUE),
         fill  = guide_legend(nrow = 2, byrow = TRUE)) +
  theme(legend.position = "bottom",
        legend.box.margin = margin(t = 2, r = 2, b = 6, l = 2))

p_two_cols <- (p_ssf2 | p_tpm2)
.save_png(p_two_cols, "ssf_tpm.png", w = 10.2, h = 6.2, dpi = 220)

# ----------------------------------------------------------------------------
# SITE RESIDENCE TIME
# ----------------------------------------------------------------------------
cat("Computing SITE residence times from Viterbi-decoded output...\n")

dt_hours <- 1

trk_joined$SITE_id_end <- get_SITE_id_vec(
  as.numeric(trk_joined$x), as.numeric(trk_joined$y))

if(!("SITE_id" %in% names(trk_joined))){
  trk_joined <- trk_joined %>%
    left_join(movement_data %>% select(ID, time, SITE_id, camp_id),
              by = c("ID","time"))
}

trk_joined <- trk_joined %>%
  arrange(ID, step_id) %>%
  mutate(is_resident = (viterbi_label == "WITHIN_SITE") & !is.na(SITE_id))

trk_joined <- trk_joined %>%
  mutate(
    start_bout = is_resident & (
      !lag(is_resident, default = FALSE) |
        lag(SITE_id, default = NA_integer_) != SITE_id
    ),
    bout_id = cumsum(start_bout)
  )

res_bouts <- trk_joined %>%
  filter(is_resident) %>%
  group_by(ID, bout_id, SITE_id, camp_id) %>%
  summarise(
    n_steps        = n(),
    duration_hours = n_steps * dt_hours,
    .groups = "drop"
  )

p_res_dist <- ggplot(res_bouts, aes(x = duration_hours)) +
  geom_density(linewidth = 1.1, fill = "grey60", alpha = 0.5, adjust = 1.1) +
  theme_minimal(base_size = 15) +
  labs(title = "Density of SITE residence times",
       x = "Residence time", y = "Density") +
  theme(plot.title = element_text(face = "bold", size = 17),
        axis.title = element_text(size = 15),
        axis.text  = element_text(size = 13))

p_res_SITE <- res_bouts %>%
  group_by(SITE_id) %>%
  summarise(mean_hours = mean(duration_hours), .groups = "drop") %>%
  ggplot(aes(x = factor(SITE_id), y = mean_hours)) +
  geom_col(alpha = 0.85, fill = "steelblue") +
  theme_minimal(base_size = 15) +
  labs(title = "Mean residence time by SITE",
       x = "SITE ID", y = "Mean residence time") +
  theme(plot.title  = element_text(face = "bold", size = 17),
        axis.title  = element_text(size = 15),
        axis.text.x = element_text(hjust = 1, size = 12),
        axis.text.y = element_text(size = 13))

p_res_combined <- p_res_dist | p_res_SITE
.save_png(p_res_combined, "SITE_residence_time_combined.png",
          w = 13, h = 5.5, dpi = 220)

# Overall accuracy
cat("Overall accuracy:", round(mean(trk_joined$match_state, na.rm = TRUE) * 100, 1), "%\n")

# Per-state confusion matrix (rows = true, cols = Viterbi)
conf <- table(True = trk_joined$true_state,
              Viterbi = trk_joined$viterbi_label_chr)
conf_prop <- prop.table(conf, margin = 1)  # row-normalised
print(round(conf_prop * 100, 1))