## =========================================================
## Consolidated stochastic Leslie-matrix workflow
## Folder: E:/LESLIE
## Input file: E:/LESLIE/LESLIE.csv
## =========================================================

## ---------- packages ----------
# install.packages(c("readr", "dplyr", "HDInterval", "ggplot2"))
library(readr)
library(dplyr)
library(HDInterval)
library(ggplot2)

## ---------- settings ----------
csv_file <- "E:/LESLIE/LESLIE.csv"
out_dir  <- "E:/LESLIE"

K <- 1000
p_bad_year <- 0.5
n_reps <- 1000
t_max_short <- 50
t_max_long  <- 2000

check_times   <- c(5, 10, 15, 20, 25, 30, 35, 40, 45, 50)
persist_times <- c(100, 1000, 1500, 2000)

set.seed(123)

## ---------- initial population vectors ----------
## IMPORTANT:
## Replace these if you want a different initial age structure.
## Length must match the number of age classes parsed from each matrix.
##
## Default assumption here: 6 age classes.
## If your matrix has different dimension, the script will stop and tell you.

start_parth <- c(50, 40, 30, 20, 10, 5)
start_sex   <- c(50, 40, 30, 20, 10, 5)

## ---------- helper functions ----------

adult_index_fun <- function(n_age_classes) {
  2:n_age_classes
}

read_leslie_from_csv <- function(file) {
  raw <- read_csv(file, col_names = TRUE, show_col_types = FALSE)
  raw <- as.data.frame(raw)
  names(raw)[1] <- "label"

  i_parth <- which(grepl("parthenogen", raw$label, ignore.case = TRUE))[1]
  i_sex   <- which(grepl("sexual breeder", raw$label, ignore.case = TRUE))[1]

  if (is.na(i_parth) || is.na(i_sex)) {
    stop("Could not find rows containing 'parthenogen' and/or 'sexual breeder'.")
  }

  block_parth <- raw[i_parth:(i_sex - 1), , drop = FALSE]
  block_sex   <- raw[i_sex:nrow(raw), , drop = FALSE]

  clean_block <- function(block) {
    num_part <- suppressWarnings(data.frame(lapply(block[, -1, drop = FALSE], as.numeric)))
    keep <- rowSums(!is.na(num_part)) > 0
    block <- block[keep, , drop = FALSE]
    num_part <- num_part[keep, , drop = FALSE]

    M <- as.matrix(num_part)
    storage.mode(M) <- "numeric"

    M <- M[, colSums(!is.na(M)) > 0, drop = FALSE]
    M[is.na(M)] <- 0

    nr <- nrow(M)
    nc <- ncol(M)

    if (nr < nc) stop("Parsed block has fewer rows than columns.")
    if (nr > nc) M <- M[1:nc, , drop = FALSE]

    M
  }

  A_parth <- clean_block(block_parth)
  A_sex   <- clean_block(block_sex)

  list(parthenogen = A_parth, sexual = A_sex)
}

step_stochastic <- function(N, A, K = 1000, p_bad_year = 0.5,
                            density_model = c("linear", "beverton_holt")) {
  density_model <- match.arg(density_model)

  n <- length(N)
  fert <- A[1, ]
  surv <- diag(A[-1, -ncol(A), drop = FALSE])
  surv <- pmin(pmax(surv, 0), 1)

  N_total <- sum(N)

  good_year <- rbinom(1, size = 1, prob = 1 - p_bad_year)

  if (density_model == "linear") {
    dd <- max(0, 1 - N_total / K)
  } else {
    dd <- 1 / (1 + N_total / K)
  }

  effective_fert <- fert * good_year * dd
  lambda_births <- sum(effective_fert * N)
  newborns <- rpois(1, lambda = max(lambda_births, 0))

  N_next <- rep(0, n)
  N_next[1] <- newborns

  for (i in 1:(n - 1)) {
    N_next[i + 1] <- rbinom(1, size = N[i], prob = surv[i])
  }

  N_next
}

simulate_population <- function(A, N0, years = 2000, K = 1000, p_bad_year = 0.5,
                                density_model = "linear", adult_classes = 2:nrow(A)) {
  n <- length(N0)
  out <- matrix(0, nrow = years + 1, ncol = n)
  out[1, ] <- N0

  extinct_time <- NA_integer_

  for (t in 1:years) {
    out[t + 1, ] <- step_stochastic(
      N = out[t, ],
      A = A,
      K = K,
      p_bad_year = p_bad_year,
      density_model = density_model
    )

    if (sum(out[t + 1, ]) == 0) {
      extinct_time <- t
      if (t < years) {
        out[(t + 2):(years + 1), ] <- 0
      }
      break
    }
  }

  adults <- rowSums(out[, adult_classes, drop = FALSE])

  list(
    full = out,
    adults = adults,
    extinct_time = extinct_time
  )
}

