"""JKP Common Task Framework submission.

Author: Abhinav Keshri <abhinav.keshri@yale.edu>

Interface (Rules 11, 12, 14):

    def main(chars, features, daily_ret) -> pd.DataFrame   # columns exactly id, eom, w

The evaluation pipeline loads the three parquet frames, passes them to main, and captures the
returned frame. This file reads no external files and uses no network, subprocess or dynamic
code execution (Rules 10, 15). Seeds are fixed, so repeated runs give identical output (Rule 18).
Every quantity used at month t is computed from data dated at or before t (Rule 1).

METHOD

A single gradient boosted model forecasts the cross section of next month returns from all 402
characteristics. The forecast is then held simultaneously in two portfolios built on different
risk models, and the combined book is scaled each month by a volatility target.

    Alpha      XGBoost on the 402 raw characteristics, missing values retained. The training
               target is arctan of the within month standardised lead excess return. Expanding
               window, annual refit, hyperparameters searched by Optuna every 36 months.
    Book 1     Ledoit and Wolf shrinkage applied to a 756 day daily covariance, then Tikhonov
               regularisation, then the tangency portfolio minus the global minimum variance
               portfolio. Weight 50 percent.
    Book 2     Fundamental factor model with 13 sign aligned JKP theme exposures, inverted by
               the Woodbury identity. Weight 50 percent.
    Overlay    Gross exposure scaled by a volatility target divided by trailing 24 month
               volatility, bounded within [0.75, 1.50].

MOTIVATION FOR THE TWO BOOK CONSTRUCTION

The conventional route to improving a machine learning equity strategy is to combine several
return forecasts. We evaluated 31 trained models across six combination schemes, including equal
weighting, correlation clustering, trailing information coefficient weighting, trailing Sharpe
weighting, ridge stacking and non negative least squares. Every scheme performed at or below the
strongest individual model. The explanation is that two forecasts evaluated through the same risk
model produce monthly return streams with a correlation of approximately 0.95, which leaves
almost no diversification available.

The same forecast evaluated through two different risk models produces streams with a correlation
of 0.65. Measured on ctff_test months over 1990 to 2023, using this model's own forecasts:

    daily covariance book alone      Sharpe 3.36, 2010 onward 2.31, 2020 onward 1.25, drawdown 9.0 percent
    both books at equal weight       Sharpe 3.34, 2010 onward 2.47, 2020 onward 1.69, drawdown 10.1 percent
    both books with the overlay      Sharpe 3.50, 2010 onward 2.65, 2020 onward 1.64, drawdown 8.7 percent

Combining the books leaves the full sample result unchanged and improves the period after 2010, and
that pattern holds across both runs of the pipeline we examined. The overlay raises the full sample
result. Turnover is approximately 0.5 in all three cases. The accompanying methodology document
reports the validation procedure and the alternatives that were rejected.

PRIOR WORK

Three components follow the CTF entry "LightGBM expected returns and daily shrunk VCV" by Paul
Geertsema and Various Agents, dated 30 August 2026. They are the arctan transform of the
standardised target, the Ledoit and Wolf daily covariance with Tikhonov regularisation, and the
tangency minus minimum variance book. Each is marked at its definition below. The underlying
statistical methods are established in the literature (Ledoit and Wolf 2004; Markowitz 1952;
Tikhonov 1963). All hyperparameters in this submission are produced by the search procedure
described above.
"""
from __future__ import annotations

import random
import warnings

import numpy as np
import pandas as pd

warnings.filterwarnings("ignore", message="X does not have valid feature names")

SEED = 42
ID_COL, DATE_COL, RET_COL, TEST_COL = "id", "eom", "ret_exc_lead1m", "ctff_test"

# ---- walk forward schedule ----
MIN_TRAIN_MONTHS = 120      # minimum of ten years of training data before the first forecast
VAL_MONTHS = 24             # tail of the training window used for early stopping / tuning
RETRAIN_EVERY = 12          # refit annually
TUNE_EVERY = 36             # re-run the hyperparameter search every third refit
OPTUNA_TRIALS = 60          # upper bound; the patience rule usually stops earlier
OPTUNA_PATIENCE = 5
OPTUNA_MIN_DELTA = 1e-4
TUNE_SUBSAMPLE = 200_000    # trials fit on the most recent rows; final refit uses everything

# ---- risk books ----
WINDOW_LONG = 756                        # trading days in the covariance window
MIN_DAILY_OBS, DAYS_PER_MONTH, MIN_NAMES, KAPPA = 250, 21, 50, 1.0
MVO_WINDOW = 60             # months of factor history
D_FLOOR_Q = 0.10
BOOK_WEIGHTS = {"dc_long": 0.50, "mvo": 0.50}

