# =====================================================================================================
#  JKP Common Task Framework submission  --  main(chars, features, daily_ret)
#  Souveer Aggarwal, EK Im, JiWhan Moon, Abdul-Kabiru Nashiru
#  Yale SOM (MGT 924, Problem Set 3)
#
#  Pipeline (everything used at month-end t is realised by t):
#   1. Cross-sectional rank transform of all 402 characteristics (NA -> median). Target = next-month
#      excess return minus its cross-sectional mean; the tree model uses the rank of that target.
#   2. Return prediction = 0.10 * ridge + 0.45 * LightGBM(ranks) + 0.45 * LightGBM(raw), per-month standardised:
#        (a) expanding-window ridge from per-month sufficient statistics X'X, X'y, refit monthly,
#            penalty chosen by rank-IC on the last 60 months of the training window;
#        (b) LightGBM on ranked characteristics, rank target, refit every 12 months, #rounds by early
#            stopping on the same validation block, then refit on the full window;
#        (c) LightGBM on RAW characteristics with missing values kept, target atan(z-score of the
#            return), 500 fixed rounds -- the recipe of Paul Geertsema's public CTF entry (2026).
#      The two tree models see different information (ranks vs. levels) and their predictions
#      correlate only ~0.6, which is where the blend's gain comes from.
#   3. Barra-style risk model: characteristics clustered (hclust, pre-test months only) into K = 40
#      themes; exposures = sign-aligned cluster means; monthly cross-sectional WLS gives factor returns
#      and residuals; Sigma_t = B Sigma_f B' + D with EWMA estimates (12-month half-lives).
#   4. Weights ~ Sigma_t^{-1} mu_t (Woodbury), demeaned within 2-digit SIC industry (no net industry
#      bets; this step follows Abhinav Keshri's CTF entry, 2026), dollar neutral, gross exposure 2,
#      then a 12-month realised-volatility overlay (10% target, leverage capped at 3x).
#  Small-data behaviour (validation datasets): with fewer than 24 months of realised targets the
#  portfolio is equal-weight long; with fewer than 120 the ridge model alone is used with a fixed
#  penalty; the factor risk model needs 24 months of factor returns, otherwise a diagonal model is used.
#  Dependencies: data.table, lightgbm (renv.lock). Deterministic: fixed seeds and thread count.
#
#  CTF Admin Modifications (2026-10-07):
#  --------------------------------------
#  1. In ctf_prepare, the kept columns now always include ivol_capm_252d when
#     the data has it.
#     Reason: the validation run failed with "object ivol_capm_252d not
#     found". The validation features list has 10 names and omits that
#     column, so it was dropped before the risk model read it. The full
#     features list already contains it, so the full run is unchanged.
# =====================================================================================================
suppressPackageStartupMessages({ library(data.table); library(lightgbm) })

CFG <- list(
  seed            = 20260928L,
  refit_every     = 12L,        # months between LightGBM refits
  ridge_refit_every = 1L,       # ridge is refit monthly (cheap)
  val_months      = 60L,        # validation block at the end of the training window
  min_train_m     = 120L,       # months of realised targets needed for the full ridge + GBM model
  min_ridge_m     = 24L,        # ...for the ridge-only fallback
  fallback_lambda = 0.01,
  lambda_grid     = 10^seq(-4, 2, by = 0.5),
  lgb = list(objective = "regression", learning_rate = 0.05, num_leaves = 63, min_data_in_leaf = 1000,
             feature_fraction = 0.5, bagging_fraction = 0.5, bagging_freq = 1, lambda_l2 = 10,
             max_bin = 63, num_threads = 8, seed = 20260928, deterministic = TRUE,
             force_row_wise = TRUE, verbose = -1),
  lgb_max_rounds  = 1000L, lgb_patience = 50L, lgb_min_rounds = 50L,
  ens_w          = c(ridge = 0.10, lgb = 0.45, lgb2 = 0.45),   # ensemble weights (chosen on the development window)
  ens_w_lgb       = 0.75,       # ridge/LightGBM split used only when the second GBM is unavailable
  lgb2 = list(objective = "regression", learning_rate = 0.02, num_leaves = 255, min_data_in_leaf = 2000,
              feature_fraction = 1.0, bagging_fraction = 1.0, bagging_freq = 0, lambda_l2 = 0,
              max_bin = 255, num_threads = 8, seed = 20260928, deterministic = TRUE,
              force_row_wise = TRUE, verbose = -1),
  lgb2_rounds     = 500L,       # fixed number of rounds for the raw-feature GBM (no early stopping)
  K_factors       = 40L,
  hl_factor       = 12,         # EWMA half-life (months), factor covariance
  hl_idio         = 12,         # EWMA half-life (months), idiosyncratic variance
  min_factor_m    = 24L,        # factor-return months needed before the factor model is used
  ind_neutral     = TRUE,       # remove net 2-digit-SIC industry bets from the weights each month (after Keshri, CTF 2026)
  vol_target      = 0.10 / sqrt(12), vol_lookback = 12L, vol_max_lev = 3,
  gross           = 2
)

