# CTF submission model
# Rank-standardize characteristics -> ensemble of elastic net (CV-selected
# alpha/lambda) and gradient-boosted trees (CV-selected depth/eta/nrounds) ->
# predicted returns -> long/short weights from a per-month mean-variance QP
# against a statistical factor risk model. Matches the required entrypoint
# in the CTF rules (Rules 11/12).
#
# Why elastic net instead of plain ridge: offline research (run entirely on
# ctff_test==FALSE data, with no look-ahead) found that a plain ridge fit on
# all ~400 characteristics is measurably worse out-of-sample than a fit that
# lets some characteristics
# get shrunk to exactly zero. Rule 2 prohibits hardcoding the specific
# characteristics that research identified as useful ("Manual selection of
# features based on historical performance knowledge is prohibited"), so
# instead of shipping a fixed feature list, this embeds the SELECTION
# MECHANISM itself: elastic net regularization, with alpha/lambda chosen by
# cross-validation, re-fit fresh from whatever training data main() is
# actually given each run. Which characteristics end up with a non-zero
# coefficient is a byproduct of that CV fit each time, not a hardcoded list --
# so this is a live feature-selection step, not a cherry-picked one.
#
# Why also gradient-boosted trees: elastic net only captures a linear
# relationship between standardized characteristics and next-month return.
# Boosted trees capture nonlinearities and interactions the linear model
# structurally cannot, and per the course brief's Guidance section ("Combine
# models. Averaging the predictions of a few different models is a cheap and
# reliable way to improve performance."), averaging its predictions with the
# elastic net's is a standard, non-cherry-picked way to combine the two.
# Hyperparameters (max_depth, eta, nrounds) are chosen by the same
# blocked/time-respecting CV discipline as the elastic net's alpha/lambda --
# never against test-period performance -- so this is a live model-selection
# step, not a hand-picked one either.
#
# data.table and glmnet: data.table is a Rule-8 pre-installed R package, used
# for the per-month standardization loop (otherwise far too slow via base R
# ave() at ~1.4M rows x 400+ characteristics). glmnet and xgboost are NOT
# pre-installed, so this submission ships an renv.lock declaring them (Rule 4/8).
#
# Portfolio construction: predicted returns and a per-month Barra-style
# statistical factor risk model (unchanged prediction/risk-model pipeline)
# feed a mean-variance QP, solved via osqp (sparse, ADMM-based) with Matrix
# for the sparse constraint/quadratic-term assembly -- replaces an earlier
# rank -> slice -> cap -> smooth heuristic, which a real walk-forward
# comparison (never against test-period Sharpe, per Rule 1) showed this QP
# construction beats decisively (Sharpe 1.9574 vs. 1.103 on the full real
# test period). Neither osqp nor Matrix beyond base is pre-installed
# (Matrix itself is base-adjacent but pinned explicitly for reproducibility);
# both are declared in renv.lock (Rule 4/8).
suppressPackageStartupMessages({
  library(data.table)
  library(glmnet)
  library(xgboost)
  library(Matrix)
  library(osqp)
})

ecdf_transform <- function(x) {
  out <- rep(NA_real_, length(x))
  ok <- !is.na(x)
  if (any(ok)) out[ok] <- rank(x[ok], ties.method = "average") / sum(ok)
  out
}

standardize_chars <- function(chars, feature_cols, eom_col = "eom") {
  # setDT (not as.data.table): converts in place, no second copy of the full
  # panel held simultaneously. At full 1990-2023 scale (~1.37M rows x 400+
  # cols, ~4-5GB) as.data.table's copy was observed empirically to spike
  # memory hard enough to threaten a crash before this function even started
  # its per-column loop -- setDT mutates the caller's `chars` object by
  # reference instead (fine here: callers always immediately reassign
  # chars <- standardize_chars(chars, ...) and never reuse the pre-call
  # object afterward).
  dt <- setDT(chars)
  for (i in seq_along(feature_cols)) {
    f <- feature_cols[i]
    is_zero <- !is.na(dt[[f]]) & dt[[f]] == 0
    dt[, (f) := ecdf_transform(get(f)), by = eom_col]
    set(dt, which(is_zero), f, 0)
    set(dt, which(is.na(dt[[f]])), f, 0.5)
    # Empirically, garbage from the per-column which()/grouping temporaries
    # accumulates across this loop faster than R's automatic GC reclaims it
    # at full 1990-2023 scale (400+ columns) -- an explicit periodic gc()
    # keeps peak memory from climbing unboundedly across the whole loop.
    if (i %% 50 == 0) gc()
  }
  as.data.frame(dt)
}

# ---------------------------------------------------------------------------
# Hierarchical clustering feature reduction: groups the ~400 candidate
# characteristics by pairwise correlation and keeps ONE representative per
# cluster, computed fresh each run from TRAINING data only. Adapted from an
# offline training-data analysis (cor() -> hclust() -> cutree() at the same
# CLUSTER_CORR_CUT threshold), simplified to drop that analysis's
# stability-selection step (a bootstrap-resampling procedure
# far too expensive to re-run inside main() itself every time). Between the
# two representative-choice options considered -- (a) the cluster member
# most correlated with the target, or (b) each cluster's first principal
# component -- (a) was chosen: it only SELECTS existing columns, so the
# exact same reduced column list applies to train and test with no separate
# transform to fit and re-apply consistently (a PCA loading vector fit on
# train would need to be carried forward and applied identically to test,
# a second place lookahead could sneak in if done carelessly). Simpler to
# implement correctly and simpler to explain.
#
# Rule 2 compliance: this is a live, algorithmic step recomputed fresh from
# whatever training data main() is given each run -- which characteristics
# end up as cluster representatives is a byproduct of that run's own
# correlation structure and its own correlation with the training target,
# not a hardcoded list. Same principle that makes the elastic net's own
# feature selection compliant.
CLUSTER_CORR_CUT <- 0.7  # features with |corr| >= this tend to land in the
                          # same cluster -- same threshold already used
                          # in the offline analysis, not tuned here
MIN_FEATURES_FOR_CLUSTERING <- 20L  # below this there's little to reduce
                                     # and a correlation matrix isn't
                                     # meaningful -- skip and use every
                                     # candidate feature unchanged (same
                                     # thin-data fallback philosophy as
                                     # elsewhere in this file)