# Volatility targeting (Moreira and Muir, 2017). Month t gross exposure is scaled by the target
# divided by trailing volatility. Realised volatility is persistent and forecastable while expected
# returns are much less so, and the evaluation rescales the whole stream to 10 percent volatility
# after the fact, so reducing exposure in volatile months raises the scored Sharpe ratio. The target
# is an expanding mean of past trailing volatilities, and the multiplier is bounded within
# [0.75, 1.50].
VOL_WINDOW, VOL_CLIP = 24, (0.75, 1.5)

THEMES = ["Accruals", "Debt Issuance", "Investment", "Low Leverage", "Low Risk",
          "Momentum", "Profit Growth", "Profitability", "Quality", "Seasonality",
          "Short-Term Reversal", "Size", "Value"]

# (characteristic, JKP theme, direction) for the 153 themed characteristics. Used only to build
# the 13 exposure columns of the fundamental risk model; the alpha model sees all 402 raw
# characteristics and is told nothing about themes.
THEME_MAP = [
    ('age', 'Low Leverage', -1.0),
    ('aliq_at', 'Investment', -1.0),
    ('aliq_mat', 'Low Leverage', -1.0),
    ('ami_126d', 'Size', +1.0),
    ('at_be', 'Low Leverage', -1.0),
    ('at_gr1', 'Investment', -1.0),
    ('at_me', 'Value', +1.0),
    ('at_turnover', 'Quality', +1.0),
    ('be_gr1a', 'Investment', -1.0),
    ('be_me', 'Value', +1.0),
    ('beta_60m', 'Low Risk', -1.0),
    ('beta_dimson_21d', 'Low Risk', -1.0),
    ('betabab_1260d', 'Low Risk', -1.0),
    ('betadown_252d', 'Low Risk', -1.0),
    ('bev_mev', 'Value', +1.0),
    ('bidaskhl_21d', 'Low Leverage', +1.0),
    ('capex_abn', 'Debt Issuance', -1.0),
    ('capx_gr1', 'Investment', -1.0),
    ('capx_gr2', 'Investment', -1.0),
    ('capx_gr3', 'Investment', -1.0),
    ('cash_at', 'Low Leverage', +1.0),
    ('chcsho_12m', 'Value', -1.0),
    ('coa_gr1a', 'Investment', -1.0),
    ('col_gr1a', 'Investment', -1.0),
    ('cop_at', 'Quality', +1.0),
    ('cop_atl1', 'Quality', +1.0),
    ('corr_1260d', 'Seasonality', -1.0),
    ('coskew_21d', 'Seasonality', -1.0),
    ('cowc_gr1a', 'Accruals', -1.0),
    ('dbnetis_at', 'Seasonality', -1.0),
    ('debt_gr3', 'Debt Issuance', -1.0),
    ('debt_me', 'Value', +1.0),
    ('dgp_dsale', 'Quality', +1.0),
    ('div12m_me', 'Value', +1.0),
    ('dolvol_126d', 'Size', -1.0),
    ('dolvol_var_126d', 'Profitability', -1.0),
    ('dsale_dinv', 'Profit Growth', +1.0),
    ('dsale_drec', 'Profit Growth', -1.0),
    ('dsale_dsga', 'Profit Growth', +1.0),
    ('earnings_variability', 'Low Risk', -1.0),
    ('ebit_bev', 'Profitability', +1.0),
    ('ebit_sale', 'Profitability', +1.0),
    ('ebitda_mev', 'Value', +1.0),
    ('emp_gr1', 'Investment', -1.0),
    ('eq_dur', 'Value', -1.0),
    ('eqnetis_at', 'Value', -1.0),
    ('eqnpo_12m', 'Value', +1.0),
    ('eqnpo_me', 'Value', +1.0),
    ('eqpo_me', 'Value', +1.0),
    ('f_score', 'Profitability', +1.0),
    ('fcf_me', 'Value', +1.0),
    ('fnl_gr1a', 'Debt Issuance', -1.0),
    ('gp_at', 'Quality', +1.0),
    ('gp_atl1', 'Quality', +1.0),
    ('inv_gr1', 'Investment', -1.0),
    ('inv_gr1a', 'Investment', -1.0),
    ('iskew_capm_21d', 'Short-Term Reversal', -1.0),
    ('iskew_ff3_21d', 'Short-Term Reversal', -1.0),
    ('iskew_hxz4_21d', 'Short-Term Reversal', -1.0),
    ('ival_me', 'Value', +1.0),
    ('ivol_capm_21d', 'Low Risk', -1.0),
    ('ivol_capm_252d', 'Low Risk', -1.0),
    ('ivol_ff3_21d', 'Low Risk', -1.0),
    ('ivol_hxz4_21d', 'Low Risk', -1.0),
    ('kz_index', 'Seasonality', +1.0),
    ('lnoa_gr1a', 'Investment', -1.0),
    ('lti_gr1a', 'Seasonality', -1.0),
    ('market_equity', 'Size', -1.0),
    ('mispricing_mgmt', 'Investment', +1.0),
    ('mispricing_perf', 'Quality', +1.0),
    ('ncoa_gr1a', 'Investment', -1.0),
    ('ncol_gr1a', 'Debt Issuance', -1.0),
    ('netdebt_me', 'Low Leverage', -1.0),
    ('netis_at', 'Value', -1.0),
    ('nfna_gr1a', 'Debt Issuance', +1.0),
    ('ni_ar1', 'Debt Issuance', +1.0),
    ('ni_be', 'Profitability', +1.0),
    ('ni_inc8q', 'Quality', +1.0),
    ('ni_ivol', 'Low Leverage', +1.0),
    ('ni_me', 'Value', +1.0),
    ('niq_at', 'Quality', +1.0),
    ('niq_at_chg1', 'Profit Growth', +1.0),
    ('niq_be', 'Profitability', +1.0),
    ('niq_be_chg1', 'Profit Growth', +1.0),
    ('niq_su', 'Profit Growth', +1.0),
    ('nncoa_gr1a', 'Investment', -1.0),
    ('noa_at', 'Debt Issuance', -1.0),
    ('noa_gr1a', 'Investment', -1.0),
    ('o_score', 'Profitability', -1.0),
    ('oaccruals_at', 'Accruals', -1.0),
    ('oaccruals_ni', 'Accruals', -1.0),
    ('ocf_at', 'Profitability', +1.0),
    ('ocf_at_chg1', 'Profit Growth', +1.0),
    ('ocf_me', 'Value', +1.0),
    ('ocfq_saleq_std', 'Low Risk', -1.0),
    ('op_at', 'Quality', +1.0),
    ('op_atl1', 'Quality', +1.0),
    ('ope_be', 'Profitability', +1.0),
    ('ope_bel1', 'Profitability', +1.0),
    ('opex_at', 'Quality', +1.0),
    ('pi_nix', 'Seasonality', +1.0),
    ('ppeinv_gr1a', 'Investment', -1.0),
    ('prc', 'Size', -1.0),
    ('prc_highprc_252d', 'Momentum', +1.0),
    ('qmj', 'Quality', +1.0),
    ('qmj_growth', 'Quality', +1.0),
    ('qmj_prof', 'Quality', +1.0),
    ('qmj_safety', 'Quality', +1.0),
    ('rd_me', 'Size', +1.0),
    ('rd_sale', 'Low Leverage', +1.0),
    ('rd5_at', 'Low Leverage', +1.0),
    ('resff3_12_1', 'Momentum', +1.0),
    ('resff3_6_1', 'Momentum', +1.0),
    ('ret_1_0', 'Short-Term Reversal', -1.0),
    ('ret_12_1', 'Momentum', +1.0),
    ('ret_12_7', 'Profit Growth', +1.0),
    ('ret_3_1', 'Momentum', +1.0),
    ('ret_6_1', 'Momentum', +1.0),
    ('ret_60_12', 'Investment', -1.0),
    ('ret_9_1', 'Momentum', +1.0),
    ('rmax1_21d', 'Low Risk', -1.0),
    ('rmax5_21d', 'Low Risk', -1.0),
    ('rmax5_rvol_21d', 'Short-Term Reversal', -1.0),
    ('rskew_21d', 'Short-Term Reversal', -1.0),
    ('rvol_21d', 'Low Risk', -1.0),
    ('sale_bev', 'Quality', +1.0),
    ('sale_emp_gr1', 'Profit Growth', +1.0),
    ('sale_gr1', 'Investment', -1.0),
    ('sale_gr3', 'Investment', -1.0),
    ('sale_me', 'Value', +1.0),
    ('saleq_gr1', 'Investment', -1.0),
    ('saleq_su', 'Profit Growth', +1.0),
    ('seas_1_1an', 'Profit Growth', +1.0),
    ('seas_1_1na', 'Momentum', +1.0),
    ('seas_11_15an', 'Seasonality', +1.0),
    ('seas_11_15na', 'Seasonality', -1.0),
    ('seas_16_20an', 'Seasonality', +1.0),
    ('seas_16_20na', 'Accruals', -1.0),
    ('seas_2_5an', 'Seasonality', +1.0),
    ('seas_2_5na', 'Investment', -1.0),
    ('seas_6_10an', 'Seasonality', +1.0),
    ('seas_6_10na', 'Low Risk', -1.0),
    ('sti_gr1a', 'Seasonality', +1.0),
    ('taccruals_at', 'Accruals', -1.0),
    ('taccruals_ni', 'Accruals', -1.0),
    ('tangibility', 'Low Leverage', +1.0),
    ('tax_gr1a', 'Profit Growth', +1.0),
    ('turnover_126d', 'Low Risk', -1.0),
    ('turnover_var_126d', 'Profitability', -1.0),
    ('z_score', 'Low Leverage', +1.0),
    ('zero_trades_126d', 'Low Risk', +1.0),
    ('zero_trades_21d', 'Low Risk', +1.0),
    ('zero_trades_252d', 'Low Risk', +1.0),
]


