"""Decay-Aware Cluster-Coherent Alpha & Risk (CCAR) -- CTF submission for MGT 924 Problem Set 3.

One hierarchical clustering of the characteristics ("themes") drives both halves of the strategy:
  * alpha: an ensemble of three learners (ridge, random-feature ridge, gradient-boosted trees) is
    split into the part spanned by the theme exposures and a stock-specific residual; the theme
    part is blended with decay-aware theme premia, estimated under several memory lengths and
    weighted by their own walk-forward track records;
  * risk:  a Barra-style factor model whose factors are a common factor plus the same themes,
    estimated from daily returns, used to turn forecasts into mean-variance weights.

The themes, learners and memory weights are re-estimated once a year on expanding past data; theme
premia and risk estimates update using only information available by the portfolio date.
Entry point: main(chars, features, daily_ret).
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
import polars as pl
from scipy.cluster.hierarchy import fcluster, linkage
from scipy.spatial.distance import squareform
from scipy.stats import norm
from sklearn.ensemble import HistGradientBoostingRegressor

# ---------------------------------------------------------------------------------------------
# Configuration. Settings were fixed before test evaluation; gamma was selected on the designated
# 1975-1989 validation period. Ridge penalties, themes and memory weights are re-estimated inside
# each refit using past data only.
# ---------------------------------------------------------------------------------------------
CFG = dict(
    seed=20251006,
    min_train_months=120,      # burn-in before the first portfolio
    refit_month=12,            # refit every December (and at the first portfolio date)
    val_months=60,             # internal validation window for ridge penalties
    ridge_grid=np.logspace(-4, 4, 17),
    rff_pairs=1000,            # random Fourier features -> 2 * 1000 regressors
    rff_bandwidth=1.0,         # target std of the random projections
    rff_max_rows=400_000,
    gbr=dict(max_iter=300, learning_rate=0.05, max_leaf_nodes=31, min_samples_leaf=2000,
             l2_regularization=1.0, max_features=0.5),
    gbr_max_rows=600_000,
    cluster_cut=0.5,           # cut the 1-|rho| dendrogram at this height
    cluster_max=25,
    cluster_months=120,
    risk_months=24,            # daily history used by the risk model
    hl_factor=252,             # EWMA half-lives in trading days
    hl_idio=63,
    idio_prior_obs=20,         # shrinkage of idiosyncratic variance toward the cross-sectional median
    min_stocks_day=50,
    pos_cap_mult=10.0,         # |w_i| <= 10 x average |w|
    # decay-aware theme premia
    gamma=1.0,                 # weight on decay-aware theme premia; chosen on 1975-1989 only (tools/select_gamma.py), frozen
    gamma_grid=[0.0, 0.25, 0.5, 0.75, 1.0],
    memories=[1, 3, 10, 30, None],   # half-lives in years for theme-premium estimates (None = full history, equal weights)
    premia_prior_months=60,    # premia shrunk toward zero: mean * n_eff / (n_eff + prior)
    ledger_months=120,         # track record used each December to weight the memories
    proj_ridge=1e-8,           # numerical ridge in the theme projection, relative to tr(E'E)/K
    target_vol=0.10,           # ex-ante annualised volatility target (no leverage cap: the leaderboard
                               # evaluates returns rescaled to 10% volatility)
)

# Column-name assumptions (see report, Section 1). The first candidate that exists is used.
COLS = dict(
    id=["id"], eom=["eom"], target=["ret_exc_lead1m"], test=["ctff_test"],
    me=["me", "market_equity"], size=["size_grp"], feature=["features"],
    date=["date"], dret=["ret_exc", "ret"],
)

VARIANTS = ["final", "ccar0", "ens_rank", "fixed_rank", "decay_rank", "ridge_rank", "rff_rank", "gbr_rank",
            "ridge_decile", "ew_market"]
GAMMA_VARIANTS = [f"g{int(100 * g):03d}" for g in CFG["gamma_grid"]]   # full model at each gamma (validation only)


def _pick(columns, key, required=True):
    for c in COLS[key]:
        if c in columns:
            return c
    if required:
        raise KeyError(f"none of {COLS[key]} found in data")
    return None


# ---------------------------------------------------------------------------------------------
# Data preparation
# ---------------------------------------------------------------------------------------------
class _Source:
    """Uniform column access to a pandas/polars frame, a polars LazyFrame or a parquet path.
    Columns are pulled a block at a time, so the full table never has to sit in memory twice."""

    def __init__(self, obj):
        if isinstance(obj, (str, Path)):
            obj = pl.scan_parquet(obj)
        self.obj = obj
        if isinstance(obj, pl.LazyFrame):
            self.columns = obj.collect_schema().names()
        else:
            self.columns = list(obj.columns)

    def get(self, cols) -> pl.DataFrame:
        if isinstance(self.obj, pd.DataFrame):
            return pl.from_pandas(self.obj[cols])
        if isinstance(self.obj, pl.LazyFrame):
            return self.obj.select(cols).collect()
        return self.obj.select(cols)


class Panel:
    """Characteristics sorted by (eom, id) with month offsets, ranked features and targets."""

    def __init__(self, chars, features, batch=40):
        src = _Source(chars)
        c_id, c_eom = _pick(src.columns, "id"), _pick(src.columns, "eom")
        c_y = _pick(src.columns, "target")
        c_test, c_me = _pick(src.columns, "test", False), _pick(src.columns, "me", False)
        c_size = _pick(src.columns, "size", False)
        fsrc = _Source(features)
        names = fsrc.get([_pick(fsrc.columns, "feature")]).to_series().cast(pl.Utf8).to_list()
        feats = [f for f in names if f in src.columns and f not in {c_id, c_eom, c_y, c_test}]

        base = src.get(list(dict.fromkeys([c_id, c_eom, c_y] + [c for c in (c_test, c_me, c_size) if c])))
        base = base.with_columns(pl.col(c_eom).cast(pl.Date), pl.col(c_id).cast(pl.Int64)).with_row_index("_row")
        base = base.sort([c_eom, c_id])
        order = pl.Series(base["_row"].to_numpy())
        eom_key = base[c_eom]

        def ranked(f):
            x = pl.col(f).cast(pl.Float64, strict=False).fill_nan(None)
            n = x.count().over("_eom")
            r = (x.rank("average").over("_eom") - 1) / pl.when(n > 1).then(n - 1).otherwise(1) - 0.5
            return pl.when(n > 1).then(r).otherwise(0.0).fill_null(0.0).cast(pl.Float32).alias(f)

        self.X = np.empty((len(base), len(feats)), dtype=np.float32)
        for b in range(0, len(feats), batch):
            fb = feats[b:b + batch]
            blk = src.get(fb).select(pl.all().gather(order)).with_columns(eom_key.alias("_eom"))
            self.X[:, b:b + len(fb)] = blk.select([ranked(f) for f in fb]).to_numpy()
            del blk

        y = pl.col(c_y).cast(pl.Float64).fill_nan(None)
        ny = y.count().over(c_eom)
        # target: within-month rank of next-month excess return, mapped to normal scores
        yq = base.select(((y.rank("average").over(c_eom) - 0.5) / ny).alias("_yq"))["_yq"].to_numpy().astype(float)

        self.features = feats
        self.ids = base[c_id].to_numpy()
        eom = base[c_eom].to_numpy().astype("datetime64[D]")
        self.months, self.month_idx = np.unique(eom, return_inverse=True)
        self.offsets = np.searchsorted(self.month_idx, np.arange(len(self.months) + 1))
        self.y_rank = np.where(np.isnan(yq), np.nan, norm.ppf(np.clip(yq, 1e-6, 1 - 1e-6)))
        self.y_raw = base[c_y].cast(pl.Float64).fill_nan(None).to_numpy()
        self.me = base[c_me].cast(pl.Float64, strict=False).to_numpy() if c_me else None
        if c_test:  # ctff_test may arrive as booleans, 0/1 or '0'/'1'
            self.test = base[c_test].cast(pl.Utf8).str.strip_chars().str.to_lowercase().is_in(["1", "true", "t"]).to_numpy()
        else:
            self.test = None
        self.size = base[c_size].cast(pl.Utf8).to_numpy() if c_size else None
        self.id_col, self.eom_col = c_id, c_eom

    def rows(self, m0, m1):
        """Row slice for months [m0, m1)."""
        return slice(self.offsets[m0], self.offsets[m1])


class Daily:
    """Daily excess returns keyed by formation month: day d belongs to the last month-end < d."""

    def __init__(self, daily_ret, months: np.ndarray):
        src = _Source(daily_ret)
        c_id, c_d = _pick(src.columns, "id"), _pick(src.columns, "date")
        c_r = _pick(src.columns, "dret")
        d = src.get([c_id, c_d, c_r])
        # compact types: 103m rows; float32 returns are ample for a risk model
        d = d.select(pl.col(c_id).cast(pl.Int32).alias("id"), pl.col(c_d).cast(pl.Date).alias("date"),
                     pl.col(c_r).cast(pl.Float64).fill_nan(None).clip(-0.5, 1.0).cast(pl.Float32).alias("ret")).drop_nulls()
        mser = pl.Series("m", months.astype("datetime64[D]")).cast(pl.Date)
        d = d.with_columns((mser.search_sorted(d["date"], side="left").cast(pl.Int32) - 1).alias("form"))
        d = d.filter(pl.col("form") >= 0).sort(["form", "id", "date"])
        self.form = d["form"].to_numpy()
        self.ids = d["id"].to_numpy().astype(np.int64)
        self.day = d["date"].cast(pl.Int32).to_numpy()
        self.ret = d["ret"].to_numpy()
        del d
        self.offsets = np.searchsorted(self.form, np.arange(len(months) + 1))


# ---------------------------------------------------------------------------------------------
# Themes: hierarchical clustering of characteristics (as in "Is There a Replication Crisis in
# Finance"), estimated on training data only.
# ---------------------------------------------------------------------------------------------
def fit_clusters(X: np.ndarray, names, rng):
    n = min(len(X), 200_000)
    S = X[np.sort(rng.choice(len(X), n, replace=False))].astype(np.float64)
    sd = S.std(0)
    live = sd > 1e-8
    C = np.eye(X.shape[1])
    C[np.ix_(live, live)] = np.corrcoef(S[:, live], rowvar=False)
    D = np.clip(1 - np.abs(C), 0, None)
    np.fill_diagonal(D, 0)
    Z = linkage(squareform(D, checks=False), method="average")
    lab = fcluster(Z, t=CFG["cluster_cut"], criterion="distance")
    if lab.max() > CFG["cluster_max"]:
        lab = fcluster(Z, t=CFG["cluster_max"], criterion="maxclust")
    clusters = []
    for k in np.unique(lab):
        mem = np.where(lab == k)[0]
        sub = C[np.ix_(mem, mem)]
        v = np.linalg.eigh(sub)[1][:, -1] if len(mem) > 1 else np.ones(1)
        sign = np.where(v >= 0, 1.0, -1.0) * (1 if v.sum() >= 0 else -1)
        medoid = mem[np.argmax(np.abs(sub).sum(0))]
        clusters.append(dict(members=mem, sign=sign, name=names[medoid], size=len(mem)))
    return dict(clusters=clusters, linkage=Z, labels=lab)


def exposures(Xm: np.ndarray, cl) -> np.ndarray:
    """Cluster exposures: signed mean of member ranks, z-scored cross-sectionally."""
    E = np.column_stack([Xm[:, c["members"]] @ c["sign"] / c["size"] for c in cl["clusters"]])
    sd = E.std(0)
    return (E - E.mean(0)) / np.where(sd > 1e-12, sd, 1.0)


def theme_returns(P: Panel, s: int, cl) -> np.ndarray:
    """Return over month s+1 of each theme's long-short portfolio formed at month-end s.
    Weights use the full formation-date universe, v = 2 e / sum|e| on centred exposures; a missing
    next-month return counts as zero after the weights are formed (no renormalisation)."""
    sl = P.rows(s, s + 1)
    E = exposures(P.X[sl], cl)
    den = np.abs(E).sum(0)
    V = 2 * E / np.where(den > 0, den, 1.0)
    return V.T @ np.nan_to_num(P.y_raw[sl])


def theme_premia(f: np.ndarray):
    """Shrunk expected theme returns under each memory, after every formation month.
    f: (S, K) theme returns of formation months 0..S-1. Returns b (H, S, K): b[h, s] uses f[:s+1] only."""
    S, K = f.shape
    out = np.zeros((len(CFG["memories"]), S, K))
    for j, h in enumerate(CFG["memories"]):
        d = 1.0 if h is None else 0.5 ** (1 / (12 * h))
        num, w, w2 = np.zeros(K), 0.0, 0.0
        for s in range(S):
            num, w, w2 = d * num + f[s], d * w + 1.0, d * d * w2 + 1.0
            n_eff = w * w / w2
            out[j, s] = num / w * n_eff / (n_eff + CFG["premia_prior_months"])
    return out


def memory_weights(f: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Walk-forward ledger: candidate h earns b[h, s-1] . f[s] / |b[h, s-1]|_1 in formation month s, a return
    realised after its premia were estimated. Weights are proportional to max(Sharpe over the last
    ledger_months, 0); all zero if no memory has a positive track record."""
    S = f.shape[0]
    if S < 13:
        return np.zeros(b.shape[0])
    lo = max(1, S - CFG["ledger_months"])
    prev = b[:, lo - 1:S - 1]
    rho = (prev * f[None, lo:S]).sum(-1) / np.maximum(np.abs(prev).sum(-1), 1e-12)
    sr = rho.mean(1) / np.maximum(rho.std(1, ddof=1), 1e-12)
    pi = np.maximum(sr, 0.0)
    return pi / pi.sum() if pi.sum() > 0 else pi