# ------------------------------------------------------------------------------------------ helpers
to_logical_flag <- function(x) {
  if (is.logical(x)) return(x)
  if (is.numeric(x)) return(x == 1)
  tolower(trimws(as.character(x))) %in% c("1", "true", "t")
}
rank01 <- function(v) {                 # cross-sectional rank in (-0.5, 0.5); NA -> 0
  ok <- !is.na(v); if (!any(ok)) return(rep(0, length(v)))
  r <- rep(0, length(v)); r[ok] <- frank(v[ok], ties.method = "average") / (sum(ok) + 1) - 0.5; r
}
ewma_weights <- function(n, hl) { w <- (0.5^(1 / hl))^(rev(seq_len(n)) - 1); w / sum(w) }
ewma_by_group <- function(x, hl) {      # running EWMA of x (used on squared residuals per stock)
  lam <- 0.5^(1 / hl); v <- numeric(length(x)); s <- NA_real_
  for (k in seq_along(x)) { s <- if (is.na(s)) x[k] else lam * s + (1 - lam) * x[k]; v[k] <- s }
  v
}

# ------------------------------------------------------------------------------------ 1. preparation
ctf_prepare <- function(chars, features) {
  chars <- as.data.table(chars)
  feats <- as.character(if (is.data.frame(features)) features[[1]] else features)
  feats <- intersect(feats, names(chars))
  chars[, `:=`(eom = as.IDate(eom), eom_ret = as.IDate(eom_ret), id = as.integer(id),
               ctff_test = to_logical_flag(ctff_test))]
  # CTF-FIX: keep ivol_capm_252d even when the features list omits it (validation dataset)
  keep <- unique(c("id", "eom", "eom_ret", "ctff_test", "ret_exc_lead1m", feats,
                   intersect(c("sic", "ivol_capm_252d"), names(chars))))
  chars <- chars[, ..keep]
  setorder(chars, eom, id)
  chars[, ind := if ("sic" %in% names(chars)) substr(as.character(sic), 1, 2) else "00"]; chars[is.na(ind) | ind == "", ind := "00"]
  chars[, ivol_raw := ivol_capm_252d]                         # raw daily idiosyncratic vol (risk model)
  chars[, ivol_m := fifelse(is.na(ivol_raw), median(ivol_raw, na.rm = TRUE), ivol_raw), by = eom]
  chars[is.na(ivol_m), ivol_m := 0.02]
  chars[, ivol_m := pmax(ivol_m, quantile(ivol_m, 0.05)), by = eom]
  raw_cols <- paste0("raw_", feats); chars[, (raw_cols) := .SD, .SDcols = feats]   # raw values, NA kept (second GBM)
  chars[, (feats) := lapply(.SD, rank01), by = eom, .SDcols = feats]
  chars[, y := ret_exc_lead1m - mean(ret_exc_lead1m, na.rm = TRUE), by = eom]   # market-relative return
  chars[, y_rank := rank01(y), by = eom]                                         # rank target (first GBM)
  chars[, y_atan := atan(y / sd(y, na.rm = TRUE)), by = eom]                     # atan(z) target (second GBM)
  list(chars = chars, feats = feats)
}