# ===========================================================================================
# 1. Data preparation
# ===========================================================================================
def _sector(chars: pd.DataFrame) -> pd.Series:
    """Coarse sector: 2-digit GICS, else 2-digit SIC, else -1."""
    for col, div in (("gics", 1_000_000), ("sic", 100)):
        if col in chars.columns and chars[col].notna().any():
            return (pd.to_numeric(chars[col], errors="coerce") // div).fillna(-1).astype("int64")
    return pd.Series(-1, index=chars.index, dtype="int64")


def _xs_rank_centered(df: pd.DataFrame, cols: list[str], eom: pd.Series) -> pd.DataFrame:
    """Within-month percentile rank, centered to [-0.5, 0.5]. NaN stays NaN."""
    return df[cols].groupby(eom, observed=True).rank(pct=True) - 0.5


def theme_exposures(chars: pd.DataFrame, eom: pd.Series) -> pd.DataFrame:
    """13 sign-aligned theme averages: rank -> center -> multiply by direction -> mean."""
    present = [(c, t, d) for c, t, d in THEME_MAP if c in chars.columns]
    cols = [c for c, _, _ in present]
    ranks = _xs_rank_centered(chars, cols, eom)
    sign = {c: d for c, _, d in present}
    out = {}
    for theme in THEMES:
        members = [c for c, t, _ in present if t == theme]
        if members:
            out[theme] = ranks[members].mul([sign[c] for c in members], axis=1).mean(axis=1)
    return pd.DataFrame(out, index=chars.index).fillna(0.0).astype(np.float32)


def prepare(chars: pd.DataFrame, features: pd.DataFrame) -> dict:
    """Sort by (eom, id) and build the model matrix, target, and risk exposures."""
    feat_list = [f for f in features["features"].astype(str).tolist() if f in chars.columns]
    df = chars.copy()
    df[DATE_COL] = pd.to_datetime(df[DATE_COL])
    df = df.sort_values([DATE_COL, ID_COL], kind="stable").reset_index(drop=True)

    # Alpha features: the raw characteristic values, missing left missing. XGBoost learns a
    # default direction per split, which beat both mean and rank imputation in our tests.
    X = df[feat_list].to_numpy(np.float32)

    # (a) Target: arctan of the within month standardised return.
    # Credit: Paul Geertsema and Various Agents, CTF entry "LightGBM expected returns and daily
    # shrunk VCV" (30 August 2026).
    r = pd.to_numeric(df[RET_COL], errors="coerce")
    g = r.groupby(df[DATE_COL], observed=True)
    y = np.arctan(((r - g.transform("mean")) / g.transform("std")).to_numpy(np.float64))

    test = (df[TEST_COL].astype(bool).to_numpy() if TEST_COL in df.columns
            else np.zeros(len(df), bool))
    themes = theme_exposures(df, df[DATE_COL])

    eom = df[DATE_COL].to_numpy()
    months, first = np.unique(eom, return_index=True)
    bounds = np.append(first, len(df))
    slc = {m: (int(bounds[i]), int(bounds[i + 1])) for i, m in enumerate(months)}
    return dict(df=df, X=X, y=y, eom=eom, months=months, slc=slc, test=test,
                ids=df[ID_COL].to_numpy(), ret=r.to_numpy(np.float64),
                themes=themes.to_numpy(np.float64), feat_list=feat_list,
                sector=_sector(df).to_numpy())


# ===========================================================================================
# 2. Alpha model. XGBoost with hyperparameters from the Optuna search described in the header.
# ===========================================================================================
def _suggest(trial) -> dict:
    """Our search space. Wide enough to reach heavily-regularised, high-capacity settings."""
    return {
        "n_estimators": 2000,
        "learning_rate": trial.suggest_float("learning_rate", 5e-3, 1.5e-1, log=True),
        "max_leaves": trial.suggest_int("max_leaves", 15, 511, log=True),
        "max_depth": trial.suggest_int("max_depth", 3, 14),
        "min_child_weight": trial.suggest_int("min_child_weight", 20, 5000, log=True),
        "subsample": trial.suggest_float("subsample", 0.6, 1.0),
        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.6, 1.0),
        "reg_lambda": trial.suggest_float("reg_lambda", 1e-2, 1e2, log=True),
    }