def theme_split(mu, E):
    """mu = E (E'E + eps I)^{-1} E' mu  +  residual."""
    G = E.T @ E
    G = G + CFG["proj_ridge"] * np.trace(G) / len(G) * np.eye(len(G))
    th = E @ np.linalg.solve(G, E.T @ mu)
    return th, mu - th


def combine(mu_res, mu_th, mu_alt, gamma):
    """mu_res + (1-gamma) mu_theme + gamma mu_alt, with mu_alt rescaled to the dispersion of mu_theme."""
    sa, st = mu_alt.std(), mu_th.std()
    alt = mu_alt * (st / sa) if sa > 1e-12 else np.zeros_like(mu_alt)
    return mu_res + (1 - gamma) * mu_th + gamma * alt


# ---------------------------------------------------------------------------------------------
# Learners
# ---------------------------------------------------------------------------------------------
def _month_ic(pred, y, groups):
    """Mean cross-sectional Pearson correlation; pred may be (n, L)."""
    pred = pred.reshape(len(y), -1)
    ics = []
    for g in np.unique(groups):
        m = groups == g
        p, t = pred[m], y[m]
        p = p - p.mean(0)
        t = t - t.mean()
        den = np.sqrt((p ** 2).sum(0) * (t ** 2).sum()) + 1e-12
        ics.append((p * t[:, None]).sum(0) / den)
    return np.mean(ics, 0)