# ------------------------------------------------------------------------------------ 2. predictions
ctf_predict <- function(chars, feats, cfg = CFG, log = message) {
  months <- sort(unique(chars$eom)); p <- length(feats)
  XtX <- vector("list", length(months)); Xty <- XtX; nobs <- integer(length(months))
  ret_month <- as.IDate(rep(NA_integer_, length(months)))
  for (i in seq_along(months)) {
    d <- chars[eom == months[i] & !is.na(y)]
    if (nrow(d) == 0) { XtX[[i]] <- matrix(0, p, p); Xty[[i]] <- numeric(p); next }
    X <- as.matrix(d[, ..feats]); XtX[[i]] <- crossprod(X); Xty[[i]] <- crossprod(X, d$y)[, 1]
    nobs[i] <- nrow(d); ret_month[i] <- d$eom_ret[1]
  }
  sum_stats <- function(idx) { A <- matrix(0, p, p); b <- numeric(p)
    for (i in idx) { A <- A + XtX[[i]]; b <- b + Xty[[i]] }; list(A = A, b = b, n = sum(nobs[idx])) }
  ridge <- function(S, lambda) solve(S$A + diag(lambda * S$n, p), S$b)

  test_months <- sort(unique(chars[ctff_test == TRUE, eom]))
  pred <- vector("list", length(test_months)); fit_log <- list(); beta <- NULL; booster <- NULL; booster2 <- NULL
  raw_cols <- paste0("raw_", feats)
  for (k in seq_along(test_months)) {
    t <- test_months[k]
    known <- which(!is.na(ret_month) & ret_month <= t)            # months whose target is realised by t
    n_val <- min(cfg$val_months, length(known) %/% 3L)
    mode  <- if (length(known) >= cfg$min_train_m) "full" else if (length(known) >= cfg$min_ridge_m) "ridge" else "none"
    d <- chars[eom == t]
    if (mode == "none") { pred[[k]] <- data.table(id = d$id, eom = d$eom, pred_ridge = NA_real_, pred_lgb = NA_real_, pred_lgb2 = NA_real_); next }
    val <- tail(known, n_val); trn <- setdiff(known, val)
    # (a) ridge
    if ((k - 1L) %% cfg$ridge_refit_every == 0L || is.null(beta)) {
      if (mode == "full") {
        S_trn <- sum_stats(trn)
        v <- chars[eom %in% months[val] & !is.na(y), c("eom", "y", "y_rank", feats), with = FALSE]; v_X <- as.matrix(v[, ..feats])
        ic <- vapply(cfg$lambda_grid, function(l) { v[, ph := as.numeric(v_X %*% ridge(S_trn, l))]
          mean(v[, cor(ph, y, method = "spearman"), by = eom]$V1) }, numeric(1))
        lam <- cfg$lambda_grid[which.max(ic)]
      } else lam <- cfg$fallback_lambda
      beta <- ridge(sum_stats(known), lam)
    }
    # (b) LightGBM
    if (mode == "full" && ((k - 1L) %% cfg$refit_every == 0L || is.null(booster))) {
      t0 <- Sys.time()
      v  <- chars[eom %in% months[val] & !is.na(y), c("eom", "y", "y_rank", feats), with = FALSE]; v_X <- as.matrix(v[, ..feats])
      tr <- chars[eom %in% months[trn] & !is.na(y), c("y_rank", feats), with = FALSE]
      d_trn <- lgb.Dataset(as.matrix(tr[, ..feats]), label = tr$y_rank, params = list(max_bin = cfg$lgb$max_bin))
      d_val <- lgb.Dataset(v_X, label = v$y_rank, reference = d_trn); rm(tr); gc()
      b1 <- lgb.train(params = cfg$lgb, data = d_trn, nrounds = cfg$lgb_max_rounds, valids = list(val = d_val),
                      early_stopping_rounds = cfg$lgb_patience, verbose = -1)
      best <- max(b1$best_iter, cfg$lgb_min_rounds); rm(b1, d_trn, d_val); gc()
      al <- chars[eom %in% months[known] & !is.na(y), c("y_rank", feats), with = FALSE]
      d_all <- lgb.Dataset(as.matrix(al[, ..feats]), label = al$y_rank, params = list(max_bin = cfg$lgb$max_bin)); rm(al); gc()
      booster <- lgb.train(params = cfg$lgb, data = d_all, nrounds = best, verbose = -1); rm(d_all); gc()
      # (c) second LightGBM: raw characteristics (NA kept), atan(z) target, fixed rounds -- after Geertsema (CTF, 2026)
      al <- chars[eom %in% months[known] & !is.na(y), c("y_atan", raw_cols), with = FALSE]
      d_raw <- lgb.Dataset(as.matrix(al[, ..raw_cols]), label = al$y_atan, params = list(max_bin = cfg$lgb2$max_bin)); rm(al); gc()
      booster2 <- lgb.train(params = cfg$lgb2, data = d_raw, nrounds = cfg$lgb2_rounds, verbose = -1); rm(d_raw); gc()
      fit_log[[length(fit_log) + 1L]] <- data.table(eom = t, lambda = lam, lgb_rounds = best, n_train = sum(nobs[known]),
                                                   mins = as.numeric(difftime(Sys.time(), t0, units = "mins")))
      log(sprintf("[%s] refit %s: lambda=%.4g lgb_rounds=%d (%.1f min)", format(Sys.time(), "%H:%M"), t, lam, best,
                  tail(fit_log, 1)[[1]]$mins))
    }
    X <- as.matrix(d[, ..feats])
    pred[[k]] <- data.table(id = d$id, eom = d$eom, pred_ridge = as.numeric(X %*% beta),
                            pred_lgb  = if (mode == "full") as.numeric(predict(booster, X)) else NA_real_,
                            pred_lgb2 = if (mode == "full") as.numeric(predict(booster2, as.matrix(d[, ..raw_cols]))) else NA_real_)
  }
  pred <- rbindlist(pred)
  pred[, pred := {
    zr <- if (all(is.na(pred_ridge))) NULL else scale(pred_ridge)[, 1]
    zl <- if (all(is.na(pred_lgb)))   NULL else scale(pred_lgb)[, 1]
    z2 <- if (all(is.na(pred_lgb2)))  NULL else scale(pred_lgb2)[, 1]
    w  <- cfg$ens_w / sum(cfg$ens_w)
    if (is.null(zr)) rep(NA_real_, .N) else if (is.null(zl)) zr
    else if (is.null(z2)) (1 - cfg$ens_w_lgb) * zr + cfg$ens_w_lgb * zl
    else w["ridge"] * zr + w["lgb"] * zl + w["lgb2"] * z2 }, by = eom]
  list(pred = pred, fit_log = rbindlist(fit_log))
}