# X_train: standardized characteristics, TRAINING rows only. y_train: the
# training target. feature_cols: the candidate column names, in the same
# order as X_train's columns. Returns the (possibly reduced) subset of
# feature_cols to actually use.
cluster_reduce_features <- function(X_train, y_train, feature_cols) {
  if (length(feature_cols) < MIN_FEATURES_FOR_CLUSTERING) return(feature_cols)

  cor_mat <- tryCatch(cor(X_train, use = "pairwise.complete.obs"), error = function(e) NULL)
  if (is.null(cor_mat) || anyNA(cor_mat)) return(feature_cols)

  hc <- tryCatch(hclust(as.dist(1 - abs(cor_mat)), method = "average"), error = function(e) NULL)
  if (is.null(hc)) return(feature_cols)

  cluster_id <- cutree(hc, h = 1 - CLUSTER_CORR_CUT)

  # each cluster's representative: the member most correlated (in absolute
  # value) with the TRAINING target -- computed once per run, from training
  # rows only, never from test-period rows
  target_cor <- abs(cor(X_train, y_train, use = "pairwise.complete.obs")[, 1])
  target_cor[!is.finite(target_cor)] <- 0

  representative_by_cluster <- tapply(feature_cols, cluster_id, function(members) {
    members[which.max(target_cor[members])]
  })
  unname(unlist(representative_by_cluster))
}

fit_ridge <- function(X, y, lambda = 10) {
  X <- cbind(1, X)
  p <- ncol(X)
  penalty <- diag(p) * lambda
  penalty[1, 1] <- 0 # do not shrink the intercept
  coeffs <- solve(t(X) %*% X + penalty, t(X) %*% y)
  as.numeric(coeffs)
}

# Elastic net with a blocked (time-respecting, never random) inner CV to pick
# alpha and lambda -- same recipe validated in the offline analysis,
# simplified to a single fit (no fold/bootstrap stability
# averaging) so it stays cheap enough to run inside main() itself.
#
# Falls back to plain ridge whenever there isn't enough data/date variety for
# a meaningful blocked CV split -- this matters because the CTF validation
# run uses a deliberately tiny (~123-observation) dataset (Rule 14), and a
# handful of dates can't support a stable inner CV.
ALPHA_GRID <- c(0.1, 0.3, 0.5, 0.7, 1.0)
N_INNER_CV_FOLDS <- 3
MIN_DATES_FOR_CV <- N_INNER_CV_FOLDS * 4
MIN_ROWS_FOR_CV <- 50

# Shared blocked (time-respecting, never random) fold assignment: dates are
# sorted and cut into N_INNER_CV_FOLDS contiguous chronological blocks, so a
# fold's "test" rows are always later in time than an earlier fold's -- the
# same scheme fit_elastic_net() always used, factored out so fit_xgboost()
# below uses the identical temporal-blocking discipline rather than a second,
# possibly-diverging implementation.
make_blocked_foldid <- function(dates, n_folds) {
  unique_dates <- sort(unique(dates))
  block_id_by_date <- setNames(
    ceiling(seq_along(unique_dates) / length(unique_dates) * n_folds),
    as.character(unique_dates)
  )
  unname(block_id_by_date[as.character(dates)])
}

fit_elastic_net <- function(X, y, dates, lambda_fallback = 10) {
  # Returns a tagged list, not a bare coefficient vector: fit_stacked_ensemble()
  # below needs to know WHICH hyperparameters (if any) were CV-selected, so it
  # can refit this same specification on each stacking fold rather than
  # re-running its own grid search. type is "zero" (no training rows),
  # "ridge_fallback" (too little data/date variety for CV), or "glmnet"
  # (CV-selected alpha/lambda).
  #
  # pathological case: no training rows at all (e.g. a truncated re-run whose
  # window falls entirely within the test period). No model can be fit, so
  # return an all-zero coefficient vector -- every prediction becomes 0, and
  # construct_qp_portfolio()'s own degenerate-month fallback (zero-weight)
  # takes over downstream instead of crashing on a singular ridge system.
  if (nrow(X) == 0) return(list(type = "zero", coeffs = rep(0, ncol(X) + 1)))

  unique_dates <- sort(unique(dates))
  if (length(unique_dates) < MIN_DATES_FOR_CV || nrow(X) < MIN_ROWS_FOR_CV) {
    return(list(type = "ridge_fallback", coeffs = fit_ridge(X, y, lambda = lambda_fallback),
                lambda_fallback = lambda_fallback))
  }

  foldid <- make_blocked_foldid(dates, N_INNER_CV_FOLDS)

  best <- list(alpha = NA, lambda = NA, cvm = Inf)
  for (a in ALPHA_GRID) {
    fit <- tryCatch(
      cv.glmnet(X, y, alpha = a, foldid = foldid, type.measure = "mse", standardize = FALSE),
      error = function(e) NULL
    )
    if (!is.null(fit)) {
      cvm_min <- min(fit$cvm)
      if (is.finite(cvm_min) && cvm_min < best$cvm) {
        best <- list(alpha = a, lambda = fit$lambda.min, cvm = cvm_min)
      }
    }
  }
  if (is.na(best$alpha)) {
    return(list(type = "ridge_fallback", coeffs = fit_ridge(X, y, lambda = lambda_fallback),
                lambda_fallback = lambda_fallback))
  }

  final_fit <- glmnet(X, y, alpha = best$alpha, lambda = best$lambda, standardize = FALSE)
  coeffs <- as.numeric(coef(final_fit))  # glmnet's coef() already includes the intercept as row 1
  cat(sprintf("[main] elastic net: alpha=%.2f, lambda=%.5g, non-zero characteristics=%d of %d\n",
              best$alpha, best$lambda, sum(coeffs[-1] != 0), length(coeffs) - 1))
  list(type = "glmnet", coeffs = coeffs, alpha = best$alpha, lambda = best$lambda)
}

