"""
Yale SOM - MGT 924 Statistical Foundations (Prof. Theis Ingerslev Jensen)
Problem Set 3: JKP Common Task Framework (CTF) Competition Model

Champion Architecture: Multi-Scale Gradient-Boosted Rank Ensemble with 
Barra 13-Theme Structural Factor Risk MVO & Market Beta Neutralization

Key Methodology:
1. Target Formulation: Cross-sectional rank transformation of 1-month forward excess returns
   y_{i,t} = rank(r_{i,t+1}) - 0.5. Bounded targets eliminate outlier sensitivity and make
   mean-variance optimization well-conditioned.
2. Feature Representation: 153 published JKP themed characteristics (Jensen, Kelly, Pedersen 2023)
   cross-sectionally percentile-ranked to [-0.5, 0.5] + one-hot SIC industry classifications.
3. Multi-Scale Alpha Engine: An ensemble of complementary LightGBM regressors:
   - Deep Interaction Model (num_leaves=63, unconstrained depth) capturing high-order non-linearities.
   - Shallow Main-Effects Model (num_leaves=15, max_depth=4) regularizing against factor overfitting.
   - 10-year exponential time decay weighting to adapt to structural regime shifts.
4. Structural Risk Model: Barra-style 13-theme factor risk model Sigma = B Omega B' + D:
   - B: Exposures to the 13 JKP sign-aligned economic themes.
   - Omega: 60-month trailing factor covariance with Ledoit-Wolf shrinkage toward diagonal.
   - D: Idiosyncratic residual variance with median fallback and 10th percentile flooring.
   - Inversion via Woodbury identity in O(N * K^2) time without forming the N x N matrix.
5. Risk Neutralization & Portfolio Sizing:
   - Market beta neutralization: Orthogonalizes weights against trailing 60-month market beta (beta_60m).
   - Strict dollar neutrality (sum w_i = 0) and gross leverage normalization (||w||_1 = 2.0).

CTF Admin Modifications (2026-09-05):
--------------------------------------
1. Added a requirements.txt containing lightgbm==4.7.0.
   Reason: the submission had no dependency file, and lightgbm is not among the
   packages pre-installed in the submission runtime (which provides only pandas,
   numpy, scipy, scikit-learn, pyarrow, boto3, polars and joblib). Without a
   dependency file the image build takes its hardcoded-fallback path, so lightgbm
   would never be installed and the run would fail at the import inside main().
   4.7.0 has no outstanding advisories and ships a Linux wheel compatible with the
   Python 3.13 runtime.

2. Commented out the local prediction cache lookup and dedented the walk-forward
   training block that followed it.
   Reason: the cache pointed at an absolute path on a developer machine, which
   cannot exist in the competition environment, so the branch could never have
   loaded and the training path is what always ran. It was removed because it made
   model behaviour conditional on filesystem state and embedded a personal
   filesystem path in a submission that is published for download. No behaviour
   changed.

3. Replaced two `float(...)` conversions in `_neutralize_vector` with `.item()`.
   Reason: both operands are matrix products of shape (1, 1), and NumPy 2.x refuses
   to convert any array with more than zero dimensions to a Python scalar. The
   submitted code raised "only 0-dimensional arrays can be converted to Python
   scalars" on the first month of beta neutralisation and the pipeline failed at the
   validation run. Our runtime installs numpy 2.5.2 as a dependency of lightgbm,
   whereas the same code succeeds under numpy 1.x where this was only deprecated.
   `.item()` extracts the same scalar, so the arithmetic is unchanged.

After these changes: no follow-up required. This model is CPU only and does not
need a GPU allocation.
"""
from __future__ import annotations

import os
import random
import warnings
import numpy as np
import pandas as pd

# Suppress cosmetic LightGBM warnings
warnings.filterwarnings("ignore")

# Determinism
SEED = 42

ID_COL = "id"
DATE_COL = "eom"
RET_COL = "ret_exc_lead1m"
TEST_COL = "ctff_test"
BETA_COL = "beta_60m"

# Walk-forward parameters
MIN_TRAIN_MONTHS = 120
VAL_MONTHS = 24
RETRAIN_MONTHS = 12
HALFLIFE_YEARS = 10.0
MVO_WINDOW = 60
MVO_D_FLOOR_Q = 0.10
OMEGA_SHRINKAGE = 0.20
GROSS = 2.0

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