class RidgeModel:
    """Ridge on a feature map, penalty picked by mean IC over the last val_months of training.
    Rows are passed as indices into the full panel so the training sample is never copied."""

    def __init__(self, fmap=None):
        self.fmap = fmap or (lambda X: X.astype(np.float64))

    def _gram(self, X, idx, y):
        G, b = 0.0, 0.0
        for i in range(0, len(idx), 20_000):
            Z = self.fmap(X[idx[i:i + 20_000]])
            G = G + (Z.T @ Z).astype(np.float64)
            b = b + Z.T.astype(np.float64) @ y[i:i + 20_000]
        return G, b

    def _solve(self, G, b, mults):
        w, V = np.linalg.eigh(G)
        scale = np.trace(G) / len(G)
        Vb = V.T @ b
        return np.column_stack([V @ (Vb / (w + m * scale)) for m in mults])

    def _predict_many(self, X, idx, B):
        idx = np.arange(len(X)) if idx is None else idx
        return np.vstack([self.fmap(X[idx[i:i + 20_000]]) @ B for i in range(0, len(idx), 20_000)])

    def fit(self, X, idx, y, groups):
        cut = np.unique(groups)[-CFG["val_months"]] if len(np.unique(groups)) > 2 * CFG["val_months"] \
            else np.quantile(groups, 0.8)
        tr, va = groups < cut, groups >= cut
        G, b = self._gram(X, idx[tr], y[tr])
        B = self._solve(G, b, CFG["ridge_grid"])
        ic = _month_ic(self._predict_many(X, idx[va], B), y[va], groups[va])
        self.penalty = float(CFG["ridge_grid"][int(np.argmax(ic))])
        self.val_ic = float(ic.max())
        G2, b2 = self._gram(X, idx[va], y[va])
        self.beta = self._solve(G + G2, b + b2, [self.penalty])[:, 0]
        return self

    def predict(self, X):
        return self._predict_many(X, None, self.beta[:, None])[:, 0]