# ---------------------------------------------------------------------------
# Gradient-boosted trees, second leg of the prediction ensemble (see file
# header). Same target, same standardized feature matrix, same blocked
# (time-respecting) CV discipline as fit_elastic_net() -- max_depth/eta are a
# small, fixed grid (not tuned on test-period performance, same reasoning as
# ALPHA_GRID above); nrounds is chosen per grid point by CV early stopping
# rather than fixed, so the "how many trees" decision is itself CV-driven.
#
# subsample/colsample_bytree < 1 make each tree's construction depend on
# xgboost's internal RNG, so every call below pins params$seed explicitly
# (alongside this file's set.seed(42) in fit_predictions()) -- otherwise two
# runs on identical data could produce different trees and violate Rule 18's
# determinism requirement.
XGB_MAX_DEPTH_GRID <- c(2L, 3L, 4L)
XGB_ETA_GRID <- c(0.05, 0.1)
XGB_NROUNDS_MAX <- 300L
XGB_EARLY_STOPPING_ROUNDS <- 20L
XGB_SEED <- 42L
XGB_FALLBACK_MAX_DEPTH <- 2L
XGB_FALLBACK_ETA <- 0.1
XGB_FALLBACK_NROUNDS <- 50L

xgb_params <- function(max_depth, eta) {
  list(max_depth = max_depth, eta = eta, objective = "reg:squarederror",
       eval_metric = "rmse", subsample = 0.8, colsample_bytree = 0.8, seed = XGB_SEED)
}

# Returns a small tagged list rather than a bare xgb.Booster, so the
# no-training-rows case can be represented without a dummy/degenerate model:
# type == "zero" means "predict 0 for everything", handled explicitly by
# predict_xgb() -- mirrors fit_elastic_net()'s all-zero-coefficient fallback
# for the same pathological (nrow(X) == 0) case.
fit_xgboost <- function(X, y, dates) {
  if (nrow(X) == 0) return(list(type = "zero"))

  dtrain <- xgb.DMatrix(data = X, label = y)
  unique_dates <- sort(unique(dates))

  fit_single <- function(max_depth, eta, nrounds) {
    xgb.train(params = xgb_params(max_depth, eta), data = dtrain, nrounds = nrounds, verbose = 0)
  }

  # same gate as fit_elastic_net()'s ridge fallback: too little data/date
  # variety (e.g. Rule 14's ~123-observation validation dataset can leave a
  # handful of training months) can't support a stable blocked CV split, so
  # fit one small, fixed-hyperparameter booster directly instead.
  if (length(unique_dates) < MIN_DATES_FOR_CV || nrow(X) < MIN_ROWS_FOR_CV) {
    return(list(type = "booster", booster = fit_single(XGB_FALLBACK_MAX_DEPTH, XGB_FALLBACK_ETA, XGB_FALLBACK_NROUNDS),
                max_depth = XGB_FALLBACK_MAX_DEPTH, eta = XGB_FALLBACK_ETA, nrounds = XGB_FALLBACK_NROUNDS))
  }

  foldid <- make_blocked_foldid(dates, N_INNER_CV_FOLDS)
  folds <- lapply(sort(unique(foldid)), function(k) which(foldid == k))

  best <- list(max_depth = NA, eta = NA, nrounds = NA, rmse = Inf)
  for (d in XGB_MAX_DEPTH_GRID) {
    for (e in XGB_ETA_GRID) {
      cvfit <- tryCatch(
        xgb.cv(params = xgb_params(d, e), data = dtrain, nrounds = XGB_NROUNDS_MAX, folds = folds,
               early_stopping_rounds = XGB_EARLY_STOPPING_ROUNDS, verbose = 0),
        error = function(e) NULL
      )
      if (!is.null(cvfit)) {
        # this xgboost version reports early-stopping results under
        # $early_stop, not top-level $best_iteration/$best_score
        best_iter <- cvfit$early_stop$best_iteration
        rmse <- unname(cvfit$early_stop$best_score)
        if (!is.null(rmse) && is.finite(rmse) && rmse < best$rmse) {
          best <- list(max_depth = d, eta = e, nrounds = best_iter, rmse = rmse)
        }
      }
    }
  }
  if (is.na(best$max_depth)) {
    return(list(type = "booster", booster = fit_single(XGB_FALLBACK_MAX_DEPTH, XGB_FALLBACK_ETA, XGB_FALLBACK_NROUNDS),
                max_depth = XGB_FALLBACK_MAX_DEPTH, eta = XGB_FALLBACK_ETA, nrounds = XGB_FALLBACK_NROUNDS))
  }

  cat(sprintf("[main] xgboost: max_depth=%d, eta=%.3f, nrounds=%d, cv_rmse=%.5g\n",
              best$max_depth, best$eta, best$nrounds, best$rmse))
  list(type = "booster", booster = fit_single(best$max_depth, best$eta, best$nrounds),
       max_depth = best$max_depth, eta = best$eta, nrounds = best$nrounds)
}

predict_xgb <- function(model, X) {
  if (model$type == "zero") return(rep(0, nrow(X)))
  as.numeric(predict(model$booster, xgb.DMatrix(data = X)))
}