def _default_params() -> dict:
    """Used on the tiny validation panel and if Optuna is unavailable."""
    return {"n_estimators": 400, "learning_rate": 0.03, "max_leaves": 63, "max_depth": 8,
            "min_child_weight": 200, "subsample": 0.9, "colsample_bytree": 0.8,
            "reg_lambda": 1.0}


def _fit(Xtr, ytr, Xva, yva, params, seed=SEED):
    import xgboost as xgb

    p = dict(params)
    n_est = p.pop("n_estimators", 2000)
    es = {"early_stopping_rounds": 50} if len(Xva) >= 50 else {}
    model = xgb.XGBRegressor(objective="reg:squarederror", tree_method="hist",
                             grow_policy="lossguide", n_jobs=-1, random_state=seed,
                             verbosity=0, n_estimators=n_est, **es, **p)
    if len(Xva) >= 50:
        model.fit(Xtr, ytr, eval_set=[(Xva, yva)], verbose=False)
    else:
        model.fit(Xtr, ytr, verbose=False)
    return model


def _fit_refit_full(Xtr, ytr, Xva, yva, params, seed=SEED):
    """Fit with early stopping, then REFIT on train+validation at the selected tree count.

    The holdout's job is to choose the number of trees and the hyperparameters; once it has done
    that, excluding it from the final fit discards the most recent two years, the years most
    like the month we are about to predict. So we refit on core+val with n_estimators fixed to
    what the holdout selected (an empty eval set disables early stopping, so the count is exact).
    Omitting this step cost 0.35 Sharpe in the 2020s in our first packaged run.
    """
    model = _fit(Xtr, ytr, Xva, yva, params, seed=seed)
    n_best = getattr(model, "best_iteration", None)
    if not n_best or len(Xva) < 50:
        return model
    p_full = {**params, "n_estimators": int(n_best) + 1}
    return _fit(np.vstack([Xtr, Xva]), np.concatenate([ytr, yva]),
                Xva[:0], yva[:0], p_full, seed=seed)