run_replicates <- function(A, N0, n_reps = 1000, years = 2000,
                           K = 1000, p_bad_year = 0.5,
                           density_model = "linear",
                           adult_classes = 2:nrow(A)) {

  adults_mat <- matrix(0, nrow = n_reps, ncol = years + 1)
  ext_times <- rep(NA_integer_, n_reps)

  for (r in 1:n_reps) {
    sim <- simulate_population(
      A = A,
      N0 = N0,
      years = years,
      K = K,
      p_bad_year = p_bad_year,
      density_model = density_model,
      adult_classes = adult_classes
    )
    adults_mat[r, ] <- sim$adults
    ext_times[r] <- sim$extinct_time
  }

  list(adults = adults_mat, extinction_times = ext_times)
}

summarize_time_series <- function(adults_mat, times) {
  res <- lapply(times, function(tt) {
    x <- adults_mat[, tt + 1]
    h <- HDInterval::hdi(x, credMass = 0.95)
    data.frame(
      year = tt,
      mean = mean(x),
      median = median(x),
      hpd_low = h[1],
      hpd_high = h[2],
      q025 = quantile(x, 0.025),
      q975 = quantile(x, 0.975)
    )
  })
  bind_rows(res)
}

summarize_extinction <- function(ext_times) {
  finite_ext <- ext_times[!is.na(ext_times)]

  data.frame(
    n_extinct = length(finite_ext),
    prop_extinct = mean(!is.na(ext_times)),
    mean_ext_time_conditional = if (length(finite_ext) > 0) mean(finite_ext) else NA_real_,
    median_ext_time_conditional = if (length(finite_ext) > 0) median(finite_ext) else NA_real_,
    hpd_ext_time_low = if (length(finite_ext) > 1) HDInterval::hdi(finite_ext, credMass = 0.95)[1] else NA_real_,
    hpd_ext_time_high = if (length(finite_ext) > 1) HDInterval::hdi(finite_ext, credMass = 0.95)[2] else NA_real_
  )
}

summarize_persistence <- function(ext_times, times) {
  res <- lapply(times, function(tt) {
    survived <- ifelse(is.na(ext_times), 1, as.integer(ext_times > tt))
    p <- mean(survived)

    a <- sum(survived) + 0.5
    b <- length(survived) - sum(survived) + 0.5
    ci <- qbeta(c(0.025, 0.975), a, b)

    data.frame(
      year = tt,
      survived_reps = sum(survived),
      total_reps = length(survived),
      proportion_survived = p,
      ci_low = ci[1],
      ci_high = ci[2]
    )
  })
  bind_rows(res)
}

summarize_persistence_curve <- function(ext_times, max_time = 2000) {
  res <- lapply(1:max_time, function(tt) {
    survived <- ifelse(is.na(ext_times), 1, as.integer(ext_times > tt))
    p <- mean(survived)

    a <- sum(survived) + 0.5
    b <- length(survived) - sum(survived) + 0.5
    ci <- qbeta(c(0.025, 0.975), a, b)

    data.frame(
      year = tt,
      survived_reps = sum(survived),
      total_reps = length(survived),
      proportion_survived = p,
      ci_low = ci[1],
      ci_high = ci[2]
    )
  })
  bind_rows(res)
}

make_band <- function(adults_mat, years, species_name) {
  bind_rows(lapply(years, function(tt) {
    x <- adults_mat[, tt + 1]
    h <- HDInterval::hdi(x, credMass = 0.95)
    data.frame(
      year = tt,
      mean = mean(x),
      low = h[1],
      high = h[2],
      species = species_name
    )
  }))
}

## ---------- read matrices ----------
mats <- read_leslie_from_csv(csv_file)
A_parth <- mats$parthenogen
A_sex   <- mats$sexual

cat("Parthenogen matrix:\n")
print(A_parth)
cat("\nSexual matrix:\n")
print(A_sex)

stopifnot(nrow(A_parth) == ncol(A_parth))
stopifnot(nrow(A_sex)   == ncol(A_sex))

if (length(start_parth) != nrow(A_parth)) {
  stop(paste0("start_parth has length ", length(start_parth),
              ", but the parthenogen matrix has ", nrow(A_parth), " age classes."))
}
if (length(start_sex) != nrow(A_sex)) {
  stop(paste0("start_sex has length ", length(start_sex),
              ", but the sexual matrix has ", nrow(A_sex), " age classes."))
}

adult_idx_parth <- adult_index_fun(nrow(A_parth))
adult_idx_sex   <- adult_index_fun(nrow(A_sex))

## ---------- run simulations ----------
res_parth <- run_replicates(
  A = A_parth,
  N0 = start_parth,
  n_reps = n_reps,
  years = t_max_long,
  K = K,
  p_bad_year = p_bad_year,
  density_model = "linear",
  adult_classes = adult_idx_parth
)

res_sex <- run_replicates(
  A = A_sex,
  N0 = start_sex,
  n_reps = n_reps,
  years = t_max_long,
  K = K,
  p_bad_year = p_bad_year,
  density_model = "linear",
  adult_classes = adult_idx_sex
)

## ---------- 1) first 50 years ----------
years_0_50 <- 0:t_max_short