# ---------------------------------------------------------------------------
# Stacked ensemble: a ridge regression learns how to combine the elastic
# net's and xgboost's predictions, replacing the earlier flat 50/50 average.
# Fit on OUT-OF-FOLD base-model predictions, not the base models' predictions
# on the training data they were themselves fit on -- the latter would be
# optimistic/leaky (xgboost in particular can nearly interpolate its own
# training data), which would teach the meta-learner to trust whichever base
# model overfits hardest, not whichever one actually generalizes. The same
# blocked (time-respecting) fold split as everywhere else (make_blocked_foldid())
# produces the out-of-fold predictions: for each fold, refit both base models
# on the OTHER folds and predict on the held-out fold.
#
# Reuses the hyperparameters fit_elastic_net()/fit_xgboost() already
# CV-selected on the FULL training set (en_fit$alpha/lambda,
# xgb_fit$max_depth/eta/nrounds) rather than re-running either model's own
# grid search inside every stacking fold -- this fold split exists only to
# get honest out-of-fold predictions for the stacker, not to re-tune the
# base models' hyperparameters. This reuse is only valid when both base
# models actually went through their CV path rather than a small-data
# fallback; that's guaranteed here because fit_stacked_ensemble()'s own gate
# below (MIN_DATES_FOR_CV / MIN_ROWS_FOR_CV) is identical to the gate inside
# fit_elastic_net()/fit_xgboost(), so whenever stacking proceeds past it, the
# base models fit on this same (X, dates) are guaranteed to have taken their
# CV branch too.
#
# Falls back to a flat 50/50 average -- exactly today's behavior -- when
# there isn't enough data/date variety for a meaningful stacked fit, same
# philosophy as fit_elastic_net()'s ridge fallback and fit_xgboost()'s
# fixed-hyperparameter fallback.
fit_stacked_ensemble <- function(X, y, dates, en_fit, xgb_fit) {
  n <- nrow(X)
  unique_dates <- sort(unique(dates))
  if (n == 0 || length(unique_dates) < MIN_DATES_FOR_CV || n < MIN_ROWS_FOR_CV ||
      en_fit$type != "glmnet" || xgb_fit$type != "booster") {
    return(list(type = "flat_average"))
  }

  foldid <- make_blocked_foldid(dates, N_INNER_CV_FOLDS)
  oof_en <- rep(NA_real_, n)
  oof_xgb <- rep(NA_real_, n)

  # Explicit rm()/gc() per fold: each iteration allocates a large
  # train-subset copy of X plus an xgb.DMatrix (which holds its own
  # internal copy of that subset) and two freshly-fit model objects. R's
  # garbage collector doesn't necessarily run between iterations of a tight
  # loop, so without this, up to N_INNER_CV_FOLDS generations of these
  # large objects can be simultaneously reachable (the loop variable
  # reassignment alone doesn't free the previous iteration's objects until
  # GC actually runs) -- found the hard way when testing an unclustered
  # (402-column, vs. the shipped 204-column) feature set OOM-killed this
  # step on a real month at 65GB resident, on a droplet with 62GB RAM.
  # Freeing explicitly keeps peak memory to roughly one fold's worth
  # regardless of column count.
  for (k in sort(unique(foldid))) {
    train_idx <- which(foldid != k)
    heldout_idx <- which(foldid == k)
    if (length(train_idx) == 0 || length(heldout_idx) == 0) next

    en_fit_k <- glmnet(X[train_idx, , drop = FALSE], y[train_idx],
                        alpha = en_fit$alpha, lambda = en_fit$lambda, standardize = FALSE)
    en_coef_k <- as.numeric(coef(en_fit_k))
    oof_en[heldout_idx] <- as.numeric(cbind(1, X[heldout_idx, , drop = FALSE]) %*% en_coef_k)
    rm(en_fit_k, en_coef_k)

    dtrain_k <- xgb.DMatrix(data = X[train_idx, , drop = FALSE], label = y[train_idx])
    xgb_fit_k <- xgb.train(params = xgb_params(xgb_fit$max_depth, xgb_fit$eta),
                            data = dtrain_k, nrounds = xgb_fit$nrounds, verbose = 0)
    oof_xgb[heldout_idx] <- as.numeric(predict(xgb_fit_k, xgb.DMatrix(data = X[heldout_idx, , drop = FALSE])))
    rm(dtrain_k, xgb_fit_k)
    gc()
  }

  ok <- is.finite(oof_en) & is.finite(oof_xgb) & is.finite(y)
  if (sum(ok) < MIN_ROWS_FOR_CV) return(list(type = "flat_average"))

  # pure ridge (alpha = 0): a lightweight 2-input meta-model, CV-selected
  # lambda via the same blocked folds used to generate the OOF predictions
  Z <- cbind(oof_en[ok], oof_xgb[ok])
  meta_fit <- tryCatch(
    cv.glmnet(Z, y[ok], alpha = 0, foldid = foldid[ok], type.measure = "mse", standardize = FALSE),
    error = function(e) NULL
  )
  if (is.null(meta_fit)) return(list(type = "flat_average"))

  meta_coef <- as.numeric(coef(meta_fit, s = "lambda.min"))
  cat(sprintf("[main] stacked ensemble: intercept=%.5g, w_elastic_net=%.4f, w_xgboost=%.4f\n",
              meta_coef[1], meta_coef[2], meta_coef[3]))
  list(type = "stacked", intercept = meta_coef[1], w_en = meta_coef[2], w_xgb = meta_coef[3])
}

predict_stacked <- function(stack_model, pred_en, pred_xgb) {
  if (stack_model$type == "flat_average") return((pred_en + pred_xgb) / 2)
  stack_model$intercept + stack_model$w_en * pred_en + stack_model$w_xgb * pred_xgb
}

# ctff_test is delivered as a string ("0"/"1") in the real WRDS-sourced
# parquet files, not a logical TRUE/FALSE -- handle both representations so
# the same code works regardless of how a given data vintage encodes it.
is_test_row <- function(x) {
  if (is.logical(x)) return(x)
  tolower(trimws(as.character(x))) %in% c("1", "true", "t")
}

TARGET_GROSS_EXPOSURE <- 2  # 100% long + 100% short, dollar-neutral -- the
                             # fixed exposure convention the QP's raw output
                             # gets rescaled to each month (see
                             # construct_qp_portfolio() below), so this stays
                             # comparable to every earlier construction
                             # approach tried in this project.

# ---------------------------------------------------------------------------
# Statistical factor covariance risk model (Barra-style), estimated fresh
# each test month from trailing daily_ret history and fed directly into the
# mean-variance QP (construct_qp_portfolio() below) as that month's Sigma.
# daily_ret is otherwise unused in this file.
#
# Design: Sigma = B %*% diag(lambda) %*% t(B) + diag(D) (loadings B, factor
# variances lambda from a statistical PCA-via-SVD fit on that month's
# trailing daily returns; D = per-name idiosyncratic variance). T (trading
# days, ~252) is the smaller dimension vs. N (~2000-3000 names), so base R
# svd() on the T x N matrix is cheap (LAPACK cost is driven by the smaller
# dimension), unlike eigen-decomposing an N x N correlation matrix. Sigma
# itself IS densified (N x N) here -- unlike an earlier version of this risk
# model that used the Woodbury identity to avoid ever forming it -- because
# the QP solver needs the actual matrix, not just its action on a vector;
# building it fresh per month and discarding it immediately keeps peak
# memory to one month's Sigma at a time (~50-75MB at realistic N), never
# accumulated across months.
#
# NO-LOOKAHEAD: the trailing window for month t is date <= t only (never
# future dates), pulled by date filtering -- the daily returns table has no
# ctff_test flag to lean on. FALLBACK: whenever a month's daily history is
# too thin (too few trading days, too few adequately-covered names, or any
# numerical failure below), this returns NULL and construct_qp_portfolio()
# falls back to a plain rank-demean for that month -- this is what keeps the
# CTF validation run (Rule 14's tiny ~123-observation dataset) safe, since
# its daily_ret slice won't have enough history to pass these gates.
RISK_LOOKBACK_CALENDAR_DAYS <- 370L  # ~252 trading days + weekend/holiday buffer
MIN_TRADING_DAYS <- 60L              # below this, a factor fit isn't credible
MIN_NAMES_FOR_RISK_MODEL <- 30L      # >= 3x N_FACTORS_MAX
MIN_STOCK_COVERAGE_FRAC <- 0.8       # per-name inclusion threshold, relative to that month's actual trading days
N_FACTORS_MAX <- 10L                 # standard statistical-factor-model range; fixed, not tuned on test-period performance (Rule 2)
IDIO_VAR_REL_FLOOR <- 1e-6           # floor D at this fraction of mean column variance -- see SPEC_VAR_FLOOR_FRAC below for why this alone isn't enough for the QP