def make_rff(n_feat, rng):
    """Random Fourier features (cf. Kelly, Malamud & Zhou, 'The Virtue of Complexity')."""
    P = CFG["rff_pairs"]
    W = rng.standard_normal((n_feat, P)) * (CFG["rff_bandwidth"] / np.sqrt(n_feat / 12.0))
    W = W.astype(np.float32)

    def fmap(X):
        H = X @ W
        return np.hstack([np.sin(H), np.cos(H)]) * np.float32(1 / np.sqrt(P))
    return fmap


def _subsample(n, k, rng):
    return np.arange(n) if n <= k else np.sort(rng.choice(n, k, replace=False))


def fit_learners(P: Panel, t: int, seed: int):
    rng = np.random.default_rng(seed)
    sl = P.rows(0, t)
    idx = np.arange(sl.start, sl.stop)[~np.isnan(P.y_rank[sl])]
    y, g = P.y_rank[idx], P.month_idx[idx]

    c0 = max(0, t - CFG["cluster_months"])
    clusters = fit_clusters(P.X[P.rows(c0, t)], P.features, rng)

    ridge = RidgeModel().fit(P.X, idx, y, g)
    s = _subsample(len(idx), CFG["rff_max_rows"], rng)
    rff = RidgeModel(make_rff(P.X.shape[1], rng)).fit(P.X, idx[s], y[s], g[s])
    s = _subsample(len(idx), CFG["gbr_max_rows"], rng)
    gbr = HistGradientBoostingRegressor(early_stopping=False, random_state=seed % (2 ** 31),
                                        **CFG["gbr"]).fit(P.X[idx[s]], y[s])
    return dict(ridge=ridge, rff=rff, gbr=gbr, clusters=clusters)


