lame_dev_path <- "/home/s7m/Research/netify_dev/lame"
if (dir.exists(lame_dev_path)) {
  pkgload::load_all(lame_dev_path, quiet = TRUE)
} else {
  library(lame)
}
stopifnot(packageVersion("lame") >= "1.3.4")
suppressPackageStartupMessages({
  library(netify)
  library(precrec)
})

set.seed(6886)
icews <- readRDS("data/icews.rds")
monadic_coverage <- unique(icews[, c("i", "year", "i_polity2", "i_log_gdp")])
coverage_counts <- aggregate(
  cbind(
    polity_years = as.integer(!is.na(i_polity2)),
    gdp_years = as.integer(!is.na(i_log_gdp))
  ) ~ i,
  monadic_coverage,
  sum
)
eligible_states <- coverage_counts$i[
  coverage_counts$polity_years == 13 &
    coverage_counts$gdp_years == 13
]
conflict_totals <- aggregate(matlConf ~ i, icews, sum)
conflict_totals <- conflict_totals[conflict_totals$i %in% eligible_states, ]
states <- head(
  conflict_totals$i[order(conflict_totals$matlConf, decreasing = TRUE)],
  18
)
icews_sub <- icews[icews$i %in% states & icews$j %in% states, ]
icews_sub$high_event_volume <- as.integer(icews_sub$matlConf > 20)
icews_sub$polity_gap <-
  abs(icews_sub$i_polity2 - icews_sub$j_polity2) / 20
icews_sub$same_region <-
  as.integer(icews_sub$i_region == icews_sub$j_region)
icews_dyad <- icews_sub[, c(
  "i", "j", "year", "matlConf", "high_event_volume",
  "polity_gap", "same_region"
)]
names(icews_dyad)[1:4] <- c(
  "source", "target", "year", "material_conflict_events"
)
icews_dyad$year <- as.integer(icews_dyad$year)
years <- sort(unique(icews_dyad$year))

icews_net <- netify(
  input = icews_dyad,
  actor1 = "source",
  actor2 = "target",
  time = "year",
  symmetric = FALSE,
  weight = "high_event_volume",
  dyad_vars = c("polity_gap", "same_region"),
  dyad_vars_symmetric = c(TRUE, TRUE),
  missing_to_zero = FALSE
)
stopifnot(all(unlist(validate_netify(icews_net))))
icews_lame <- to_lame(
  icews_net,
  lame = TRUE,
  family = "binary",
  fit_method = "als"
)
Y_icews <- icews_lame$Y
X_icews <- icews_lame$Xdyad

fit_r0 <- lame(
  Y = Y_icews,
  Xdyad = X_icews,
  family = "binary",
  symmetric = FALSE,
  R = 0,
  dynamic_ab = TRUE,
  dynamic_beta = "intercept",
  dynamic_beta_kind = "rw1",
  method = "als",
  als_max_iter = 1000,
  als_tol = 1e-5,
  als_stability = "validation",
  seed = 6886,
  verbose = FALSE
)

fit_r2 <- lame(
  Y = Y_icews,
  Xdyad = X_icews,
  family = "binary",
  symmetric = FALSE,
  R = 2,
  dynamic_ab = TRUE,
  dynamic_uv = TRUE,
  dynamic_beta = "intercept",
  dynamic_beta_kind = "rw1",
  method = "als",
  als_max_iter = 1000,
  als_tol = 1e-5,
  als_stability = "validation",
  bootstrap = 100,
  bootstrap_type = "parametric",
  bootstrap_seed = 6886,
  seed = 6886,
  verbose = FALSE
)

change_stats <- do.call(
  rbind,
  lapply(seq_along(years)[-1], function(tt) {
    before <- Y_icews[[tt - 1]]
    now <- Y_icews[[tt]]
    data.frame(
      year = years[tt],
      changes = sum(now != before, na.rm = TRUE),
      onsets = sum(now == 1 & before == 0, na.rm = TRUE),
      dissolutions = sum(now == 0 & before == 1, na.rm = TRUE)
    )
  })
)

fit_dynamic_rank <- function(Y, rank, seed) {
  lame(
    Y = Y,
    Xdyad = X_icews,
    family = "binary",
    symmetric = FALSE,
    R = rank,
    dynamic_ab = TRUE,
    dynamic_uv = rank > 0,
    dynamic_beta = "intercept",
    dynamic_beta_kind = "rw1",
    method = "als",
    als_max_iter = 1000,
    als_tol = 1e-5,
    als_stability = "none",
    seed = seed,
    verbose = FALSE
  )
}

pairs <- t(utils::combn(seq_len(nrow(Y_icews[[1]])), 2))
set.seed(6886)
fold_id <- sample(rep(seq_len(5), length.out = nrow(pairs)))
rank_fold_results <- list()
result_index <- 0L

