# =============================================================================
# Estimation for discrete-time Markov chains: GeoLife trajectories
# =============================================================================
#
# GeoLife trajectories as a random walk on Z^2. An infinite state space, so
# we impose parametric structure: the walk is described by five step
# probabilities rather than a transition matrix.
#
# See also precip-models.R and nyctaxi-models.R for the other two real data
# examples (sibling scripts, split out of what was originally one combined
# dtmc-estimation.R).
#
# FIRST DRAFT: originally an archived script, repaired enough to run top to
# bottom. Comments flag [CUT?] expensive or peripheral parts, [FIX] repairs
# to the originals.

library(tidyverse)
library(markovchain)
library(sf)          # maps
library(ggspatial)   # basemap tiles

## ---------------------------------------------------------------------------
## 1. DATA
## ---------------------------------------------------------------------------

# helper functions (defines sample_geolife_same_locale_fast, etc.)
# [FIX] the helpers are still in the archive; move them to scripts/ if we keep this part
source("_archive/scripts/sample_geolife.R")

geolife_root <- "_data/Geolife Trajectories 1.3/Data"

# knobs
n_traj   <- 100   # how many trajectories to pull
n_plot   <- 10    # how many to show on the map
cell_deg <- 0.02  # grid resolution (~2.2km in latitude)

# sample trajectories from a single locale  [CUT?] slow: reads many raw files
pts_local <- sample_geolife_same_locale_fast(
  n = n_traj,
  radius_km = 60,
  n_candidates = 1500,
  head_pts = 50,
  seed = 42,
  root = geolife_root
)

# plot a few example trajectories on an OSM basemap  [CUT?]
set.seed(12)
traj_keep <- pts_local |>
  distinct(traj_id) |>
  slice_sample(n = min(n_plot, n_distinct(pts_local$traj_id))) |>
  pull(traj_id)

traj_lines <- pts_local |>
  filter(traj_id %in% traj_keep) |>
  st_as_sf(coords = c("lon", "lat"), crs = 4326, remove = FALSE) |>
  arrange(traj_id, t) |>
  group_by(traj_id) |>
  summarise(do_union = FALSE, .groups = "drop") |>
  st_cast("LINESTRING") |>
  st_transform(3857)

ggplot() +
  annotation_map_tile(type = "osm", zoomin = 0) +
  geom_rect(
    inherit.aes = FALSE,
    aes(xmin = -Inf, xmax = Inf, ymin = -Inf, ymax = Inf),
    fill = "white", alpha = 0.35
  ) +
  geom_sf(data = traj_lines, aes(color = factor(traj_id)),
          linewidth = 1.0, alpha = 0.9) +
  coord_sf(crs = st_crs(3857), expand = FALSE) +
  theme_minimal() +
  labs(title = paste0("GeoLife trajectories (", length(traj_keep), " sampled)")) +
  guides(color = "none")

## ---------------------------------------------------------------------------
## 2. DISCRETIZE TO A GRID, THEN COMPLETE THE PATHS
## ---------------------------------------------------------------------------

pts_states <- pts_local |>
  arrange(traj_id, t) |>
  mutate(
    gx = floor(lon / cell_deg),
    gy = floor(lat / cell_deg),
    state = paste(gx, gy, sep = "_")
  )

# drop consecutive duplicates in cell space (repeated GPS points in one cell)
pts_states <- pts_states |>
  group_by(traj_id) |>
  filter(gx != lag(gx) | gy != lag(gy) | is.na(lag(gx))) |>
  ungroup()