# ---------------------------------------------------------------------------------------------
# Risk model (Barra-style, common factor + theme factors) and portfolio construction
# ---------------------------------------------------------------------------------------------
def _factor_month(P: Panel, D: Daily, k: int, cl):
    """Daily cross-sectional regressions of returns in month k+1 on exposures at month-end k."""
    sl = P.rows(k, k + 1)
    ids = P.ids[sl]
    B = np.column_stack([np.ones(len(ids)), exposures(P.X[sl], cl)])
    ds = slice(D.offsets[k], D.offsets[k + 1])
    did, dday, dret = D.ids[ds], D.day[ds], D.ret[ds]
    pos = np.searchsorted(ids, did)
    pos = np.clip(pos, 0, len(ids) - 1)
    hit = ids[pos] == did
    pos, dday, dret = pos[hit], dday[hit], dret[hit]
    out_f, out_day, res_id, res_day, res_e = [], [], [], [], []
    for d in np.unique(dday):
        m = dday == d
        if m.sum() < CFG["min_stocks_day"]:
            continue
        Bd, r = B[pos[m]], dret[m]
        f = np.linalg.lstsq(Bd, r, rcond=None)[0]
        out_f.append(f)
        out_day.append(d)
        res_id.append(ids[pos[m]])
        res_day.append(np.full(m.sum(), d))
        res_e.append(r - Bd @ f)
    if not out_f:
        return None
    return dict(f=np.array(out_f), day=np.array(out_day), rid=np.concatenate(res_id),
                rday=np.concatenate(res_day), e=np.concatenate(res_e))