# ------------------------------------------------------------------------------------- 3. risk model
ctf_risk_model <- function(chars, feats, cfg = CFG, log = message) {
  pre_months <- if (any(!chars$ctff_test)) chars[ctff_test == FALSE, unique(eom)] else head(sort(unique(chars$eom)), 24L)
  pre <- chars[eom %in% pre_months, ..feats]
  C <- suppressWarnings(cor(as.matrix(pre))); C[is.na(C)] <- 0; diag(C) <- 1; rm(pre)
  hc <- hclust(as.dist(1 - abs(C)), method = "ward.D2"); cl <- cutree(hc, min(cfg$K_factors, length(feats)))
  cdefs <- lapply(split(names(cl), cl), function(m) {
    ref <- m[which.max(rowSums(abs(C[m, m, drop = FALSE])))]; list(feature = m, sign = sign(C[ref, m])) })
  fcols <- paste0("F", seq_along(cdefs))
  expo <- chars[, .(id, eom, eom_ret, ctff_test, ret_exc_lead1m, ivol = ivol_m, ind)]
  for (j in seq_along(cdefs)) { cd <- cdefs[[j]]
    expo[, (fcols[j]) := as.numeric(as.matrix(chars[, cd$feature, with = FALSE]) %*% cd$sign) / max(1, sum(abs(cd$sign)))] }
  months <- sort(unique(expo$eom)); fr <- list(); resid <- list()
  for (i in seq_along(months)) {
    d <- expo[eom == months[i] & !is.na(ret_exc_lead1m)]
    if (nrow(d) < 50) next
    B <- cbind(1, as.matrix(d[, ..fcols])); wgt <- 1 / d$ivol^2
    f <- solve(crossprod(B * sqrt(wgt)) + diag(1e-8, ncol(B)), crossprod(B * wgt, d$ret_exc_lead1m))
    fr[[i]] <- data.table(eom = months[i], eom_ret = d$eom_ret[1], t(f[, 1]))
    resid[[i]] <- data.table(id = d$id, avail = d$eom_ret[1], e2 = (d$ret_exc_lead1m - B %*% f)[, 1]^2)
  }
  fr <- rbindlist(fr); resid <- rbindlist(resid)
  if (nrow(resid)) { setorder(resid, id, avail); resid[, e2_ewma := ewma_by_group(e2, cfg$hl_idio), by = id]; setkey(resid, id, avail) }
  log(sprintf("risk model: %d factors, %d factor-return months", length(fcols), nrow(fr)))
  list(expo = expo, fcols = fcols, fr = fr, fmat = if (nrow(fr)) as.matrix(fr[, -(1:2)]) else NULL, resid = resid, clusters = cdefs)
}