# path completion: insert intermediate cells along a Manhattan path so that
# every step is a 4-neighbor move, |dx| + |dy| in {0, 1}
complete_path_xy <- function(gx, gy) {
  stopifnot(length(gx) == length(gy))
  n <- length(gx)
  if (n == 0) return(list(gx = integer(), gy = integer()))
  if (n == 1) return(list(gx = gx, gy = gy))

  out_x <- integer(0)
  out_y <- integer(0)

  # start at the first observed cell
  x <- gx[1]
  y <- gy[1]
  out_x <- c(out_x, x)
  out_y <- c(out_y, y)

  for (i in 2:n) {
    x1 <- gx[i]
    y1 <- gy[i]

    dx <- x1 - x
    dy <- y1 - y

    # move in x direction one unit at a time
    if (dx != 0) {
      sx <- sign(dx)
      for (k in seq_len(abs(dx))) {
        x <- x + sx
        out_x <- c(out_x, x)
        out_y <- c(out_y, y)
      }
    }

    # then move in y direction one unit at a time
    if (dy != 0) {
      sy <- sign(dy)
      for (k in seq_len(abs(dy))) {
        y <- y + sy
        out_x <- c(out_x, x)
        out_y <- c(out_y, y)
      }
    }
  }

  list(gx = out_x, gy = out_y)
}

# completed state sequences, one per trajectory (no stitching across them)
seq_tbl <- pts_states |>
  group_by(user, traj_id, traj_file) |>
  summarise(gx = list(gx), gy = list(gy), .groups = "drop") |>
  mutate(
    completed = purrr::map2(gx, gy, complete_path_xy),
    gx_c = purrr::map(completed, "gx"),
    gy_c = purrr::map(completed, "gy"),
    state_seq = purrr::map2(gx_c, gy_c, ~ paste(.x, .y, sep = "_"))
  ) |>
  select(user, traj_id, traj_file, state_seq, gx_c, gy_c)

## ---------------------------------------------------------------------------
## 3. MLE FOR THE STEP DISTRIBUTION
## ---------------------------------------------------------------------------

# one row per unit step
steps <- seq_tbl |>
  transmute(traj_id, gx = gx_c, gy = gy_c) |>
  tidyr::unnest(c(gx, gy)) |>
  group_by(traj_id) |>
  mutate(
    dx = gx - lag(gx),
    dy = gy - lag(gy)
  ) |>
  filter(!is.na(dx), !is.na(dy)) |>
  ungroup()

step_counts <- steps |>
  mutate(step = case_when(
    dx ==  0 & dy ==  0 ~ "stay",
    dx ==  1 & dy ==  0 ~ "E",
    dx == -1 & dy ==  0 ~ "W",
    dx ==  0 & dy ==  1 ~ "N",
    dx ==  0 & dy == -1 ~ "S",
    TRUE                ~ "other"
  )) |>
  count(step, name = "n") |>
  arrange(desc(n))

step_counts

# under path completion, "other" should be zero (or extremely close)
step_counts |>
  summarise(frac_other = sum(n[step == "other"], na.rm = TRUE) / sum(n))

# MLE step probabilities for the lazy random walk: phat_j = N_j / N
rw_hat <- step_counts |>
  filter(step %in% c("stay", "N", "S", "E", "W")) |>
  mutate(p = n / sum(n)) |>
  arrange(match(step, c("stay", "N", "S", "E", "W")))

rw_hat

# estimated drift (mean increment)
# [FIX] this used `steps` before it existed in the original script
drift_hat <- steps |>
  summarise(mu_dx = mean(dx), mu_dy = mean(dy))

drift_hat

## ---------------------------------------------------------------------------
## 4. VISIT HEATMAP
## ---------------------------------------------------------------------------
## [CUT?] a picture of where the walkers spend their time, not an estimate

visit_counts <- seq_tbl |>
  tidyr::unnest(state_seq) |>
  count(state_seq, name = "visits") |>
  rename(state = state_seq)

heat_sf <- visit_counts |>
  tidyr::separate(state, into = c("gx", "gy"), sep = "_", convert = TRUE) |>
  mutate(
    xmin = gx * cell_deg,
    xmax = (gx + 1) * cell_deg,
    ymin = gy * cell_deg,
    ymax = (gy + 1) * cell_deg,
    geometry = purrr::pmap(
      list(xmin, xmax, ymin, ymax),
      function(xmin, xmax, ymin, ymax) {
        st_polygon(list(matrix(
          c(xmin, ymin,
            xmax, ymin,
            xmax, ymax,
            xmin, ymax,
            xmin, ymin),
          ncol = 2, byrow = TRUE
        )))
      }
    )
  ) |>
  st_as_sf(crs = 4326) |>
  st_transform(3857)