def risk_model(cache, P, D, m, cl, ids_m):
    """Factor covariance F (monthly) and idiosyncratic variances for the stocks of month m."""
    parts = []
    for k in range(max(0, m - CFG["risk_months"]), m):
        if k not in cache:
            cache[k] = _factor_month(P, D, k, cl)
        if cache[k] is not None:
            parts.append(cache[k])
    if not parts:
        return None
    f = np.vstack([p["f"] for p in parts])
    day = np.concatenate([p["day"] for p in parts])
    w = 0.5 ** ((day.max() - day) / (CFG["hl_factor"] * 7 / 5))
    fc = f - (w[:, None] * f).sum(0) / w.sum()
    F = 21 * (w[:, None] * fc).T @ fc / w.sum()

    rid = np.concatenate([p["rid"] for p in parts])
    rday = np.concatenate([p["rday"] for p in parts])
    e2 = np.concatenate([p["e"] for p in parts]) ** 2
    wi = 0.5 ** ((day.max() - rday) / (CFG["hl_idio"] * 7 / 5))
    u, inv = np.unique(rid, return_inverse=True)
    num = np.bincount(inv, wi * e2)
    den = np.bincount(inv, wi)
    cnt = np.bincount(inv)
    var_u = 21 * num / den
    med = np.median(var_u[cnt >= CFG["idio_prior_obs"]]) if (cnt >= CFG["idio_prior_obs"]).any() else np.median(var_u)
    var_u = (cnt * var_u + CFG["idio_prior_obs"] * med) / (cnt + CFG["idio_prior_obs"])
    pos = np.clip(np.searchsorted(u, ids_m), 0, len(u) - 1)
    have = u[pos] == ids_m
    dvar = np.where(have, var_u[pos], med)
    return F, np.maximum(dvar, 1e-6)


def _cap_and_norm(w):
    w = w / (np.abs(w).sum() + 1e-18)
    cap = CFG["pos_cap_mult"] / len(w)
    for _ in range(5):
        w = np.clip(w, -cap, cap)
        w = w / (np.abs(w).sum() + 1e-18)
    return w


def _vol_target(w, B, F, dvar):
    v = (B.T @ w) @ F @ (B.T @ w) + (dvar * w ** 2).sum()
    k = (CFG["target_vol"] / np.sqrt(12)) / np.sqrt(max(v, 1e-12))
    return w * k, float(np.sqrt(v) * k * np.sqrt(12))


def mv_direction(s, B, F, dvar):
    """Sigma^{-1} s with Sigma = B F B' + diag(dvar), via Woodbury."""
    Di = 1.0 / dvar
    BtDi = B.T * Di
    M = np.eye(B.shape[1]) + F @ (BtDi @ B)
    return Di * s - Di * (B @ np.linalg.solve(M, F @ (BtDi @ s)))


def _z(x):
    sd = x.std()
    return (x - x.mean()) / (sd if sd > 1e-12 else 1.0)