# ---------------------------------------------------------------------------------------- 4. weights
ctf_weights <- function(pred, RM, cfg = CFG) {
  test_months <- sort(unique(pred$eom)); out <- vector("list", length(test_months))
  for (k in seq_along(test_months)) {
    t <- test_months[k]
    d <- merge(RM$expo[eom == t, c("id", "ivol", "ind", RM$fcols), with = FALSE], pred[eom == t, .(id, pred)], by = "id")
    if (all(is.na(d$pred))) { out[[k]] <- data.table(id = d$id, eom = t, w = 1 / nrow(d)); next }   # no model yet
    m_ivol <- (d$ivol * sqrt(21))^2
    if (nrow(RM$resid)) { d[, avail := t]; d[, e2_ewma := RM$resid[d[, .(id, avail)], roll = TRUE, on = .(id, avail)]$e2_ewma]
    } else d[, e2_ewma := NA_real_]
    d[, dvar := fifelse(is.na(e2_ewma), m_ivol, 0.5 * e2_ewma + 0.5 * m_ivol)]
    d[, dvar := pmax(dvar, quantile(dvar, 0.05))]
    mu <- d$pred - mean(d$pred); Di <- 1 / d$dvar
    F <- if (is.null(RM$fmat)) NULL else RM$fmat[RM$fr$eom_ret <= t, , drop = FALSE]
    if (!is.null(F) && nrow(F) >= cfg$min_factor_m) {
      wts <- ewma_weights(nrow(F), cfg$hl_factor)
      Sf <- crossprod(sqrt(wts) * (F - rep(colSums(F * wts), each = nrow(F)))); Sf <- Sf + diag(1e-6 * mean(diag(Sf)), ncol(Sf))
      B <- cbind(1, as.matrix(d[, RM$fcols, with = FALSE]))
      M <- solve(Sf) + crossprod(B * sqrt(Di))
      w <- Di * mu - Di * (B %*% solve(M, crossprod(B, Di * mu)))[, 1]                # Sigma^{-1} mu via Woodbury
    } else w <- Di * mu                                                                # diagonal fallback
    if (cfg$ind_neutral) w <- w - ave(w, d$ind, FUN = mean)                           # no net industry bets
    w <- w - mean(w); w <- w / sum(abs(w)) * cfg$gross
    out[[k]] <- data.table(id = d$id, eom = t, w = w)
  }
  pf <- rbindlist(out)
  # volatility overlay from the strategy's own realised returns (all realised before t)
  x <- merge(pf, RM$expo[, .(id, eom, ret_exc_lead1m)], by = c("id", "eom"), all.x = TRUE)
  m <- x[, .(ret = sum(w * fifelse(is.na(ret_exc_lead1m), 0, ret_exc_lead1m))), by = eom][order(eom)]
  m[, vol_prev := shift(frollapply(ret, cfg$vol_lookback, sd), 1)]
  m[, k := fifelse(is.na(vol_prev) | !(vol_prev > 0), 1, pmin(cfg$vol_target / vol_prev, cfg$vol_max_lev))]
  pf <- merge(pf, m[, .(eom, k)], by = "eom")[, .(id, eom, w = w * k)]
  pf[w != 0]
}

# -------------------------------------------------------------------------------------- entry point
main <- function(chars, features, daily_ret) {
  set.seed(CFG$seed)
  log <- function(m) { cat(m, "\n"); flush.console() }
  prep <- ctf_prepare(chars, features); rm(chars); gc()
  out  <- ctf_predict(prep$chars, prep$feats, CFG, log)
  RM   <- ctf_risk_model(prep$chars, prep$feats, CFG, log)
  pf   <- ctf_weights(out$pred, RM, CFG)
  data.frame(id = as.integer(pf$id), eom = as.Date(pf$eom), w = as.numeric(pf$w))
}

# ---- local test (same convention as the competition template; not executed by the harness) -------
if (interactive()) {
  library(arrow)
  chars     <- read_parquet("ctff_chars.parquet")
  features  <- read_parquet("ctff_features.parquet")
  daily_ret <- read_parquet("ctff_daily_ret.parquet")
  pf <- main(chars, features, daily_ret)
  write.csv(pf, "output.csv", row.names = FALSE)
}