# 153 Themed Characteristics -> (Theme, Direction Sign) from Jensen, Kelly, Pedersen (2023)
THEME_MAP = {
    'age': ('Low Leverage', -1), 'aliq_at': ('Investment', -1), 'aliq_mat': ('Low Leverage', -1),
    'ami_126d': ('Size', 1), 'at_be': ('Low Leverage', -1), 'at_gr1': ('Investment', -1),
    'at_me': ('Value', 1), 'at_turnover': ('Quality', 1), 'be_gr1a': ('Investment', -1),
    'be_me': ('Value', 1), 'beta_60m': ('Low Risk', -1), 'beta_dimson_21d': ('Low Risk', -1),
    'betabab_1260d': ('Low Risk', -1), 'betadown_252d': ('Low Risk', -1), 'bev_mev': ('Value', 1),
    'bidaskhl_21d': ('Low Leverage', 1), 'capex_abn': ('Debt Issuance', -1), 'capx_gr1': ('Investment', -1),
    'capx_gr2': ('Investment', -1), 'capx_gr3': ('Investment', -1), 'cash_at': ('Low Leverage', 1),
    'chcsho_12m': ('Value', -1), 'coa_gr1a': ('Investment', -1), 'col_gr1a': ('Investment', -1),
    'cop_at': ('Quality', 1), 'cop_atl1': ('Quality', 1), 'corr_1260d': ('Seasonality', -1),
    'coskew_21d': ('Seasonality', -1), 'cowc_gr1a': ('Accruals', -1), 'dbnetis_at': ('Seasonality', -1),
    'debt_gr3': ('Debt Issuance', -1), 'debt_me': ('Value', 1), 'dgp_dsale': ('Quality', 1),
    'div12m_me': ('Value', 1), 'dolvol_126d': ('Size', -1), 'dolvol_var_126d': ('Profitability', -1),
    'dsale_dinv': ('Profit Growth', 1), 'dsale_drec': ('Profit Growth', -1), 'dsale_dsga': ('Profit Growth', 1),
    'earnings_variability': ('Low Risk', -1), 'ebit_bev': ('Profitability', 1), 'ebit_sale': ('Profitability', 1),
    'ebitda_mev': ('Value', 1), 'emp_gr1': ('Investment', -1), 'eq_dur': ('Value', -1),
    'eqnetis_at': ('Value', -1), 'eqnpo_12m': ('Value', 1), 'eqnpo_me': ('Value', 1),
    'eqpo_me': ('Value', 1), 'f_score': ('Profitability', 1), 'fcf_me': ('Value', 1),
    'fnl_gr1a': ('Debt Issuance', -1), 'gp_at': ('Quality', 1), 'gp_atl1': ('Quality', 1),
    'inv_gr1': ('Investment', -1), 'inv_gr1a': ('Investment', -1), 'iskew_capm_21d': ('Short-Term Reversal', -1),
    'iskew_ff3_21d': ('Short-Term Reversal', -1), 'iskew_hxz4_21d': ('Short-Term Reversal', -1), 'ival_me': ('Value', 1),
    'ivol_capm_21d': ('Low Risk', -1), 'ivol_capm_252d': ('Low Risk', -1), 'ivol_ff3_21d': ('Low Risk', -1),
    'ivol_hxz4_21d': ('Low Risk', -1), 'kz_index': ('Seasonality', 1), 'lnoa_gr1a': ('Investment', -1),
    'lti_gr1a': ('Seasonality', -1), 'market_equity': ('Size', -1), 'mispricing_mgmt': ('Investment', 1),
    'mispricing_perf': ('Quality', 1), 'ncoa_gr1a': ('Investment', -1), 'ncol_gr1a': ('Debt Issuance', -1),
    'netdebt_me': ('Low Leverage', -1), 'netis_at': ('Value', -1), 'nfna_gr1a': ('Debt Issuance', 1),
    'ni_ar1': ('Debt Issuance', 1), 'ni_be': ('Profitability', 1), 'ni_inc8q': ('Quality', 1),
    'ni_ivol': ('Low Leverage', 1), 'ni_me': ('Value', 1), 'niq_at': ('Quality', 1),
    'niq_at_chg1': ('Profit Growth', 1), 'niq_be': ('Profitability', 1), 'niq_be_chg1': ('Profit Growth', 1),
    'niq_su': ('Profit Growth', 1), 'nncoa_gr1a': ('Investment', -1), 'noa_at': ('Debt Issuance', -1),
    'noa_gr1a': ('Investment', -1), 'o_score': ('Profitability', -1), 'oaccruals_at': ('Accruals', -1),
    'oaccruals_ni': ('Accruals', -1), 'ocf_at': ('Profitability', 1), 'ocf_at_chg1': ('Profit Growth', 1),
    'ocf_me': ('Value', 1), 'ocfq_saleq_std': ('Low Risk', -1), 'op_at': ('Quality', 1),
    'op_atl1': ('Quality', 1), 'ope_be': ('Profitability', 1), 'ope_bel1': ('Profitability', 1),
    'opex_at': ('Quality', 1), 'pi_nix': ('Seasonality', 1), 'ppeinv_gr1a': ('Investment', -1),
    'prc': ('Size', -1), 'prc_highprc_252d': ('Momentum', 1), 'qmj': ('Quality', 1),
    'qmj_growth': ('Quality', 1), 'qmj_prof': ('Quality', 1), 'qmj_safety': ('Quality', 1),
    'rd_me': ('Size', 1), 'rd_sale': ('Low Leverage', 1), 'rd5_at': ('Low Leverage', 1),
    'resff3_12_1': ('Momentum', 1), 'resff3_6_1': ('Momentum', 1), 'ret_1_0': ('Short-Term Reversal', -1),
    'ret_12_1': ('Momentum', 1), 'ret_12_7': ('Profit Growth', 1), 'ret_3_1': ('Momentum', 1),
    'ret_6_1': ('Momentum', 1), 'ret_60_12': ('Investment', -1), 'ret_9_1': ('Momentum', 1),
    'rmax1_21d': ('Low Risk', -1), 'rmax5_21d': ('Low Risk', -1), 'rmax5_rvol_21d': ('Short-Term Reversal', -1),
    'rskew_21d': ('Short-Term Reversal', -1), 'rvol_21d': ('Low Risk', -1), 'sale_bev': ('Quality', 1),
    'sale_emp_gr1': ('Profit Growth', 1), 'sale_gr1': ('Investment', -1), 'sale_gr3': ('Investment', -1),
    'sale_me': ('Value', 1), 'saleq_gr1': ('Investment', -1), 'saleq_su': ('Profit Growth', 1),
    'seas_1_1an': ('Profit Growth', 1), 'seas_1_1na': ('Momentum', 1), 'seas_11_15an': ('Seasonality', 1),
    'seas_11_15na': ('Seasonality', -1), 'seas_16_20an': ('Seasonality', 1), 'seas_16_20na': ('Accruals', -1),
    'seas_2_5an': ('Seasonality', 1), 'seas_2_5na': ('Investment', -1), 'seas_6_10an': ('Seasonality', 1),
    'seas_6_10na': ('Low Risk', -1), 'sti_gr1a': ('Seasonality', 1), 'taccruals_at': ('Accruals', -1),
    'taccruals_ni': ('Accruals', -1), 'tangibility': ('Low Leverage', 1), 'tax_gr1a': ('Profit Growth', 1),
    'turnover_126d': ('Low Risk', -1), 'turnover_var_126d': ('Profitability', -1), 'z_score': ('Low Leverage', 1),
    'zero_trades_126d': ('Low Risk', 1), 'zero_trades_21d': ('Low Risk', 1), 'zero_trades_252d': ('Low Risk', 1),
}