def _rank_ic(pred, y, months) -> float:
    """Mean monthly Spearman correlation between prediction and target."""
    d = pd.DataFrame({"m": months, "p": pred, "y": y})
    ics = d.groupby("m").apply(lambda x: x["p"].corr(x["y"], method="spearman"),
                               include_groups=False)
    return float(ics.mean()) if len(ics) else -1.0


def _tune(Xtr, ytr, Xva, yva, va_months, prev_best, seed=SEED):
    """Maximise validation rank-IC. Returns (params, raw_trial_dict)."""
    try:
        import optuna
    except ImportError:
        return _default_params(), None
    optuna.logging.set_verbosity(optuna.logging.WARNING)

    if TUNE_SUBSAMPLE and len(Xtr) > TUNE_SUBSAMPLE:
        Xtr, ytr = Xtr[-TUNE_SUBSAMPLE:], ytr[-TUNE_SUBSAMPLE:]

    def objective(trial):
        params = _suggest(trial)
        model = _fit(Xtr, ytr, Xva, yva, params, seed=seed)
        trial.set_user_attr("full_params", params)
        return _rank_ic(model.predict(Xva), yva, va_months)

    study = optuna.create_study(direction="maximize",
                                sampler=optuna.samplers.TPESampler(seed=seed))
    if prev_best:
        study.enqueue_trial(prev_best)      # warm-start from the previous window's choice

    state = {"best": -np.inf, "since": 0}

    def patience(st, tr):
        if tr.value is None:
            return
        if tr.value > state["best"] + OPTUNA_MIN_DELTA:
            state["best"], state["since"] = tr.value, 0
        else:
            state["since"] += 1
            if state["since"] >= OPTUNA_PATIENCE:
                st.stop()

    study.optimize(objective, n_trials=OPTUNA_TRIALS, callbacks=[patience],
                   show_progress_bar=False)
    best = study.best_trial
    return best.user_attrs.get("full_params", _default_params()), dict(best.params)