prepare_daily_ret <- function(daily_ret) {
  dt <- as.data.table(daily_ret)
  setkey(dt, id, date)
  dt
}

# One non-equi join per test month against the full daily_ret table (rather
# than batching every test month's window into a single join up front). The
# batched version materializes every month's window rows at once, which at
# full 1990-2023 scale pushed peak R heap past what this machine has
# available; this per-month version pays a larger constant per-join cost
# (empirically ~15s/call, ~400+ calls across the full test period) but keeps
# only one month's window in memory at a time.
build_month_return_matrix <- function(daily_ret_keyed, ids_month, eom_d) {
  lo <- eom_d - RISK_LOOKBACK_CALENDAR_DAYS
  hi <- eom_d - 1L
  target_dt <- data.table(id = ids_month, lo = lo, hi = hi)
  # x.date, as before: without it, the non-equi join silently returns i's
  # bound value instead of x's actual matched date.
  window_month <- daily_ret_keyed[target_dt, .(id, date = x.date, ret_exc),
                                   on = .(id, date >= lo, date <= hi), nomatch = NULL]
  if (nrow(window_month) == 0) return(NULL)

  n_trading_days <- length(unique(window_month$date))
  if (n_trading_days < MIN_TRADING_DAYS) return(NULL)

  coverage <- window_month[, .N, by = id]
  covered_ids <- coverage[N >= MIN_STOCK_COVERAGE_FRAC * n_trading_days, id]
  if (length(covered_ids) < MIN_NAMES_FOR_RISK_MODEL) return(NULL)

  wide <- dcast(window_month[id %in% covered_ids], date ~ id, value.var = "ret_exc", fun.aggregate = mean)
  ret_mat <- as.matrix(wide[, -1, with = FALSE])
  ret_mat[is.na(ret_mat)] <- 0                       # residual gaps within an otherwise-covered name
  ret_mat <- scale(ret_mat, center = TRUE, scale = FALSE)  # demean only -- keep raw return units for variance

  list(mat = ret_mat, covered_ids = as.numeric(colnames(wide)[-1]))
}

fit_statistical_factor_model <- function(ret_mat) {
  n_days <- nrow(ret_mat)
  n_names <- ncol(ret_mat)
  k <- max(1L, min(N_FACTORS_MAX, floor(n_days / 10), n_names - 1L))

  sv <- tryCatch(svd(ret_mat, nu = 0, nv = k), error = function(e) NULL)
  if (is.null(sv)) return(NULL)

  B <- sv$v[, seq_len(k), drop = FALSE]
  lambda <- sv$d[seq_len(k)]^2 / (n_days - 1)
  var_i <- colSums(ret_mat^2) / (n_days - 1)
  explained_i <- as.numeric((B^2) %*% lambda)
  D <- pmax(var_i - explained_i, IDIO_VAR_REL_FLOOR * mean(var_i))

  ok <- all(is.finite(B)) && all(is.finite(lambda)) && all(lambda > 0) &&
    all(is.finite(D)) && all(D > 0)
  if (!ok) return(NULL)

  list(B = B, lambda = lambda, D = D)
}

# Specific-risk (idiosyncratic variance) floor, as a fraction of that
# month's OWN median fitted D -- deliberately relative, not an absolute
# constant, so it scales with the market's overall volatility level instead
# of binding differently in calm vs. crisis regimes.
#
# Why this exists: an early real-data test of this QP construction produced
# a degenerate, multi-million-dollar-per-dollar position, traced to a single
# stock's fitted specific risk being far below that month's typical value --
# Sigma's smallest eigenvalue is bounded below by min(D) (a PSD matrix added
# to diag(D) can only raise eigenvalues), so one stock's tiny D directly
# makes Sigma's smallest eigenvalue tiny, and the optimizer's risk penalty
# along that direction becomes negligible. Checked across 9 real months
# spanning calm and crisis periods before treating this as genuine rather
# than a data artifact: the offending stock had full trading-day coverage,
# no exact-zero returns, no abnormal repeated-value runs, and was a
# DIFFERENT stock each time -- a genuinely well-explained residual, not thin
# data or stale prices. IDIO_VAR_REL_FLOOR above already floors D, but at
# 3-4 orders of magnitude below where these real cases land -- far too small
# to ever bind. 5% of the month's median D binds on the bottom half of the
# observed worst-case-ratio range (0.37%-13.5% across the 9-month check)
# while leaving the bulk of legitimately-varying D values untouched. (This
# floor alone was later found NOT sufficient by itself to fix the degenerate
# case above -- see QP_LAMBDA/QP_KAPPA below for the actual root cause and
# fix -- but is independently justified and kept regardless.)
SPEC_VAR_FLOOR_FRAC <- 0.05

# Daily-to-monthly variance scaling. fit_statistical_factor_model() fits on
# DAILY returns, so its raw lambda/D are daily variances -- but this QP is a
# monthly rebalancing decision against a monthly expected-return prediction
# (test$pred). A mean-variance objective needs alpha and Sigma in the same
# holding-period units, so the daily factor variances are scaled up to
# monthly here by the standard trading-days-per-month convention (252/12),
# under the same no-autocorrelation approximation already relied on
# elsewhere in this project. A fixed unit conversion, not a tuned parameter.
MONTHLY_TRADING_DAYS <- 21

