# =====================================================================================
# CTF submission -- "Risk-Neutralised Ridge + LightGBM Ensemble"
# MGT 924 Statistical Foundations, Problem Set 3
#
# Entry point: main(chars, features, daily_ret) -> data.frame(id, eom, w)
#
# Pipeline (all steps use only information available at the portfolio formation date t):
#   1. Features: 402 supplied characteristics + 2 daily-return features (5-day return and 5-day RMS
#      return, from daily returns dated inside month t). Centred cross-sectional ECDF ranks; NA -> 0.
#   2. Alpha models, refit every December on all rows whose next-month return is realised (eom <= t-1):
#        a) ridge on the raw excess-return target, lambda chosen by Sharpe ratio of the linear
#           characteristic portfolio over the last 60 training months (inner chronological validation)
#        b) LightGBM on the cross-sectional rank of returns, #trees by early stopping on the last 60 months
#   3. Forecast combination: equal-weight average of within-month ranks -> Gaussian score z.
#   4. Risk model Sigma = B F B' + D: 12 Fama-French industries + 30 hierarchical-cluster style factors,
#      EWMA factor covariance (half-life 36m), idiosyncratic variance from ivol_capm_252d.
#   5. Weights w = Sigma^{-1} z, scaled to 1% ex-ante monthly volatility.
# Deterministic: fixed seeds, LightGBM deterministic mode with a fixed thread count.
# =====================================================================================
suppressPackageStartupMessages({
  library(data.table)
  library(lightgbm)
})

PARAMS <- list(
  V          = 60L,                     # inner validation window (months)
  min_train  = 24L,                     # minimum training months for any fit
  ridge_grid = 10^seq(-3, 2, by = 0.5),
  ridge_default = 0.01,
  lgb = list(objective = "regression", num_leaves = 31L, min_data_in_leaf = 2000L, learning_rate = 0.05,
             feature_fraction = 0.5, bagging_fraction = 0.5, bagging_freq = 1L, lambda_l2 = 10,
             seed = 1L, deterministic = TRUE, force_col_wise = TRUE, num_threads = 16L, verbose = -1L),
  max_bin    = 63L, lgb_max_rounds = 1000L, lgb_es = 50L, lgb_min_rounds = 20L, lgb_default_rounds = 100L,
  K          = 30L,                     # style clusters in the risk model
  LB         = 120L,                    # lookback for clustering and factor covariance (months)
  HL         = 36,                      # EWMA half-life of factor covariance (months)
  idio_lb    = 60L,                     # lookback for idiosyncratic-variance calibration (months)
  target_vol = 0.01,                    # ex-ante monthly portfolio volatility
  SCALE      = 1e6                      # ranks stored as integers (rank * SCALE) to save memory
)

log_msg <- function(...) { cat(format(Sys.time(), "%H:%M:%S"), ..., "\n"); flush.console() }

# ---------- helpers ----------
cs_rank <- function(x, g) {               # centred within-group ECDF rank, NA -> 0
  dt <- data.table(x = as.numeric(x), g = g)
  dt[, r := {
    k <- sum(!is.na(x))
    if (k == 0L) rep(0, .N) else (frank(x, ties.method = "average", na.last = "keep") - 0.5) / k - 0.5
  }, by = g]
  r <- dt$r; r[is.na(r)] <- 0; r
}
gauss_score <- function(p) qnorm((frank(p, ties.method = "average") - 0.5) / length(p))

as_test_flag <- function(x) {
  if (is.logical(x)) return(!is.na(x) & x)
  if (is.numeric(x)) return(!is.na(x) & x == 1)
  x <- toupper(trimws(as.character(x))); x %in% c("1", "TRUE", "T")
}

ff12_from_sic <- function(sic) {
  s <- sic; out <- rep(12L, length(s)); b <- function(lo, hi) !is.na(s) & s >= lo & s <= hi
  out[b(100, 999) | b(2000, 2399) | b(2700, 2749) | b(2770, 2799) | b(3100, 3199) | b(3940, 3989)] <- 1L
  out[b(2500, 2519) | b(2590, 2599) | b(3630, 3659) | b(3710, 3711) | b(3714, 3714) | b(3716, 3716) |
      b(3750, 3751) | b(3792, 3792) | b(3900, 3939) | b(3990, 3999)] <- 2L
  out[b(2520, 2589) | b(2600, 2699) | b(2750, 2769) | b(3000, 3099) | b(3200, 3569) | b(3580, 3629) |
      b(3700, 3709) | b(3712, 3713) | b(3715, 3715) | b(3717, 3749) | b(3752, 3791) | b(3793, 3799) |
      b(3830, 3839) | b(3860, 3899)] <- 3L
  out[b(1200, 1399) | b(2900, 2999)] <- 4L
  out[b(2800, 2829) | b(2840, 2899)] <- 5L
  out[b(3570, 3579) | b(3660, 3692) | b(3694, 3699) | b(3810, 3829) | b(7370, 7379)] <- 6L
  out[b(4800, 4899)] <- 7L
  out[b(4900, 4949)] <- 8L
  out[b(5000, 5999) | b(7200, 7299) | b(7600, 7699)] <- 9L
  out[b(2830, 2839) | b(3693, 3693) | b(3840, 3859) | b(8000, 8099)] <- 10L
  out[b(6000, 6999)] <- 11L
  out
}