def walk_forward(data: dict, verbose: bool = True) -> np.ndarray:
    """Expanding-window predictions for every month we need to hold a position in."""
    X, y, months, slc = data["X"], data["y"], data["months"], data["slc"]
    n_months = len(months)
    small = n_months < 150      # the ~123-month validation panel (Rule 14)
    if small:
        min_train, val_months = max(12, n_months // 3), max(6, n_months // 10)
        min_tr_obs, min_va_obs, use_optuna = 50, 10, False
    else:
        min_train, val_months = MIN_TRAIN_MONTHS, VAL_MONTHS
        min_tr_obs, min_va_obs, use_optuna = 500, 50, True

    if data["test"].any():
        predict_set = {m for m in months if data["test"][slc[m][0]:slc[m][1]].any()}
    else:
        predict_set = set(months[min_train:])

    preds = np.full(len(X), np.nan)
    params, prev_best, model, n_done = _default_params(), None, None, 0
    for ti, t in enumerate(months):
        if t not in predict_set:
            continue
        avail = months[:ti]
        if len(avail) < min_train:
            continue
        core = avail[:-val_months] if len(avail) > val_months else avail
        vals = avail[-val_months:] if len(avail) > val_months else avail
        c_lo, c_hi = slc[core[0]][0], slc[core[-1]][1]
        v_lo, v_hi = slc[vals[0]][0], slc[vals[-1]][1]
        Xc, yc = X[c_lo:c_hi], y[c_lo:c_hi]
        Xv, yv = X[v_lo:v_hi], y[v_lo:v_hi]
        m_tr, m_va = np.isfinite(yc), np.isfinite(yv)
        if m_tr.sum() < min_tr_obs or m_va.sum() < min_va_obs:
            continue
        va_months = data["eom"][v_lo:v_hi][m_va]
        Xc, yc, Xv, yv = Xc[m_tr], yc[m_tr], Xv[m_va], yv[m_va]

        if use_optuna and (n_done % TUNE_EVERY == 0):
            params, prev_best = _tune(Xc, yc, Xv, yv, va_months, prev_best)
            if verbose:
                print(f"  [tune {pd.Timestamp(t).date()}] {prev_best}", flush=True)
            model = None
        if (n_done % RETRAIN_EVERY == 0) or model is None:
            model = _fit_refit_full(Xc, yc, Xv, yv, params)
            if verbose:
                print(f"  [fit  {pd.Timestamp(t).date()}] train={len(Xc):,}", flush=True)
        lo, hi = slc[t]
        preds[lo:hi] = model.predict(X[lo:hi])
        n_done += 1
    return preds


def standardize(mu: np.ndarray, eom: np.ndarray) -> np.ndarray:
    """Per-month z-score of mu-hat. Free for tangency-GMV, worth real Sharpe for factor MVO."""
    s = pd.Series(mu, index=eom)
    g = s.groupby(level=0)
    z = (s - g.transform("mean")) / g.transform("std")
    return np.nan_to_num(z.to_numpy(), nan=0.0)


# ===========================================================================================
# 3. Risk books
# ===========================================================================================
def build_daily_matrix(daily_ret: pd.DataFrame, panel_ids: np.ndarray):
    dr = daily_ret.rename(columns={c: c.lower() for c in daily_ret.columns})
    rid = dr["id"].to_numpy()
    keep = np.isin(rid, panel_ids)
    rid = rid[keep]
    rdate = pd.to_datetime(dr["date"]).to_numpy()[keep].astype("datetime64[D]")
    rret = dr["ret_exc"].to_numpy()[keep].astype(np.float32)
    dates, di = np.unique(rdate, return_inverse=True)
    ids, ii = np.unique(rid, return_inverse=True)
    M = np.full((len(dates), len(ids)), np.nan, dtype=np.float32)
    M[di, ii] = rret
    return dates, ids, M


def trailing_window(M, end, cols, window):
    lo = max(0, end - window)
    X = M[lo:end, cols]
    obs = np.isfinite(X)
    enough = obs.sum(axis=0) >= min(MIN_DAILY_OBS, max(20, (end - lo) // 3))
    X = np.where(obs, X, np.float32(0.0))
    X -= X.mean(axis=0, keepdims=True)
    return X, enough


# (b) Ledoit and Wolf shrinkage of the daily covariance.
# Credit: Paul Geertsema and Various Agents, CTF entry "LightGBM expected returns and daily
# shrunk VCV" (30 August 2026).
def shrunk_covariance(X):
    """(b) Ledoit-Wolf shrinkage toward a scaled identity (Ledoit & Wolf 2004), as used by
    Geertsema et al., see header. X is the demeaned trailing daily return matrix."""
    T, n = X.shape
    S = (X.T @ X) / T
    tr_S = float(np.trace(S, dtype=np.float64))
    mu_lw = tr_S / n
    norm_S2 = float(np.square(S).sum(dtype=np.float64))
    d2 = norm_S2 - 2.0 * mu_lw * tr_S + n * mu_lw ** 2
    row_sq = np.square(X).sum(axis=1, dtype=np.float64)
    b2 = float(np.square(row_sq).sum()) / T ** 2 - norm_S2 / T
    delta = min(max(b2, 0.0), d2) / d2 if d2 > 1e-300 else 1.0
    S *= np.float32(1.0 - delta)
    S[np.diag_indices(n)] += np.float32(delta * mu_lw)
    return S


# (c) Tikhonov regularisation, then tangency minus global minimum variance portfolio.
# Credit: Paul Geertsema and Various Agents, CTF entry "LightGBM expected returns and daily
# shrunk VCV" (30 August 2026).
def tan_minus_gmv(mu, Sigma, kappa=KAPPA):
    """(c) Tikhonov ridge then tangency minus GMV, scaled to unit gross (Geertsema et al.)."""
    n = mu.shape[0]
    lam = np.float32(kappa * float(np.trace(Sigma, dtype=np.float64)) / n)
    A = Sigma + lam * np.eye(n, dtype=np.float32)
    one = np.ones(n, dtype=np.float32)
    sol = np.linalg.solve(A, np.column_stack([mu, one]))
    a, b = sol[:, 0], sol[:, 1]
    w = a - np.float32(float(one @ a) / float(one @ b)) * b
    den = np.abs(w).sum()
    return w / den if den > 0 else w


def _woodbury_sigma_inv_mu(B, Omega, d, mu):
    """Sigma^-1 mu for Sigma = B Omega B' + diag(d), without forming the N x N matrix."""
    K = B.shape[1]
    dinv = 1.0 / d
    BtDinv = B.T * dinv
    M = np.linalg.inv(Omega + 1e-8 * np.eye(K)) + BtDinv @ B
    return dinv * mu - dinv * (B @ np.linalg.solve(M, BtDinv @ mu))


def _month_book(data, dates, daily_ids, M, mu, fn):
    """Apply a per-month weight builder over the panel, returning unit-gross weights."""
    w_full = np.zeros(len(mu), np.float32)
    ids = data["ids"].astype(np.float64)
    col = np.clip(np.searchsorted(daily_ids, ids), 0, len(daily_ids) - 1)
    has = daily_ids[col] == ids
    for t in data["months"]:
        lo, hi = data["slc"][t]
        if not np.isfinite(mu[lo:hi]).any():
            continue
        cols = col[lo:hi][has[lo:hi]]
        end = int(np.searchsorted(dates, np.datetime64(pd.Timestamp(t), "D"), "right"))
        if end <= 20 or len(cols) < 2:
            continue
        w = fn(end, cols, np.nan_to_num(mu[lo:hi], nan=0.0), has[lo:hi], hi - lo)
        if w is not None:
            w_full[lo:hi] = w
    return w_full


def book_daily_cov(data, dates, daily_ids, M, mu, window):
    """Books (b)+(c): shrunk daily covariance -> ridge -> tangency minus GMV."""
    def fn(end, cols, mu_m, has_m, n_rows):
        Xd, enough = trailing_window(M, end, cols, window)
        usable = np.zeros(n_rows, bool)
        usable[np.flatnonzero(has_m)[enough]] = True
        if usable.sum() < min(MIN_NAMES, max(5, n_rows // 4)):
            return None
        Sig = shrunk_covariance(Xd[:, enough] * np.float32(np.sqrt(DAYS_PER_MONTH)))
        w = np.zeros(n_rows, np.float32)
        w[usable] = tan_minus_gmv(mu_m[usable], Sig)
        return w
    return _month_book(data, dates, daily_ids, M, mu, fn)


def book_factor_mvo(data, mu):
    """Fundamental factor MVO: Sigma = B Omega B' + D with B the 13 theme exposures, Omega the
    covariance of trailing cross sectional factor returns, D idiosyncratic residual variance."""
    months, slc = data["months"], data["slc"]
    Th, ret, ids = data["themes"], data["ret"], data["ids"]
    K = Th.shape[1]
    reg = 1e-4 * np.eye(K)

    # cross sectional factor returns and residuals, month by month
    F, resid = {}, np.full(len(ret), np.nan)
    for t in months:
        lo, hi = slc[t]
        Xm, rm = Th[lo:hi], ret[lo:hi]
        ok = np.isfinite(rm) & np.isfinite(Xm).all(axis=1)
        if ok.sum() <= K:
            continue
        Xo, ro = Xm[ok], rm[ok]
        f = np.linalg.solve(Xo.T @ Xo + reg, Xo.T @ ro)
        F[t] = f
        idx = np.arange(lo, hi)[ok]
        resid[idx] = ro - Xo @ f

    fmonths = np.array(sorted(F))
    fpos = {m: i for i, m in enumerate(fmonths)}
    Fv = np.array([F[m] for m in fmonths]) if len(fmonths) else np.zeros((0, K))
    window = MVO_WINDOW if len(fmonths) > 3 * MVO_WINDOW else max(12, len(fmonths) // 3)

    w_full = np.zeros(len(mu), np.float32)
    for t in months:
        lo, hi = slc[t]
        if t not in fpos or not np.isfinite(mu[lo:hi]).any():
            continue
        i = fpos[t]
        win = fmonths[max(0, i - window):i]              # strictly before t
        if len(win) < max(12, K):
            continue
        Omega = np.cov(Fv[max(0, i - window):i], rowvar=False)
        # residual variance per stock over the same window
        rows = np.concatenate([np.arange(*slc[m]) for m in win])
        rv = pd.Series(resid[rows], index=ids[rows]).groupby(level=0).var()
        d = pd.Series(ids[lo:hi]).map(rv).to_numpy(np.float64)
        fin = np.isfinite(d)
        if fin.sum() < 5:
            continue
        floor = max(float(np.nanquantile(d[fin], D_FLOOR_Q)), 1e-8)
        d = np.where(fin, np.maximum(d, floor), floor)
        B = Th[lo:hi].astype(np.float64)
        w = _woodbury_sigma_inv_mu(B, Omega, d, np.nan_to_num(mu[lo:hi], nan=0.0))
        w -= w.mean()
        den = np.abs(w).sum()
        if den > 0:
            w_full[lo:hi] = w / den
    return w_full


def book_fallback(data, mu):
    """Last resort so a month is never empty: rank-proportional, dollar-neutral, unit gross."""
    s = pd.Series(mu, index=data["eom"])
    r = s.groupby(level=0).rank(pct=True) - 0.5
    w = np.nan_to_num(r.to_numpy(), nan=0.0)
    gross = pd.Series(np.abs(w), index=data["eom"]).groupby(level=0).transform("sum").to_numpy()
    return np.divide(w, gross, out=np.zeros_like(w), where=gross > 0).astype(np.float32)


# ===========================================================================================
# 4. Entry point
# ===========================================================================================
def main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    random.seed(SEED)
    np.random.seed(SEED)

    data = prepare(chars, features)
    print(f"panel {data['X'].shape}, {len(data['months'])} months, "
          f"{len(data['feat_list'])} features", flush=True)

    mu_raw = walk_forward(data)
    mu = standardize(mu_raw, data["eom"])
    held = np.isfinite(mu_raw)
    print(f"predictions for {held.sum():,} rows", flush=True)

    dates, daily_ids, M = build_daily_matrix(daily_ret, np.unique(data["ids"]))
    print(f"daily matrix {M.shape[0]:,}d x {M.shape[1]:,} ids", flush=True)

    books = {
        "dc_long": book_daily_cov(data, dates, daily_ids, M, mu, WINDOW_LONG),
        "mvo": book_factor_mvo(data, mu),
    }
    for k, v in books.items():
        print(f"  book {k:9s} nonzero rows {int((v != 0).sum()):,}", flush=True)

    w = np.zeros(len(mu), np.float64)
    for name, alpha in BOOK_WEIGHTS.items():
        w += alpha * books[name]
    w[~held] = 0.0

    # renormalize each month to unit gross; fall back where every book was unavailable
    eom = data["eom"]
    gross = pd.Series(np.abs(w), index=eom).groupby(level=0).transform("sum").to_numpy()
    w = np.divide(w, gross, out=np.zeros_like(w), where=gross > 0)
    empty = pd.Series(gross <= 0, index=eom).groupby(level=0).transform("all").to_numpy()
    if empty.any():
        fb = book_fallback(data, np.where(held, mu, 0.0))
        w = np.where(empty & held, fb, w)

    # ---- volatility targeting overlay ----
    # The book formed at month m earns its return over m+1, so at month t only months <= t-1 are
    # observable: hence .shift(1) on the realised series. The target is an expanding mean of past
    # trailing vols, so nothing here uses a statistic from the future.
    # Only months in which the book holds positions enter the volatility estimates. The panel
    # begins decades before the first forecast month, and those empty months have zero return and
    # zero volatility; including them drags the expanding target toward zero and pins the
    # multiplier at its lower bound, which disables the overlay entirely.
    gross_m = pd.Series(np.abs(w), index=eom).groupby(level=0).sum()
    r_hist = pd.Series(w * np.nan_to_num(data["ret"]), index=eom).groupby(level=0).sum()
    live = r_hist[gross_m > 0]
    tv = live.rolling(VOL_WINDOW).std().shift(1)
    lever = (tv.expanding().mean() / tv).clip(*VOL_CLIP).fillna(1.0)
    lever = lever.reindex(r_hist.index).fillna(1.0)
    w = w * lever.reindex(pd.DatetimeIndex(eom)).to_numpy()
    print(f"vol-target leverage: mean {lever.mean():.2f}, "
          f"range {lever.min():.2f}-{lever.max():.2f}", flush=True)

    out = pd.DataFrame({ID_COL: data["ids"], DATE_COL: pd.to_datetime(eom), "w": w})
    if data["test"].any():
        out = out[data["test"]]
    out = out.dropna(subset=[ID_COL, DATE_COL])
    out["w"] = out["w"].fillna(0.0).replace([np.inf, -np.inf], 0.0)
    out[DATE_COL] = out[DATE_COL].dt.strftime("%Y-%m-%d")
    print(f"returning {len(out):,} rows, gross/month "
          f"{out.groupby(DATE_COL)['w'].apply(lambda x: x.abs().sum()).mean():.3f}", flush=True)
    return out.reset_index(drop=True)