# Builds one month's QP inputs from the SAME per-month trailing-window
# factor fit already used above -- returns the small FACTORS (B_full,
# lambda, D_full), not a densified Sigma: for n ~ 2,500 names that's ~220KB
# vs. ~50MB for the dense n x n Sigma, so this is cheap to hold if ever
# needed across months, though construct_qp_portfolio() below builds and
# discards Sigma within the same month's iteration either way.
build_month_qp_inputs <- function(alpha, ids_month, daily_ret_keyed, eom_d) {
  tryCatch({
    if (length(ids_month) < 2) return(NULL)

    rm_data <- build_month_return_matrix(daily_ret_keyed, ids_month, eom_d)
    if (is.null(rm_data)) return(NULL)

    fm <- fit_statistical_factor_model(rm_data$mat)
    if (is.null(fm)) return(NULL)

    k <- length(fm$lambda)
    match_pos <- match(ids_month, rm_data$covered_ids)
    covered <- !is.na(match_pos)

    # names that fail the coverage threshold get idiosyncratic-only treatment
    # (zero factor exposure, average idiosyncratic variance) rather than
    # being dropped -- every name in the month must still get a weight (Rule 12).
    B_full <- matrix(0, nrow = length(ids_month), ncol = k)
    B_full[covered, ] <- fm$B[match_pos[covered], , drop = FALSE]
    D_full <- rep(mean(fm$D), length(ids_month))
    D_full[covered] <- fm$D[match_pos[covered]]

    spec_floor <- SPEC_VAR_FLOOR_FRAC * median(fm$D)
    D_full <- pmax(D_full, spec_floor)

    # daily -> monthly variance units (see MONTHLY_TRADING_DAYS comment
    # above). Applied after the floor -- floor is a fraction of that
    # month's own (still-daily) median D, so scaling before or after is
    # equivalent; doing it after keeps the floor's own logic untouched.
    list(alpha = alpha, B_full = B_full,
         lambda = fm$lambda * MONTHLY_TRADING_DAYS,
         D_full = D_full * MONTHLY_TRADING_DAYS,
         ids = ids_month)
  }, error = function(e) NULL)
}

# Builds the dense n x n Sigma for ONE month from its small cached factors.
# Called once per month inside construct_qp_portfolio()'s loop, used for
# that month's QP solve, then allowed to go out of scope -- never held
# across months.
build_sigma_from_factors <- function(qp_inputs) {
  k <- length(qp_inputs$lambda)
  Sigma <- qp_inputs$B_full %*% diag(qp_inputs$lambda, k) %*% t(qp_inputs$B_full)
  diag(Sigma) <- diag(Sigma) + qp_inputs$D_full
  (Sigma + t(Sigma)) / 2  # defensive symmetrize, float noise only
}

# ---------------------------------------------------------------------------
# Mean-variance QP portfolio construction.
#
#   maximize   w'alpha - lambda * w'Sigma*w - kappa * ||w - w_prev||_1
#   subject to sum(w) = 0                      (dollar-neutral)
#
# Solved via osqp by introducing slack variables u, v >= 0 such that
# w - w_prev = u - v, so ||w - w_prev||_1 = sum(u + v). Variable vector
# x = [w (n), u (n), v (n)], length 3n. osqp minimizes (1/2)x'Px + q'x
# subject to l <= Ax <= u (two-sided row bounds: an equality is l_i = u_i,
# a one-sided constraint sets the other bound to +-Inf).
#
# QP_LAMBDA (risk-aversion) and QP_KAPPA (turnover penalty) were chosen via
# a lightweight walk-forward check (a separate development driver, not part
# of this file): 5 expanding-window folds INSIDE the training period only
# (never the real test months this file is scored on -- Rule 1 applies to
# hyperparameter tuning, not just prediction), each refitting the full
# ensemble from scratch on data strictly before its own held-out block, 15
# combinations of the two dials tested per fold. Result: QP_KAPPA = 0 won
# decisively (median Sharpe 4.32 across the 5 folds, vs. 0.36 for the best
# nonzero-kappa combination) and every QP_LAMBDA value tied exactly at
# kappa = 0 (expected: with no turnover term, a month's answer cannot
# depend on w_prev at all, so risk-aversion only rescales the raw answer's
# SIZE before the fixed-gross-exposure rescaling below erases the
# difference). Nonzero kappa's walk-forward results were erratic across
# folds -- strongly positive in two folds, strongly negative in a third,
# for reasons traced to kappa comparing each month's raw pre-rescaling
# solution (which can naturally sit at very different scales month to
# month) against a previous month's answer that's already been rescaled to
# a fixed convention -- a real design wrinkle worth revisiting with a
# cleaner formulation later, not something to route around by hand-picking
# a "safer-looking" nonzero kappa now. QP_LAMBDA's specific value is
# therefore inconsequential given QP_KAPPA = 0 -- kept at a middle-of-grid
# value for documentation clarity, not because it was found to matter.
QP_LAMBDA <- 2
QP_KAPPA <- 0

