# MGT 924 PS3: Portfolio-ML for the JKP Common Task Framework
#
# Idea: instead of predicting returns and then forming weights, choose weights
# that maximize the Sharpe ratio directly (Portfolio-ML / ridge Markowitz on
# characteristic-managed portfolios).
#
#   1. Each month, rank every characteristic across stocks to [-0.5, 0.5]
#      (missing = 0). Optionally add random Fourier features of the ranks,
#      sin(S W) and cos(S W), to allow nonlinear effects. Every feature is
#      demeaned and scaled to unit sd across stocks within the month.
#      A constant column (the equal-weight market) is added.
#   2. Each feature k defines a portfolio with weights Z[, k] / N. Its return
#      next month is f_k = Z[, k]' r / N.
#   3. The portfolio is w = Z beta / N with
#        beta = (F'F / T + lambda I)^-1  F'1 / T  (solved in dual form),
#      the ridge-shrunk mean-variance combination of the factor portfolios,
#      estimated on all past months (expanding window, refit every January).
#   4. lambda is chosen each year by 5-fold CV over contiguous blocks of past
#      years, maximizing the out-of-fold Sharpe ratio.
#
# Timing: data is processed in calendar order. beta for year Y uses only
# factor returns from rows with eom in years < Y. The last of these is the
# return over January Y, realized at eom January Y, the first month that uses
# beta. So no lead return of month t is ever used for month t weights.
#
# Base R only. set.seed fixes the random features, so output is deterministic.

CFG <- list(
  n_rf = 3000,     # random features (each gives a sin and a cos column); 0 = linear
  gamma = 1,       # scale of the random weights
  market = TRUE,   # include the equal-weight market as a factor
  min_months = 60, # months of history needed before fitting
  seed = 1
)

# Rank a vector to [-0.5, 0.5]; NA -> 0; a constant column -> 0
rank_std <- function(x) {
  ok <- !is.na(x)
  n <- sum(ok)
  out <- numeric(length(x))
  if (n > 1) out[ok] <- (rank(x[ok], ties.method = "average") - 1) / (n - 1) - 0.5
  out
}

# Feature matrix for one month: N x P
make_features <- function(d, feats, W, market) {
  S <- vapply(feats, function(f) rank_std(as.numeric(d[[f]])), numeric(nrow(d)))
  S <- matrix(S, nrow = nrow(d))
  if (!is.null(W)) {
    A <- S %*% W
    S <- cbind(S, sin(A), cos(A))
  }
  mu <- colMeans(S)
  sdv <- sqrt(colMeans(S^2) - mu^2)
  sdv[!is.finite(sdv) | sdv < 1e-12] <- Inf  # constant columns become 0
  S <- sweep(sweep(S, 2, mu), 2, sdv, "/")
  if (market) S <- cbind(1, S)
  S
}

# Ridge Markowitz in dual form, using the T x T Gram matrix G = F F':
#   beta = (F'F / T + lam I)^-1 F'1 / T = F' (G + lam T I)^-1 1
# The dual form is fast when there are more factors than months.
pml_beta <- function(F, G, lam) {
  n <- nrow(F)
  drop(crossprod(F, solve(G + diag(lam * n, n), rep(1, n))))
}

# Choose lambda by blocked CV over past years, maximizing out-of-fold Sharpe.
# Out-of-fold returns F_te beta_tr = G[te, tr] (G[tr, tr] + lam n I)^-1 1.
cv_lambda <- function(G, F_year, grid, n_folds = 5) {
  yrs <- sort(unique(F_year))
  n_folds <- min(n_folds, length(yrs))
  if (n_folds < 2) return(grid[ceiling(length(grid) / 2)])
  fold_of_year <- cut(seq_along(yrs), n_folds, labels = FALSE)
  fold <- fold_of_year[match(F_year, yrs)]
  oos <- matrix(NA_real_, nrow(G), length(grid))
  for (k in seq_len(n_folds)) {
    tr <- fold != k
    e <- eigen(G[tr, tr, drop = FALSE], symmetric = TRUE)
    ev <- pmax(e$values, 0)
    u1 <- colSums(e$vectors)
    A <- G[!tr, tr, drop = FALSE] %*% e$vectors
    for (g in seq_along(grid)) {
      oos[!tr, g] <- A %*% (u1 / (ev + grid[g] * sum(tr)))
    }
  }
  sr <- apply(oos, 2, function(x) mean(x) / sd(x))
  sr[!is.finite(sr)] <- -Inf
  grid[which.max(sr)]
}