# ---------------------------------------------------------------------------------------------
# Engine
# ---------------------------------------------------------------------------------------------
def run(P: Panel, D: Daily, out_months, variants=("final",), verbose=False):
    """Yield one dict per portfolio month with ids, predictions and weights for each variant."""
    out_months = sorted(m for m in out_months if m >= CFG["min_train_months"])
    fitted, cache = None, {}
    for m in out_months:
        is_dec = pd.Timestamp(P.months[m]).month == CFG["refit_month"]
        if fitted is None or (is_dec and fitted["t"] != m):
            seed = CFG["seed"] + int(str(P.months[m])[:4])
            fitted = fit_learners(P, m, seed)
            fitted["t"] = m
            cache = {}
            cl = fitted["clusters"]
            # theme-return history under the current themes, formation months < m (all realised by m)
            fitted["fh"] = [theme_returns(P, s, cl) for s in range(m)]
            fh = np.array(fitted["fh"])
            fitted["pi"] = memory_weights(fh, theme_premia(fh))
            if verbose:
                print(f"refit {P.months[m]}: {len(fitted['clusters']['clusters'])} themes, "
                      f"ridge pen {fitted['ridge'].penalty:g}, rff pen {fitted['rff'].penalty:g}", flush=True)
        sl = P.rows(m, m + 1)
        ids, Xm = P.ids[sl], P.X[sl]
        preds = dict(ridge=fitted["ridge"].predict(Xm), rff=fitted["rff"].predict(Xm),
                     gbr=fitted["gbr"].predict(Xm))
        ens = _z(np.mean([_z(p) for p in preds.values()], 0))
        cl = fitted["clusters"]
        E = exposures(Xm, cl)
        B = np.column_stack([np.ones(len(ids)), E])
        while len(fitted["fh"]) < m:  # formation month m-1 is realised at m
            fitted["fh"].append(theme_returns(P, len(fitted["fh"]), cl))
        bh = theme_premia(np.array(fitted["fh"]))[:, -1]           # (H, K), data up to formation month m-1
        pi = fitted["pi"]
        mu_th, mu_res = theme_split(ens, E)
        mu_dec = E @ (pi @ bh)
        mu_fix = E @ bh[CFG["memories"].index(None)]
        g = CFG["gamma"]
        rm = risk_model(cache, P, D, m, cl, ids)
        if rm is None:
            F, dvar = np.eye(B.shape[1]) * 1e-3, np.full(len(ids), 0.01)
        else:
            F, dvar = rm

        W, vols = {}, {}
        for v in variants:
            if v == "final":
                w = _cap_and_norm(mv_direction(combine(mu_res, mu_th, mu_dec, g), B, F, dvar))
            elif v == "ccar0":
                w = _cap_and_norm(mv_direction(ens, B, F, dvar))
            elif v in GAMMA_VARIANTS:
                gv = CFG["gamma_grid"][GAMMA_VARIANTS.index(v)]
                w = _cap_and_norm(mv_direction(combine(mu_res, mu_th, mu_dec, gv), B, F, dvar))
            elif v == "fixed_rank":
                w = _cap_and_norm(_z(combine(mu_res, mu_th, mu_fix, g)))
            elif v == "decay_rank":
                w = _cap_and_norm(_z(combine(mu_res, mu_th, mu_dec, g)))
            elif v == "ens_rank":
                w = _cap_and_norm(ens)
            elif v.endswith("_rank"):
                w = _cap_and_norm(_z(preds[v[:-5]]))
            elif v == "ridge_decile":
                q = np.quantile(preds["ridge"], [0.1, 0.9])
                w = (preds["ridge"] >= q[1]).astype(float) - (preds["ridge"] <= q[0])
                w = _cap_and_norm(w)
            elif v == "ew_market":
                W[v], vols[v] = np.full(len(ids), 1.0 / len(ids)), np.nan
                continue
            W[v], vols[v] = _vol_target(w, B, F, dvar)
        yield dict(m=m, eom=P.months[m], ids=ids, preds=preds, ens=ens, weights=W, exante_vol=vols,
                   n_themes=len(cl["clusters"]), fitted=fitted, pi=pi,
                   theme_share=float(mu_th.var() / max(ens.var(), 1e-12)))


def main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    """Competition entry point: returns portfolio weights with columns id, eom, w.
    The inputs may also be polars frames or parquet paths (used locally to save memory)."""
    assert CFG["gamma"] is not None, "gamma must be frozen from the validation period"
    np.random.seed(CFG["seed"])  # explicit seed as the rules require; every draw below uses seeded Generators
    P = Panel(chars, features)
    D = Daily(daily_ret, P.months)
    if P.test is not None and P.test.any():
        out = np.unique(P.month_idx[P.test])
    else:
        out = np.arange(CFG["min_train_months"], len(P.months))
    frames = []
    for r in run(P, D, out):
        frames.append(pd.DataFrame({"id": r["ids"].astype(np.int64), "eom": pd.Timestamp(r["eom"]).date(),
                                    "w": r["weights"]["final"]}))
    res = pd.concat(frames, ignore_index=True)
    return res.sort_values(["id", "eom"], ignore_index=True)


if __name__ == "__main__":
    import sys
    root = Path(sys.argv[1] if len(sys.argv) > 1 else "data")
    w = main(root / "ctff_chars.parquet", root / "ctff_features.parquet", root / "ctff_daily_ret.parquet")
    w.to_csv("output.csv", index=False)
    print(w.groupby("eom")["w"].agg(["count", lambda s: s.abs().sum()]).tail())