# ridge_eps: small ridge on the u/v slack block of the QP's quadratic term
# so it stays positive definite even when Sigma is only PSD -- several
# orders of magnitude below typical Sigma eigenvalues, so it doesn't
# meaningfully change the objective.
solve_portfolio_qp <- function(alpha, Sigma, w_prev, lambda, kappa, ridge_eps = 1e-8) {
  n <- length(alpha)
  nx <- 3 * n
  idx_w <- 1:n; idx_u <- (n + 1):(2 * n); idx_v <- (2 * n + 1):(3 * n)

  Sigma_sym <- (Sigma + t(Sigma)) / 2
  P <- Matrix::bdiag(
    2 * lambda * Matrix::Matrix(Sigma_sym, sparse = TRUE),
    Matrix::Diagonal(n, x = ridge_eps),
    Matrix::Diagonal(n, x = ridge_eps)
  )

  q <- numeric(nx)
  q[idx_w] <- -alpha
  q[idx_u] <- kappa
  q[idx_v] <- kappa

  n_rows <- n + 1 + n + n
  tri_i <- integer(0); tri_j <- integer(0); tri_x <- numeric(0)
  l <- numeric(n_rows); u <- numeric(n_rows)

  # (1) turnover link equalities: w_i - u_i + v_i = w_prev_i, rows 1..n
  tri_i <- c(tri_i, 1:n, 1:n, 1:n)
  tri_j <- c(tri_j, idx_w, idx_u, idx_v)
  tri_x <- c(tri_x, rep(1, n), rep(-1, n), rep(1, n))
  l[1:n] <- w_prev; u[1:n] <- w_prev

  # (2) dollar-neutral equality: sum(w) = 0, row n+1
  row_dn <- n + 1
  tri_i <- c(tri_i, rep(row_dn, n)); tri_j <- c(tri_j, idx_w); tri_x <- c(tri_x, rep(1, n))
  l[row_dn] <- 0; u[row_dn] <- 0

  # (3) u_i >= 0: rows n+2 .. 2n+1
  rows_u <- (n + 2):(2 * n + 1)
  tri_i <- c(tri_i, rows_u); tri_j <- c(tri_j, idx_u); tri_x <- c(tri_x, rep(1, n))
  l[rows_u] <- 0; u[rows_u] <- Inf

  # (4) v_i >= 0: rows 2n+2 .. 3n+1
  rows_v <- (2 * n + 2):(3 * n + 1)
  tri_i <- c(tri_i, rows_v); tri_j <- c(tri_j, idx_v); tri_x <- c(tri_x, rep(1, n))
  l[rows_v] <- 0; u[rows_v] <- Inf

  A <- Matrix::sparseMatrix(i = tri_i, j = tri_j, x = tri_x, dims = c(n_rows, nx))

  pars <- osqp::osqpSettings(verbose = FALSE, eps_abs = 1e-8, eps_rel = 1e-8,
                              max_iter = 20000L, polishing = TRUE)
  sol <- osqp::solve_osqp(P = P, q = q, A = A, l = l, u = u, pars = pars)
  if (!(sol$info$status %in% c("solved", "solved inaccurate"))) {
    stop(sprintf("osqp did not solve: status = %s", sol$info$status))
  }
  sol$x[idx_w]
}


# Rule 17 explicitly invites stdout progress logging ("Consider logging
# progress updates, timing information, and intermediate results to aid
# troubleshooting") -- this is the *only* logging mechanism used in main(),
# never direct file writes (Rule 12 forbids those). Timestamped so an
# external log file (produced by redirecting this process's stdout, not by
# main() itself) can be used to tell how far a run got if it's interrupted.
log_stage <- function(msg) {
  cat(sprintf("[main %s] %s\n", format(Sys.time(), "%Y-%m-%d %H:%M:%S"), msg))
  flush(stdout())
}

# Timestamped "starting <step>" marker, so an unattended run's log shows which
# step is in progress, not only which ones have finished. message() writes to
# stderr, which the CTF also captures (Rule 17). Pure logging: no effect on
# any computed value.
log_step <- function(msg) {
  message(sprintf("[step %s] starting: %s", format(Sys.time(), "%Y-%m-%d %H:%M:%S"), msg))
}

# Everything through the raw predicted return: standardization, the
# elastic net + xgboost fits, and the stacked-ensemble prediction. Factored
# out from construct_qp_portfolio() so the two stay independently testable
# and can never numerically diverge.
#
# Deliberately does NOT take daily_ret -- it's never used in this function
# (only construct_qp_portfolio(), via prepare_daily_ret(), needs it). At full
# 1990-2023 scale, an empirical run showed standardize_chars() alone
# pushing free system memory to the edge of a crash; daily_ret is a ~2.4GB
# table that callers should defer loading until after this function
# returns, rather than holding it resident during the riskiest step for no
# reason.
fit_predictions <- function(chars, features) {
  set.seed(42)

  eom_col <- "eom"
  target_col <- "ret_exc_lead1m"
  feature_cols <- as.character(features$features)
  feature_cols <- intersect(feature_cols, names(chars))

  log_stage(sprintf("main() started: %d chars rows, %d candidate feature cols", nrow(chars), length(feature_cols)))

  log_step("standardizing characteristics (monthly ranks)")
  chars <- standardize_chars(chars, feature_cols, eom_col)
  log_stage("standardization done")

  test_flag <- is_test_row(chars$ctff_test)
  train <- chars[!test_flag & !is.na(chars[[target_col]]), ]
  test  <- chars[test_flag, ]

  X_train_full <- as.matrix(train[, feature_cols])
  y_train <- train[[target_col]]
  dates_train <- train[[eom_col]]

  # Clustering feature reduction (see cluster_reduce_features() header):
  # computed from X_train_full / y_train only, i.e. training rows -- the
  # reduced column list is then used identically for both X_train and
  # X_test below, so both models see the same, smaller feature set.
  #
  # Re-checked with an on/off toggle (separate development driver) after
  # the QP rebuild, since this had never been cleanly A/B tested on its
  # own before -- a genuine walk-forward comparison came back mixed (2 of
  # 5 real historical slices favored OFF by a wide margin, including the
  # 1987-crash slice, but the project's own median-based decision
  # statistic favored keeping it ON, if only by a hair) and didn't clear
  # the bar to change anything -- kept ON, matching every prior real-data
  # result. See the methodology document.
  log_step("clustering characteristics (training data)")
  feature_cols <- cluster_reduce_features(X_train_full, y_train, feature_cols)
  log_stage(sprintf("clustering: reduced to %d representative characteristics", length(feature_cols)))
  X_train <- X_train_full[, feature_cols, drop = FALSE]
  rm(X_train_full)

  log_step("elastic net (blocked CV)")
  en_fit <- fit_elastic_net(X_train, y_train, dates_train)
  log_stage("elastic net fit done")

  log_step("gradient-boosted trees (blocked CV)")
  xgb_fit <- fit_xgboost(X_train, y_train, dates_train)
  log_stage("xgboost fit done")

  # Stacked ensemble (see fit_stacked_ensemble() header for the leakage
  # rationale): learns how to combine the two base models' predictions via a
  # ridge meta-model fit on out-of-fold predictions, rather than a
  # hand-fixed average. Falls back to a flat average when there's too
  # little data/date variety for a meaningful stacked fit.
  log_step("stacked ensemble (out-of-fold forecasts)")
  stack_model <- fit_stacked_ensemble(X_train, y_train, dates_train, en_fit, xgb_fit)
  log_stage("stacked ensemble fit done")

  rm(X_train, y_train, dates_train)  # all fits are done; ~400-col training matrix no longer needed
  gc()

  log_step("predicting test-period returns")
  X_test <- as.matrix(test[, feature_cols])
  pred_en <- as.numeric(cbind(1, X_test) %*% en_fit$coeffs)
  pred_xgb <- predict_xgb(xgb_fit, X_test)
  # Raw predicted return -- the course brief's Guidance section states the intended
  # convention explicitly ("go long stocks with high predicted returns and
  # short stocks with low ones"), and construct_qp_portfolio() below feeds
  # this directly to the optimizer as alpha (in ret_exc_lead1m's own units,
  # not an arbitrary-scale rank transform -- see QP_LAMBDA/QP_KAPPA comment
  # for why that distinction matters).
  test$pred <- predict_stacked(stack_model, pred_en, pred_xgb)

  list(test = test, eom_col = eom_col)
}