ggplot() +
  annotation_map_tile(type = "osm", zoomin = 0) +
  geom_rect(
    inherit.aes = FALSE,
    aes(xmin = -Inf, xmax = Inf, ymin = -Inf, ymax = Inf),
    fill = "white", alpha = 0.60
  ) +
  geom_sf(data = heat_sf, aes(fill = log1p(visits)), color = NA, alpha = 0.85) +
  coord_sf(crs = st_crs(3857), expand = FALSE) +
  theme_minimal() +
  labs(title = paste0("GeoLife state visit heatmap (cell_deg = ", cell_deg, ")"),
       fill = "log(1 + visits)")

## ---------------------------------------------------------------------------
## 5. SIMULATE FROM THE FITTED WALK
## ---------------------------------------------------------------------------

# [FIX] the original built a 4-state markovchain with identical rows to draw
# i.i.d. steps, which also dropped "stay"; sampling rw_hat directly is the same
# thing and keeps the lazy step
set.seed(1)
sim_steps <- sample(rw_hat$step, size = 200, replace = TRUE, prob = rw_hat$p)

# map step labels -> (dx, dy)
step_to_dxdy <- function(step) {
  switch(step,
         "stay" = c(0L,  0L),
         "N"    = c(0L,  1L),
         "S"    = c(0L, -1L),
         "E"    = c(1L,  0L),
         "W"    = c(-1L, 0L),
         stop("Unknown step: ", step))
}

reconstruct_path <- function(steps, gx0 = 0L, gy0 = 0L) {
  inc <- t(vapply(steps, step_to_dxdy, FUN.VALUE = c(0L, 0L)))
  tibble(
    n  = 0:length(steps),
    gx = c(gx0, gx0 + cumsum(inc[, 1])),
    gy = c(gy0, gy0 + cumsum(inc[, 2]))
  )
}

path_sim <- reconstruct_path(sim_steps)

ggplot(path_sim, aes(gx, gy)) +
  geom_path(linewidth = 0.8, alpha = 0.9) +
  geom_point(size = 0.8, alpha = 0.5) +
  coord_equal() +
  theme_minimal() +
  labs(title = "Simulated trajectory", x = "gx", y = "gy")

## ---------------------------------------------------------------------------
## 6. ARE THE INCREMENTS REALLY I.I.D.?
## ---------------------------------------------------------------------------
##
## The random walk model says each step is drawn fresh. Fit a chain to the step
## sequence itself and see whether the rows differ.

step_states <- c("N", "S", "E", "W")

steps_labeled <- steps |>
  mutate(step = case_when(
    dx ==  1 & dy ==  0 ~ "E",
    dx == -1 & dy ==  0 ~ "W",
    dx ==  0 & dy ==  1 ~ "N",
    dx ==  0 & dy == -1 ~ "S",
    TRUE ~ NA_character_
  )) |>
  filter(!is.na(step)) |>
  arrange(traj_id)  # keep trajectory grouping stable

# step-to-step transition counts (no stitching across trajectories)
N_step <- with(steps_labeled |>
                 group_by(traj_id) |>
                 mutate(step_next = lead(step)) |>
                 filter(!is.na(step_next)) |>
                 ungroup(),
               table(factor(step, levels = step_states),
                     factor(step_next, levels = step_states)))

# transition matrix
P_step_dep <- prop.table(N_step, margin = 1) |> unclass()

mc_steps_dep <- new("markovchain",
                    states = step_states,
                    byrow = TRUE,
                    transitionMatrix = P_step_dep,
                    name = "4-neighbor step process (step-to-step Markov)")

mc_steps_dep

# under the random walk model every row would be the same
# long-run step-direction frequencies
steadyStates(mc_steps_dep)