# Daily features for (id, calendar month): compounded return and RMS return over the last 5 trading
# days of the month. Uses only daily returns dated inside month t (all <= eom of t).
daily_features <- function(daily_ret, keys, first_date) {
  id <- as.integer(daily_ret$id); date <- as.Date(daily_ret$date); ret <- as.numeric(daily_ret$ret_exc)
  keep <- which(date >= first_date & !is.na(ret)); id <- id[keep]; date <- date[keep]; ret <- ret[keep]
  o <- order(id, date); id <- id[o]; date <- date[o]; ret <- ret[o]; rm(o, keep)
  ym <- year(date) * 12L + month(date); rm(date)
  brk <- c(which(diff(id) != 0L | diff(ym) != 0L), length(id))     # last row of each (id, ym) group
  grp_end <- rep(brk, diff(c(0L, brk)))
  last5 <- (grp_end - seq_along(id)) < 5L
  d <- data.table(id = id[last5], ym = ym[last5], ret = ret[last5]); rm(id, ym, ret, grp_end, last5)
  f <- d[, .(ret_5d = exp(sum(log1p(pmax(ret, -0.99)))) - 1, vol_5d = sqrt(mean(ret^2))), keyby = .(id, ym)]
  f[keys, on = .(id, ym)][, .(ret_5d, vol_5d)]
}

# Sigma^{-1} a with Sigma = B F B' + diag(d), via Woodbury (no N x N matrices)
sigma_inv_a <- function(a, B, Fm, d) {
  Da <- a / d; DB <- B / d
  G <- diag(ncol(B)) + Fm %*% crossprod(B, DB)
  drop(Da - DB %*% solve(G, Fm %*% crossprod(B, Da)))
}
port_var <- function(w, B, Fm, d) { bw <- crossprod(B, w); drop(crossprod(bw, Fm %*% bw) + sum(w^2 * d)) }