def _woodbury_sigma_inv_mu(B: np.ndarray, Omega: np.ndarray, d: np.ndarray, mu: np.ndarray) -> np.ndarray:
    """Solve (B Omega B^T + diag(d))^-1 mu using the Woodbury matrix inversion lemma.
    Time complexity O(N * K^2 + K^3), eliminating N x N matrix inversion.
    """
    K = B.shape[1]
    dinv = 1.0 / d
    BtDinv = B.T * dinv
    Omega_inv = np.linalg.inv(Omega + 1e-8 * np.eye(K))
    M = Omega_inv + BtDinv @ B
    inner = np.linalg.solve(M, BtDinv @ mu)
    return dinv * mu - dinv * (B @ inner)


def _neutralize_vector(w: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Orthogonalize vector w against vector b (removes component along b)."""
    if b is None or len(b) == 0:
        return w
    b_col = b.reshape(-1, 1)
    denom = float((b_col.T @ b_col).item())
    if denom < 1e-12:
        return w
    proj = ((b_col.T @ w.reshape(-1, 1)).item() / denom) * b_col.ravel()
    return w - proj


def main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    """Main competition function according to JKP CTF specifications.

    Args:
        chars (pd.DataFrame): Monthly characteristics dataset.
        features (pd.DataFrame): Characteristic feature names.
        daily_ret (pd.DataFrame): Daily returns dataset.

    Returns:
        pd.DataFrame: Portfolio weights with columns ['id', 'eom', 'w'].
    """
    # Determinism
    random.seed(SEED)
    np.random.seed(SEED)

    # Feature selection
    if "features" in features.columns:
        feat_names = features["features"].astype(str).tolist()
    else:
        feat_names = features[features.columns[0]].astype(str).tolist()
    feat_set = set(feat_names)

    chars = chars.copy()
    chars[DATE_COL] = pd.to_datetime(chars[DATE_COL])
    chars = chars.sort_values([DATE_COL, ID_COL]).reset_index(drop=True)

    char_cols = [c for c in THEME_MAP if c in feat_set and c in chars.columns]
    print(f"[Champion Model] Data rows: {len(chars):,}, Themed features: {len(char_cols)}/153", flush=True)

    # Cross-sectional percentile ranking [-0.5, 0.5]
    ranks = chars.groupby(DATE_COL)[char_cols].rank(pct=True) - 0.5
    ranks = ranks.fillna(0.0).astype(np.float32)

    # Coarse industry one-hot encoding from SIC
    sic_num = pd.to_numeric(chars["sic"], errors="coerce").fillna(-1) // 100
    ind_div = np.where(sic_num >= 0, sic_num // 10, -1)
    ind_dummies = pd.get_dummies(pd.Series(ind_div, index=chars.index), prefix="ind").astype(np.float32)

    X = pd.concat([ranks, ind_dummies], axis=1)

    # Target: within-month centered percentile rank of next-month excess return
    r = pd.to_numeric(chars[RET_COL], errors="coerce")
    y = (r.groupby(chars[DATE_COL]).rank(pct=True) - 0.5).values.astype(np.float64)

    # Compute the 13 JKP theme exposures
    theme_dfs = {}
    for theme in THEMES:
        members = [c for c in char_cols if THEME_MAP[c][0] == theme]
        if not members:
            continue
        signs = [THEME_MAP[c][1] for c in members]
        theme_dfs[theme] = ranks[members].mul(signs, axis=1).mean(axis=1)
    themes = pd.DataFrame(theme_dfs, index=chars.index).fillna(0.0)
    theme_cols = list(themes.columns)

    # Cross-sectional factor returns regression
    K = len(theme_cols)
    reg = 1e-4 * np.eye(K)
    Xall = themes.values.astype(np.float64)
    rall = r.values.astype(np.float64)
    f_by_eom = {}
    resid = np.full(len(chars), np.nan)

    for eom, rowpos in chars.groupby(DATE_COL, sort=True).indices.items():
        Xt, rt = Xall[rowpos], rall[rowpos]
        ok = np.isfinite(rt) & np.isfinite(Xt).all(axis=1)
        if ok.sum() <= K:
            continue
        Xo, ro = Xt[ok], rt[ok]
        f = np.linalg.solve(Xo.T @ Xo + reg, Xo.T @ ro)
        f_by_eom[eom] = f
        resid[rowpos[ok]] = ro - Xo @ f

    F_df = pd.DataFrame.from_dict(f_by_eom, orient="index", columns=theme_cols).sort_index()
    chars["resid"] = resid

    # Check test mask
    if TEST_COL in chars.columns:
        test_mask = chars[TEST_COL].astype(str).str.strip().isin(["1", "True", "true"])
    else:
        test_mask = pd.Series(True, index=chars.index)

    all_months = np.array(sorted(chars[DATE_COL].unique()))
    test_months = np.array(sorted(chars.loc[test_mask, DATE_COL].unique()))
    first_test_m = test_months.min() if len(test_months) else all_months[MIN_TRAIN_MONTHS]

    # Time-indexed sorted arrays for zero-copy lookups
    order = np.argsort(chars[DATE_COL].values, kind="stable")
    eom_s = chars[DATE_COL].values[order]
    ids = chars[ID_COL].values[order]
    Xv = X.values.astype(np.float32)[order]
    yv = y[order]
    retv = rall[order]

    eom_s_ns = eom_s.astype("datetime64[ns]")
    all_months_ns = all_months.astype("datetime64[ns]")
    starts = np.searchsorted(eom_s_ns, all_months_ns, side="left")
    ends = np.searchsorted(eom_s_ns, all_months_ns, side="right")
    slc = {m: (int(a), int(b)) for m, a, b in zip(all_months, starts, ends)}

    # Import lightgbm inside main to adhere to modularity
    import lightgbm as lgb

    p_deep = {
        "n_estimators": 250, "learning_rate": 0.04, "num_leaves": 63,
        "max_depth": -1, "min_child_samples": 100, "subsample": 0.8,
        "colsample_bytree": 0.8, "reg_lambda": 1.0, "n_jobs": -1,
        "force_col_wise": True
    }
    p_shallow = {
        "n_estimators": 250, "learning_rate": 0.04, "num_leaves": 15,
        "max_depth": 4, "min_child_samples": 250, "subsample": 0.8,
        "colsample_bytree": 0.8, "reg_lambda": 5.0, "n_jobs": -1,
        "force_col_wise": True
    }

    # CTF admin: the local prediction cache below was disabled. The path is a
    # developer machine path that cannot exist in the competition environment, so
    # this branch could never have loaded and the walk-forward training below is
    # what always ran. It is commented out rather than left in place because it
    # made behaviour conditional on filesystem state and embedded a personal path
    # in a submission that is published for download.
    # # Check local cache for fast experimentation
    # cache_path = "<developer machine path>/preds_cache.parquet"  # path redacted by CTF admin
    # if os.path.exists(cache_path):
    # print(f"[Champion Model] Loading cached walk-forward predictions from {cache_path}...", flush=True)
    # preds = pd.read_parquet(cache_path)
    # else:
    print("[Champion Model] Executing point-in-time walk-forward training...", flush=True)
    m1_prev, m2_prev = None, None
    preds_rows = []

    for ti, t in enumerate(all_months):
        if t < first_test_m:
            continue

        avail = all_months[:ti]
        if len(avail) < MIN_TRAIN_MONTHS:
            continue

        n_test_done = np.searchsorted(test_months, t)
        needs_retrain = (m1_prev is None) or (n_test_done % RETRAIN_MONTHS == 0)

        if needs_retrain:
            core = avail[:-VAL_MONTHS]
            vals = avail[-VAL_MONTHS:]
            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 = Xv[c_lo:c_hi], yv[c_lo:c_hi]
            Xva, yva = Xv[v_lo:v_hi], yv[v_lo:v_hi]
            m_tr, m_va = np.isfinite(yc), np.isfinite(yva)
            Xc, yc = Xc[m_tr], yc[m_tr]
            Xva, yva = Xva[m_va], yva[m_va]

            # Exponential time-decay sample weighting
            tr_months = eom_s[c_lo:c_hi][m_tr]
            yrs = pd.DatetimeIndex(tr_months).year.values.astype(np.float64)
            sw = 0.5 ** ((yrs.max() - yrs) / HALFLIFE_YEARS)
            sw = sw / sw.mean()

            # Fit Deep Model
            m1 = lgb.LGBMRegressor(
                objective="regression", verbose=-1, random_state=42,
                deterministic=True, **p_deep
            )
            m1.fit(Xc, yc, sample_weight=sw, eval_set=[(Xva, yva)],
                   callbacks=[lgb.early_stopping(30, verbose=False)])
            m1_prev = m1

            # Fit Shallow Regularized Model
            m2 = lgb.LGBMRegressor(
                objective="regression", verbose=-1, random_state=142,
                deterministic=True, **p_shallow
            )
            m2.fit(Xc, yc, sample_weight=sw, eval_set=[(Xva, yva)],
                   callbacks=[lgb.early_stopping(30, verbose=False)])
            m2_prev = m2

        p_lo, p_hi = slc[t]
        X_test = Xv[p_lo:p_hi]

        y_p1 = m1_prev.predict(X_test).astype(np.float64)
        y_p2 = m2_prev.predict(X_test).astype(np.float64)
        y_pred = 0.60 * y_p1 + 0.40 * y_p2

        preds_rows.append(pd.DataFrame({
            ID_COL: ids[p_lo:p_hi],
            DATE_COL: t,
            "y_pred": y_pred,
            RET_COL: retv[p_lo:p_hi]
        }))

    preds = pd.concat(preds_rows, ignore_index=True)

    # -----------------------------------------------------------------------
    # Portfolio Construction: Barra Structural MVO with Factor Shrinkage & Beta Neutralization
    # -----------------------------------------------------------------------
    print("[Champion Model] Constructing Barra MVO portfolios...", flush=True)
    chars_map = pd.concat([chars[[ID_COL, DATE_COL, "resid", BETA_COL]], themes], axis=1).copy()
    all_f_months = list(F_df.index)
    f_pos = {m: i for i, m in enumerate(all_f_months)}

    w_out = []

    for t, sub in preds.groupby(DATE_COL):
        if t not in f_pos:
            continue
        win = all_f_months[max(0, f_pos[t] - MVO_WINDOW):f_pos[t]]
        if len(win) < 15:
            continue

        # Factor covariance matrix with Ledoit-Wolf diagonal shrinkage
        sample_cov = np.cov(F_df.loc[win].values, rowvar=False)
        if OMEGA_SHRINKAGE > 0:
            target = np.diag(np.diag(sample_cov))
            Omega = (1.0 - OMEGA_SHRINKAGE) * sample_cov + OMEGA_SHRINKAGE * target
        else:
            Omega = sample_cov

        cur = chars_map[chars_map[DATE_COL] == t]
        hist_resids = chars_map[chars_map[DATE_COL].isin(win)]
        rv = hist_resids.groupby(ID_COL)["resid"].var()

        joined = sub.merge(cur, on=[ID_COL, DATE_COL], how="inner").dropna(subset=theme_cols)
        if joined.empty:
            continue

        B = joined[theme_cols].values.astype(np.float64)
        mu = joined["y_pred"].values.astype(np.float64)

        # Idiosyncratic residual variance with median fallback and flooring
        d = joined[ID_COL].map(rv).values.astype(np.float64)
        finite_d = d[np.isfinite(d)]
        med_d = np.nanmedian(finite_d) if len(finite_d) else 0.01
        q10_d = np.nanquantile(finite_d, MVO_D_FLOOR_Q) if len(finite_d) else 0.001

        d = np.where(np.isnan(d), med_d, d)
        d = np.maximum(d, q10_d)

        # Closed-form Woodbury inversion
        w_raw = _woodbury_sigma_inv_mu(B, Omega, d, mu)

        # Dollar Neutrality
        w_raw = w_raw - w_raw.mean()

        # Market Beta Neutralization
        if BETA_COL in joined.columns:
            betas = pd.to_numeric(joined[BETA_COL], errors="coerce").fillna(1.0).values.astype(np.float64)
            w_raw = _neutralize_vector(w_raw, betas)
            w_raw = w_raw - w_raw.mean()

        # Normalization to target gross exposure
        denom = np.abs(w_raw).sum()
        w_final = GROSS * w_raw / denom if denom > 1e-10 else np.zeros_like(w_raw)

        w_out.append(pd.DataFrame({
            ID_COL: joined[ID_COL].values,
            DATE_COL: t,
            "w": w_final
        }))

    weights_df = pd.concat(w_out, ignore_index=True)

    # Filter strictly to ctff_test == True rows
    if TEST_COL in chars.columns:
        test_keys = chars.loc[test_mask, [ID_COL, DATE_COL]]
        out = weights_df.merge(test_keys, on=[ID_COL, DATE_COL], how="inner")
    else:
        out = weights_df

    # Strictly enforce types and columns
    out[ID_COL] = out[ID_COL].astype("int64")
    out[DATE_COL] = pd.to_datetime(out[DATE_COL]).dt.strftime("%Y-%m-%d")
    out["w"] = pd.to_numeric(out["w"], errors="coerce").fillna(0.0).astype("float64")
    out = out[[ID_COL, DATE_COL, "w"]].dropna().sort_values([DATE_COL, ID_COL]).reset_index(drop=True)

    print(f"[Champion Model] Returning {len(out):,} rows across {out[DATE_COL].nunique()} months.", flush=True)
    return out


if __name__ == "__main__":
    import time
    from pathlib import Path

    DATA_DIR = Path(__file__).resolve().parent / "Data"
    chars_p = DATA_DIR / "ctff_chars.parquet"
    feats_p = DATA_DIR / "ctff_features.parquet"
    daily_p = DATA_DIR / "ctff_daily_ret.parquet"

    print("Running end-to-end model verification...")
    t0 = time.time()
    chars_df = pd.read_parquet(chars_p)
    feats_df = pd.read_parquet(feats_p)
    daily_df = pd.DataFrame()  # Daily return not required for structural factor MVO

    weights = main(chars_df, feats_df, daily_df)
    
    # Save output
    output_path = Path(__file__).resolve().parent / "portfolio_weights_submission.csv"
    weights.to_csv(output_path, index=False)
    print(f"Saved submission weights to {output_path} ({len(weights):,} rows) in {time.time()-t0:.1f}s")