# Core loop. get_year(Y) returns that year's rows (id, eom, ret_exc_lead1m, feats).
pml_engine <- function(get_year, years, out_months, feats, cfg = CFG, verbose = TRUE) {
  set.seed(cfg$seed)
  K <- length(feats)
  W <- if (cfg$n_rf > 0) matrix(rnorm(K * cfg$n_rf, sd = cfg$gamma * sqrt(12 / K)), K) else NULL
  out_years <- sort(unique(as.integer(format(out_months, "%Y"))))
  F_rows <- list(); F_year <- integer(0)
  out <- list()
  for (Y in years) {
    beta <- NULL
    if (Y %in% out_years && length(F_year) >= cfg$min_months) {
      F <- do.call(rbind, F_rows)
      G <- tcrossprod(F)
      scale <- sum(diag(G)) / nrow(F) / ncol(F)  # average eigenvalue of F'F / T
      grid <- scale * 10^seq(-3, 2, by = 0.5)
      lam <- cv_lambda(G, F_year, grid)
      beta <- pml_beta(F, G, lam)
      if (verbose) cat("Year", Y, " months:", nrow(F), " factors:", ncol(F),
                       " lambda/scale:", signif(lam / scale, 2), "\n")
    }
    d_year <- get_year(Y)
    if (is.null(d_year) || nrow(d_year) == 0) next
    for (r in split(seq_len(nrow(d_year)), d_year$eom)) {
      d <- d_year[r, , drop = FALSE]
      Z <- make_features(d, feats, W, cfg$market)
      if (d$eom[1] %in% out_months) {
        w <- if (is.null(beta)) rep(1 / nrow(d), nrow(d)) else drop(Z %*% beta) / nrow(d)
        out[[length(out) + 1]] <- data.frame(id = as.integer(d$id), eom = d$eom, w = w)
      }
      y <- d$ret_exc_lead1m
      ok <- !is.na(y)
      if (sum(ok) >= 2) {
        F_rows[[length(F_rows) + 1]] <- drop(crossprod(Z[ok, , drop = FALSE], y[ok])) / sum(ok)
        F_year <- c(F_year, Y)
      }
    }
  }
  out <- do.call(rbind, out)
  out[order(out$eom, out$id), ]
}

# ctff_test can arrive as logical, number or text ("1"/"0", "True"/"False")
as_flag <- function(x) {
  if (is.logical(x)) return(!is.na(x) & x)
  toupper(trimws(as.character(x))) %in% c("1", "TRUE", "T")
}

# Months to output: test months (all months if none are flagged)
output_months <- function(eom, ctff_test) {
  m <- sort(unique(eom[as_flag(ctff_test)]))
  if (length(m) == 0) m <- sort(unique(eom))
  m
}

main <- function(chars, features, daily_ret) {
  set.seed(1)
  chars <- as.data.frame(chars)
  feats <- sort(intersect(as.character(features$features), names(chars)))
  chars$eom <- as.Date(chars$eom)
  year <- as.integer(format(chars$eom, "%Y"))
  rows_by_year <- split(seq_len(nrow(chars)), year)
  years <- sort(as.integer(names(rows_by_year)))
  cat("Rows:", nrow(chars), " Years:", length(years), " Features:", length(feats), "\n")
  get_year <- function(Y) chars[rows_by_year[[as.character(Y)]], , drop = FALSE]
  months <- output_months(chars$eom, chars$ctff_test)
  pml_engine(get_year, years, months, feats, CFG)
}