# Per-month QP construction loop, over `months` (defaults to every test
# month). With QP_KAPPA = 0 (see above) each month's solve is fully
# independent of every other month -- no turnover state to thread, unlike
# the rank/slice/cap/smooth chain this replaced -- so months could in
# principle run in any order or be split across separate calls with no
# special handling; kept as an explicit sorted loop anyway for
# determinism (Rule 18) and straightforward progress logging.
#
# Falls back to a plain rank-demeaned weight (same zero-sum,
# TARGET_GROSS_EXPOSURE-scaled convention as every earlier construction
# approach in this project) for any month where the factor-model fit fails
# (thin history, degenerate SVD) or the QP solver itself errors -- same
# fallback philosophy as the risk model's own degenerate-month handling
# throughout this file.
#
# Returns data.frame(id, eom, w) -- Rule 12's required output shape.
construct_qp_portfolio <- function(test, eom_col, daily_ret_keyed, months = NULL) {
  t0 <- Sys.time()
  if (is.null(months)) months <- unique(test[[eom_col]])
  months <- sort(months)
  test <- test[test[[eom_col]] %in% months, ]

  test$w <- NA_real_
  n_qp_months <- 0L
  n_fallback_months <- 0L
  n_months <- length(months)
  log_stage(sprintf("entering QP construction loop: %d months", n_months))
  for (i in seq_along(months)) {
    d <- months[i]
    idx <- which(test[[eom_col]] == d)
    ids_month <- test$id[idx]
    n <- length(idx)
    alpha <- test$pred[idx]  # the ensemble's raw predicted return -- same
                              # units as ret_exc_lead1m and, after the
                              # daily->monthly scaling above, the same
                              # holding-period units as Sigma. NOT
                              # test$rank: an arbitrary-scale ordinal
                              # transform with no natural relationship to
                              # Sigma's units was the root cause of an
                              # earlier degenerate real-data solution.

    qp_inputs <- build_month_qp_inputs(alpha, ids_month, daily_ret_keyed, d)
    w_month <- NULL
    if (!is.null(qp_inputs)) {
      Sigma <- build_sigma_from_factors(qp_inputs)
      w_raw <- tryCatch(
        solve_portfolio_qp(alpha = qp_inputs$alpha, Sigma = Sigma, w_prev = rep(0, n),
                            lambda = QP_LAMBDA, kappa = QP_KAPPA),
        error = function(e) NULL
      )
      rm(Sigma)
      if (!is.null(w_raw)) {
        # Re-center BEFORE rescaling, not just rescale: osqp's dollar-neutral
        # constraint only holds to its solver tolerance, and dividing by
        # sum(abs(w_raw)) can amplify that tiny residual past a meaningful
        # dollar-neutral tolerance if the raw solution happens to be small
        # in magnitude. Explicitly demeaning first guarantees exact
        # sum(w)=0 regardless of the raw solution's scale.
        w_c <- w_raw - mean(w_raw)
        denom <- sum(abs(w_c))
        if (is.finite(denom) && denom > 0) {
          w_month <- w_c / denom * TARGET_GROSS_EXPOSURE
          n_qp_months <- n_qp_months + 1L
        }
      }
    }
    if (is.null(w_month)) {
      n_fallback_months <- n_fallback_months + 1L
      a <- alpha - mean(alpha)
      denom <- sum(abs(a))
      w_month <- if (is.finite(denom) && denom > 0) a / denom * TARGET_GROSS_EXPOSURE else rep(0, n)
    }

    test$w[idx] <- w_month

    if (i %% 50 == 0 || i == n_months) {
      elapsed_min <- as.numeric(difftime(Sys.time(), t0, units = "mins"))
      log_stage(sprintf("QP construction loop progress: %d/%d months (elapsed %.1f min, ~%.2f sec/month)",
                         i, n_months, elapsed_min, elapsed_min * 60 / i))
    }
  }
  # final safety net regardless of the logic above: the output must never
  # contain NA/Inf (Rule 12), so any leftover non-finite value becomes 0
  test$w[!is.finite(test$w)] <- 0

  cat(sprintf("[main] QP solve used: %d of %d months (rest used plain rank-demean fallback)\n",
              n_qp_months, n_months))

  # the source id column comes through as double (e.g. Postgres bigint via
  # RPostgres); Rule 12 requires an integer id column in the output
  data.frame(id = as.integer(test$id), eom = test[[eom_col]], w = test$w)
}

main <- function(chars, features, daily_ret) {
  # Rule 18: seed at the start of main(). fit_predictions() re-seeds with the
  # same value and xgboost pins its own seed via params$seed, so this adds
  # no new randomness; it only makes the requirement explicit.
  set.seed(42)
  log_step("forecasting model (fit on training months, predict test months)")
  fp <- fit_predictions(chars, features)
  log_step("indexing daily returns for the risk model")
  daily_ret_keyed <- prepare_daily_ret(daily_ret)
  log_step("monthly risk model + mean-variance optimization")
  out <- construct_qp_portfolio(fp$test, fp$eom_col, daily_ret_keyed)
  log_stage("main() completed")
  out
}