# ---------- main ----------
main <- function(chars, features, daily_ret) {
  set.seed(1L)
  P <- PARAMS
  log_msg("start")
  feat <- as.character(features$features)

  # shallow copy (no data duplication); caller's objects are never modified
  dt <- as.list(chars); rm(chars); setDT(dt)
  feat <- feat[feat %in% names(dt)]
  dt[, eom := as.Date(eom)]
  dt[, id := as.integer(id)]
  o <- dt[, order(eom, id)]
  meta <- data.table(id = dt$id[o], eom = dt$eom[o], ret = as.numeric(dt$ret_exc_lead1m[o]),
                     test = as_test_flag(dt$ctff_test[o]),
                     ff12 = ff12_from_sic(suppressWarnings(as.integer(as.character(dt$sic[o])))),
                     ivol = as.numeric(dt$ivol_capm_252d[o]))
  months <- sort(unique(meta$eom)); M <- length(months)
  meta[, mi := match(eom, months)]
  n <- nrow(meta); p0 <- length(feat); p <- p0 + 2L
  log_msg("rows", n, "months", M, "features", p0, "test months", meta[test == TRUE, uniqueN(mi)])

  # 1) rank features -> integer matrix (rank * SCALE)
  Xi <- matrix(0L, n, p)
  for (j in seq_len(p0)) {
    Xi[, j] <- as.integer(round(cs_rank(dt[[feat[j]]][o], meta$mi) * P$SCALE))
    set(dt, j = feat[j], value = NULL)
  }
  rm(dt); gc(verbose = FALSE)
  log_msg("characteristics ranked")
  df <- daily_features(daily_ret, meta[, .(id, ym = year(eom) * 12L + month(eom))],
                       first_date = min(months) - 40)
  rm(daily_ret); gc(verbose = FALSE)
  Xi[, p0 + 1L] <- as.integer(round(cs_rank(df$ret_5d, meta$mi) * P$SCALE))
  Xi[, p0 + 2L] <- as.integer(round(cs_rank(df$vol_5d, meta$mi) * P$SCALE))
  rm(df)
  log_msg("daily features built")
  Xd <- function(r) Xi[r, , drop = FALSE] / P$SCALE      # double copy of the requested rows only

  rows_of <- split(seq_len(n), factor(meta$mi, levels = seq_len(M)))
  nm <- lengths(rows_of)
  first_row <- c(1L, cumsum(nm)[-M] + 1L); last_row <- cumsum(nm)      # rows are sorted by month
  yrank <- cs_rank(meta$ret, meta$mi)
  ret0 <- meta$ret; ret0[is.na(ret0)] <- 0

  # 2) refit schedule: every December from the last December at/before the first test month
  test_mi <- sort(unique(meta[test == TRUE, mi]))
  if (!length(test_mi)) stop("no test months")
  dec <- which(month(months) == 12L)
  r0 <- if (any(dec <= test_mi[1])) max(dec[dec <= test_mi[1]]) else test_mi[1]
  refits <- sort(unique(c(r0, dec[dec > r0 & dec <= max(test_mi)])))
  blocks <- lapply(seq_along(refits), function(i) {
    hi <- if (i < length(refits)) refits[i + 1] - 1L else M
    intersect(refits[i]:hi, test_mi) })

  # ---- monthly sufficient statistics for ridge (cumulative snapshots only where needed)
  need <- sort(unique(c(refits - 1L, refits - 1L - P$V, refits - 1L - P$LB)))
  need <- need[need >= 1L]
  snapS <- list(); snaps <- list(); snapN <- list()
  Xy <- matrix(0, M, p)
  S <- matrix(0, p, p); s <- numeric(p); N <- 0
  for (m in seq_len(max(need, 1L))) {
    r <- rows_of[[m]]; ok <- r[!is.na(meta$ret[r])]
    if (length(ok)) {
      Xm <- Xd(ok); S <- S + crossprod(Xm); Xy[m, ] <- crossprod(Xm, meta$ret[ok])
      s <- s + Xy[m, ]; N <- N + length(ok)
    }
    if (m %in% need) { k <- as.character(m); snapS[[k]] <- S; snaps[[k]] <- s; snapN[[k]] <- N }
  }
  getS <- function(m) if (m < 1L) list(S = matrix(0, p, p), s = numeric(p), N = 0) else
    list(S = snapS[[as.character(m)]], s = snaps[[as.character(m)]], N = snapN[[as.character(m)]])
  ridge_solve <- function(st, lambda) {
    A <- st$S / st$N; b <- st$s / st$N; sc <- mean(diag(A))
    if (length(lambda) == 1L) return(solve(A + diag(lambda * sc, p), b))
    e <- eigen(A, symmetric = TRUE); q <- crossprod(e$vectors, b)
    sapply(lambda, function(l) e$vectors %*% (q / (e$values + l * sc)))
  }
  log_msg("ridge statistics done")

  out <- vector("list", length(refits))
  for (i in seq_along(refits)) {
    r <- refits[i]; blk <- blocks[[i]]
    if (!length(blk)) next
    ntrain_m <- r - 1L
    rows_blk <- unlist(rows_of[blk], use.names = FALSE)
    if (ntrain_m < P$min_train) {                       # not enough history: no position
      out[[i]] <- data.table(row = rows_blk, w = 0); next
    }
    # ---- ridge
    lam <- P$ridge_default
    if (r - 1L - P$V >= P$min_train) {
      Bg <- ridge_solve(getS(r - 1L - P$V), P$ridge_grid)
      vm <- (r - P$V):(r - 1L)
      sc <- apply(Bg, 2, function(b) { pr <- (Xy[vm, , drop = FALSE] %*% b) / nm[vm]; mean(pr) / sd(pr) })
      sc[!is.finite(sc)] <- -Inf
      lam <- P$ridge_grid[which.max(sc)]
    }
    b_ridge <- ridge_solve(getS(r - 1L), lam)
    # ---- LightGBM (training rows are a contiguous prefix because rows are sorted by month)
    nr <- P$lgb_default_rounds
    if (r - 1L - P$V >= P$min_train) {
      tr_in <- 1L:last_row[r - 1L - P$V]; va <- first_row[r - P$V]:last_row[r - 1L]
      dtr <- lgb.Dataset(Xd(tr_in), label = yrank[tr_in], params = list(max_bin = P$max_bin))
      dva <- lgb.Dataset.create.valid(dtr, Xd(va), label = yrank[va])
      # Early-stopping metric computed in R: LightGBM's multithreaded metric reduction is not
      # bit-reproducible (differences ~1e-17 can flip the argmin on a flat validation curve),
      # whereas its predictions are. This keeps the selected number of trees deterministic.
      yva <- yrank[va]
      l2_eval <- function(preds, dtrain) list(name = "l2r", value = mean((preds - yva)^2), higher_better = FALSE)
      m1 <- lgb.train(c(P$lgb, metric = "None"), dtr, nrounds = P$lgb_max_rounds, valids = list(va = dva),
                      eval = l2_eval, early_stopping_rounds = P$lgb_es, verbose = -1L)
      nr <- max(P$lgb_min_rounds, m1$best_iter); rm(dtr, dva, m1); gc(verbose = FALSE)
    }
    tr <- 1L:last_row[r - 1L]
    dfull <- lgb.Dataset(Xd(tr), label = yrank[tr], params = list(max_bin = P$max_bin))
    mod <- lgb.train(P$lgb, dfull, nrounds = nr, verbose = -1L); rm(dfull); gc(verbose = FALSE)

    # ---- risk model: clusters from rank correlations over the last LB training months
    st_hi <- getS(r - 1L); st_lo <- getS(r - 1L - P$LB)
    Sc <- (st_hi$S - st_lo$S)[1:p0, 1:p0]
    okf <- diag(Sc) > 1e-12
    C <- cov2cor(Sc[okf, okf])
    hc <- hclust(as.dist(1 - abs(C)), method = "ward.D2")
    K <- min(P$K, sum(okf) - 1L)
    cl <- cutree(hc, K)
    L <- matrix(0, p0, K)
    for (k in 1:K) {
      jj <- which(cl == k)
      v <- if (length(jj) == 1L) 1 else eigen(C[jj, jj, drop = FALSE], symmetric = TRUE)$vectors[, 1]
      if (sum(v) < 0) v <- -v
      L[which(okf)[jj], k] <- sign(v) / length(jj)
    }
    expo <- function(m) {
      rr <- rows_of[[m]]
      Z <- Xd(rr)[, 1:p0, drop = FALSE] %*% L
      Z <- apply(Z, 2, function(z) ((frank(z) - 0.5) / length(z) - 0.5) / sqrt(1 / 12))
      cbind(outer(meta$ff12[rr], 1:12, "==") * 1, Z)
    }
    nf <- 12L + K
    hist <- max(1L, r - P$LB):(max(blk) - 1L)
    fr <- matrix(NA_real_, M, nf); e2 <- iv <- rep(NA_real_, M)
    ivol2 <- 21 * meta$ivol^2
    for (m in hist) {
      rr <- rows_of[[m]]; B <- expo(m); y <- ret0[rr]
      G <- crossprod(B); diag(G) <- diag(G) + 1e-8
      f <- solve(G, crossprod(B, y)); fr[m, ] <- f
      e2[m] <- mean((y - B %*% f)^2)
      v <- ivol2[rr]; v[!is.finite(v)] <- NA; v[is.na(v)] <- median(v, na.rm = TRUE); iv[m] <- mean(v)
    }

    # ---- forecasts and weights for each test month in the block
    res <- vector("list", length(blk))
    for (j in seq_along(blk)) {
      k <- blk[j]; rr <- rows_of[[k]]; Xk <- Xd(rr)
      a1 <- drop(Xk %*% b_ridge); a2 <- predict(mod, Xk)
      cr <- function(x) (frank(x) - 0.5) / length(x) - 0.5
      z <- gauss_score((cr(a1) + cr(a2)) / 2)
      past <- max(1L, k - P$LB):(k - 1L); past <- past[!is.na(fr[past, 1])]
      wt <- 0.5^((k - 1L - past) / P$HL); wt <- wt / sum(wt)
      Fm <- crossprod(fr[past, , drop = FALSE] * sqrt(wt))
      pp <- max(1L, k - P$idio_lb):(k - 1L); pp <- pp[!is.na(e2[pp])]
      cf <- if (length(pp)) mean(e2[pp]) / mean(iv[pp]) else 1
      d <- ivol2[rr]; d[!is.finite(d)] <- NA; d[is.na(d)] <- median(d, na.rm = TRUE)
      if (!is.finite(cf) || cf <= 0) cf <- 1
      d <- d * cf; d[!is.finite(d) | d <= 0] <- median(d[is.finite(d) & d > 0])
      B <- expo(k)
      w <- sigma_inv_a(z, B, Fm, d)
      w <- w / sum(abs(w))
      w <- w * P$target_vol / sqrt(port_var(w, B, Fm, d))
      if (any(!is.finite(w))) w <- (z - mean(z)) / sum(abs(z - mean(z)))
      res[[j]] <- data.table(row = rr, w = w)
    }
    out[[i]] <- rbindlist(res)
    log_msg("refit", format(months[r]), "lambda", lam, "trees", nr, "months", length(blk))
    rm(mod); gc(verbose = FALSE)
  }
  out <- rbindlist(out)
  res <- data.frame(id = meta$id[out$row], eom = meta$eom[out$row], w = as.numeric(out$w))
  res <- res[order(res$eom, res$id), ]; rownames(res) <- NULL
  log_msg("done:", nrow(res), "weights")
  res
}

# Local testing only (not executed by the CTF pipeline)
#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)
#}