traj_parth <- data.frame(
  year = years_0_50,
  mean_adults = colMeans(res_parth$adults[, years_0_50 + 1, drop = FALSE]),
  median_adults = apply(res_parth$adults[, years_0_50 + 1, drop = FALSE], 2, median),
  species = "parthenogen"
)

traj_sex <- data.frame(
  year = years_0_50,
  mean_adults = colMeans(res_sex$adults[, years_0_50 + 1, drop = FALSE]),
  median_adults = apply(res_sex$adults[, years_0_50 + 1, drop = FALSE], 2, median),
  species = "sexual"
)

traj_all <- bind_rows(traj_parth, traj_sex)

## ---------- 2) extinction summaries ----------
ext_summary_parth <- summarize_extinction(res_parth$extinction_times) %>%
  mutate(species = "parthenogen")

ext_summary_sex <- summarize_extinction(res_sex$extinction_times) %>%
  mutate(species = "sexual")

ext_summary_all <- bind_rows(ext_summary_parth, ext_summary_sex)

## ---------- 3) selected-year HDIs ----------
hpd_parth <- summarize_time_series(res_parth$adults, check_times) %>%
  mutate(species = "parthenogen")

hpd_sex <- summarize_time_series(res_sex$adults, check_times) %>%
  mutate(species = "sexual")

hpd_all <- bind_rows(hpd_parth, hpd_sex)

## ---------- yearly 0-50 HDI band for plot ----------
band_parth <- make_band(res_parth$adults, years_0_50, "parthenogen")
band_sex   <- make_band(res_sex$adults, years_0_50, "sexual")
band_all   <- bind_rows(band_parth, band_sex)

## ---------- 4) persistence at requested years ----------
persist_parth <- summarize_persistence(res_parth$extinction_times, persist_times) %>%
  mutate(species = "parthenogen")

persist_sex <- summarize_persistence(res_sex$extinction_times, persist_times) %>%
  mutate(species = "sexual")

persist_all <- bind_rows(persist_parth, persist_sex)

## ---------- full persistence curve ----------
persist_curve_parth <- summarize_persistence_curve(res_parth$extinction_times, max_time = t_max_long) %>%
  mutate(species = "parthenogen")

persist_curve_sex <- summarize_persistence_curve(res_sex$extinction_times, max_time = t_max_long) %>%
  mutate(species = "sexual")

persist_curve_all <- bind_rows(persist_curve_parth, persist_curve_sex)

## ---------- save outputs ----------
write.csv(traj_all, file.path(out_dir, "adult_population_trajectory_0_50y.csv"), row.names = FALSE)
write.csv(ext_summary_all, file.path(out_dir, "extinction_summary_all.csv"), row.names = FALSE)
write.csv(hpd_all, file.path(out_dir, "adult_population_HPD_selected_years.csv"), row.names = FALSE)
write.csv(band_all, file.path(out_dir, "adult_population_HDI_0_50y_band.csv"), row.names = FALSE)
write.csv(persist_all, file.path(out_dir, "persistence_proportions.csv"), row.names = FALSE)
write.csv(persist_curve_all, file.path(out_dir, "persistence_curve_1_2000y.csv"), row.names = FALSE)

## ---------- plot 1: trajectory 0-50 ----------
p1 <- ggplot(traj_all, aes(x = year, y = mean_adults, linetype = species)) +
  geom_line(linewidth = 1) +
  labs(
    x = "Year",
    y = "Mean adult population size",
    title = "Adult population size through first 50 years"
  ) +
  theme_bw()

ggsave(
  file.path(out_dir, "adult_population_trajectory_0_50y.png"),
  p1, width = 7, height = 5, dpi = 300
)

## ---------- plot 2: trajectory + 95% HDI ----------
p2 <- ggplot(band_all, aes(x = year, y = mean, linetype = species)) +
  geom_ribbon(
    aes(x = year, ymin = low, ymax = high, fill = species),
    alpha = 0.18,
    colour = NA
  ) +
  geom_line(linewidth = 1) +
  labs(
    x = "Year",
    y = "Adult population size",
    title = "Mean and 95% HDI of adult population size"
  ) +
  theme_bw()

ggsave(
  file.path(out_dir, "adult_population_HDI_0_50y.png"),
  p2, width = 7, height = 5, dpi = 300
)

## ---------- plot 3: long-term persistence log-scale + 95% intervals ----------
p3 <- ggplot(
  persist_curve_all,
  aes(x = year, y = proportion_survived, linetype = species)
) +
  geom_ribbon(
    aes(x = year, ymin = ci_low, ymax = ci_high, fill = species),
    alpha = 0.18,
    colour = NA
  ) +
  geom_line(linewidth = 1.1) +
  scale_x_log10(
    breaks = c(1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000)
  ) +
  coord_cartesian(ylim = c(0, 1)) +
  labs(
    x = "Year (log scale)",
    y = "Probability of population persistence",
    title = "Long-term persistence probability"
  ) +
  theme_bw()

ggsave(
  file.path(out_dir, "population_persistence_curve_logx_CI.png"),
  p3, width = 7, height = 5, dpi = 300
)

cat("\nDone.\n")
cat("All outputs written to:\n", out_dir, "\n")