for (rank in 0:4) {
  for (fold in seq_len(5)) {
    held_pairs <- pairs[fold_id == fold, , drop = FALSE]
    held_mask <- lapply(Y_icews, function(M) {
      out <- matrix(FALSE, nrow(M), ncol(M))
      for (kk in seq_len(nrow(held_pairs))) {
        ii <- held_pairs[kk, 1]
        jj <- held_pairs[kk, 2]
        out[ii, jj] <- TRUE
        out[jj, ii] <- TRUE
      }
      diag(out) <- FALSE
      out
    })
    Y_train <- Map(function(M, mask) {
      M[mask] <- NA_real_
      M
    }, Y_icews, held_mask)
    fit_fold <- fit_dynamic_rank(
      Y_train,
      rank = rank,
      seed = 6886 + 100L * rank + fold
    )
    score <- lame::evaluate_heldout(
      y_obs = Y_icews,
      y_pred = fit_fold$fitted,
      mask = held_mask,
      family = "binary"
    )
    result_index <- result_index + 1L
    rank_fold_results[[result_index]] <- data.frame(
      fold = fold,
      rank = rank,
      converged = fit_fold$converged,
      iterations = fit_fold$iterations,
      objective = tail(fit_fold$objective_history, 1),
      score
    )
  }
}

rank_folds <- do.call(rbind, rank_fold_results)
rank_summary <- do.call(
  rbind,
  lapply(split(rank_folds, rank_folds$rank), function(dd) {
    data.frame(
      rank = dd$rank[1],
      n_eval = mean(dd$n_eval),
      auroc = mean(dd$auroc),
      auprc = mean(dd$auprc),
      brier = mean(dd$brier),
      logloss = mean(dd$logloss),
      converged = sum(dd$converged)
    )
  })
)
rownames(rank_summary) <- NULL
rank_cv <- list(
  folds = rank_folds,
  summary = rank_summary,
  fold_id = fold_id,
  pairs = pairs
)

panel_gof <- function(fit, n_sim = 500, seed = 6886) {
  stat_one <- function(M) {
    c(
      density = mean(M, na.rm = TRUE),
      lame::gof_stats(M, mode = "unipartite")
    )
  }
  observed <- t(vapply(
    Y_icews,
    stat_one,
    numeric(6)
  ))
  simulations <- simulate(
    fit,
    nsim = n_sim,
    seed = seed
  )
  # lame 1.3.4 returns simulation metadata alongside the draws. Earlier
  # releases returned the draw list directly.
  if (inherits(simulations, "ame.sim")) {
    simulations <- simulations$Y
  }
  stat_array <- array(
    NA_real_,
    dim = c(length(years), 6, n_sim),
    dimnames = list(
      years,
      names(stat_one(Y_icews[[1]])),
      NULL
    )
  )
  changes <- matrix(
    NA_real_,
    nrow = length(years) - 1,
    ncol = n_sim,
    dimnames = list(years[-1], NULL)
  )
  for (ss in seq_len(n_sim)) {
    for (tt in seq_along(years)) {
      stat_array[tt, , ss] <- stat_one(simulations[[ss]][[tt]])
      if (tt > 1) {
        changes[tt - 1, ss] <- sum(
          simulations[[ss]][[tt]] !=
            simulations[[ss]][[tt - 1]],
          na.rm = TRUE
        )
      }
    }
  }
  lower <- apply(stat_array, c(1, 2), quantile, 0.025)
  upper <- apply(stat_array, c(1, 2), quantile, 0.975)
  coverage <- colSums(observed >= lower & observed <= upper)
  change_lower <- apply(changes, 1, quantile, 0.025)
  change_upper <- apply(changes, 1, quantile, 0.975)
  change_cover <- sum(
    change_stats$changes >= change_lower &
      change_stats$changes <= change_upper
  )
  list(
    sims = simulations,
    arr = stat_array,
    changes = changes,
    coverage = coverage,
    change_cover = change_cover
  )
}

gof <- list(
  r0 = panel_gof(fit_r0, seed = 6886),
  r2 = panel_gof(fit_r2, seed = 6890)
)

dir.create("cache", showWarnings = FALSE)
save(
  fit_r0,
  file = "cache/icews_dynamic_srm_r0_v1.rda",
  compress = "xz"
)
save(
  fit_r2,
  file = "cache/icews_dynamic_ame_r2_boot100_v1.rda",
  compress = "xz"
)
save(
  rank_cv,
  file = "cache/icews_dynamic_rank_cv_v1.rda",
  compress = "xz"
)
save(
  gof,
  file = "cache/icews_dynamic_gof_v1.rda",
  compress = "xz"
)

summary_out <- list(
  states = states,
  years = years,
  dyad = icews_dyad,
  Y = Y_icews,
  X = X_icews,
  change_stats = change_stats
)
save(
  summary_out,
  file = "cache/icews_dynamic_data_v1.rda",
  compress = "xz"
)

cat(
  "R0 converged:", fit_r0$converged,
  "\nR2 converged:", fit_r2$converged,
  "\nBootstrap successes:", fit_r2$bootstrap_dynamic$n_success,
  "\n"
)
