"""JKP Common Task Framework -- group working draft prepared by Jerry Cao.

This version preserves the shared strategy specification and adds strict input
validation, complete warm-up output, a causal no-flag fallback and bounded
thread use. See the accompanying report for methods and verification.

Model: "Two risk models, one signal: boosted-tree alpha with an hclust-Barra and a
shrunk-sample covariance, selected on a pre-test decade."

Interface (Rules 11-12): main(chars, features, daily_ret) -> DataFrame[id, eom, w].
Everything below is a pure function of the three inputs; no files, no network.

PIPELINE (every quantity at month t uses data dated <= t only; Rule 1)

 0. Canonical input order: rows sorted by (eom, id), features sorted by name, daily
    returns sorted by (date, id). Output does not depend on how the data is fed in.
 1. Features: within-month percentile ranks of every characteristic in `features`
    (average ties, NaN kept). Target: arctan of the within-month z-score of
    ret_exc_lead1m  [target transform from P. Geertsema, CTF entry 30 Aug 2026].
 2. Alpha (expanding blocks, refit every 12 months on the trailing 240 months):
      (a) histogram gradient-boosted trees (scikit-learn), large capacity
          (255 leaves, 500 rounds) regularised by a 2,000-observation leaf minimum
          [leaf-minimum idea from Geertsema];
      (b) ridge regression on centred ranks; the penalty is chosen inside each
          block by time-series cross-validation (last 60 training months, rank IC).
    mu_gbm = z(a); mu_blend = z(0.7 z(a) + 0.3 z(b)).
 3. Risk model A: Ledoit-Wolf shrunk covariance of the trailing 756 daily returns plus
    a Tikhonov ridge equal to the average variance  [Geertsema].
    Risk model B (own construction, following Jensen-Kelly-Malamud-Pedersen's Barra
    model): factors = 12 Fama-French industries + K=13 characteristic clusters.
    The clusters are re-estimated every 12 months by average-linkage hierarchical
    clustering of the trailing 60-month rank-correlation matrix (distance 1-|rho|),
    each cluster's members sign-aligned with its first principal component. Daily
    cross-sectional ridge regressions give factor returns and residuals; the factor
    covariance is EWMA (half-life 504 days for correlations, 84 for variances, 1,260
    days of data) and specific variance is an EWMA (half-life 84) of residuals.
 4. Portfolio per risk model: tangency minus global-minimum-variance portfolio under
    Sigma + lambda*I, lambda = mean variance (dollar neutral, unit gross)  [Geertsema].
    Book options: A, B, or 50/50 A+B. Every book is then made market-beta neutral
    (trailing 252-day beta on the equal-weighted universe) and rescaled to unit gross.
 5. Optional risk overlay: scale month t by median(past ex-ante vol)/ex-ante vol_t,
    ex-ante vol from risk model B, clipped to [0.5, 2].
 6. Model selection (Rule 1 compliant): the 2 x 3 x 2 = 12 configurations are run over
    the 120 months that precede the first ctff_test month. The configuration with the
    highest Sharpe ratio there is frozen and used for every test month.

Prior work disclosed: arctan z-score target, LW + Tikhonov covariance and the
tangency-minus-GMV portfolio follow Paul Geertsema's public CTF entry ("LightGBM with
penalised VCV (MV-GMV)", 30 Aug 2026); the beta-neutral step follows
Abhinav Keshri's public entry ("XGB+MLP + two risk models", Sep 2026). The Barra model
follows Jensen, Kelly, Malamud & Pedersen (2024) and the course benchmark code;
clustering follows Jensen, Kelly & Pedersen (2023).
"""
from __future__ import annotations

import time
import warnings
from typing import Callable

import numpy as np
import pandas as pd
from threadpoolctl import threadpool_limits

# ---------------------------------------------------------------------------------
# Frozen configuration (literals only; Rules 1 and 18)
# ---------------------------------------------------------------------------------
SEED = 42
MIN_TRAIN_MONTHS = 36        # first month that can be traded (validation panel: 123 months)
VAL_MONTHS = 120             # pre-test selection window
REFIT_EVERY = 12             # months between refits (alpha, clusters)
TRAIN_WINDOW = 240           # months of training data for the alpha models
MAX_TRAIN_ROWS = 600_000     # cap on training rows (most recent rows kept)

GBM_PARAMS = dict(loss="squared_error", learning_rate=0.02, max_iter=500,
                  max_leaf_nodes=255, l2_regularization=0.0, max_bins=255,
                  early_stopping=False, random_state=SEED)
GBM_MIN_LEAF = 2000          # capped at n_train/50 for tiny samples
RIDGE_GRID = (1e-3, 1e-2, 1e-1, 1.0, 10.0)   # multiples of the rank variance 1/12
RIDGE_VAL_MONTHS = 60
BLEND_W_GBM = 0.7

LW_WINDOW = 756              # daily obs in the sample covariance
LW_MIN_OBS = 250
KAPPA = 1.0                  # Tikhonov ridge as a multiple of the mean variance
DAYS_PER_MONTH = 21

N_CLUSTERS = 13
CLUSTER_MONTHS = 60
CLUSTER_MIN_COVERAGE = 0.5
COV_OBS = 1260               # daily factor returns used for the factor covariance
HL_COR, HL_VAR, HL_SPEC = 504, 84, 84
SPEC_WINDOW, SPEC_MIN_OBS = 252, 200
SPEC_FLOOR_Q = 0.05
REG_RIDGE = 1e-4

BETA_WINDOW, BETA_MIN_OBS = 252, 60
MIN_NAMES = 20
OVERLAY_CLIP = (0.5, 2.0)
OVERLAY_MIN_HIST = 12

MU_VARIANTS = ("gbm", "blend")
BOOK_VARIANTS = ("A", "B", "AB")
OVERLAY_VARIANTS = ("raw", "vol")
DEFAULT_VARIANT = "blend|AB|vol"

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

# Optional injection point for tests (None -> scikit-learn HistGradientBoostingRegressor)
GBM_FACTORY: Callable | None = None

_T0 = time.perf_counter()


def log(msg: str) -> None:
    print(f"[{time.perf_counter() - _T0:7.0f}s] {msg}", flush=True)


# ---------------------------------------------------------------------------------
# 0. Input handling
# ---------------------------------------------------------------------------------
def integer_ids(values) -> np.ndarray:
    """Preserve integer security IDs; never round them through floating point."""
    values = pd.to_numeric(values, errors="raise")
    if pd.isna(values).any():
        raise ValueError("Security identifiers cannot be missing")
    if pd.api.types.is_integer_dtype(values):
        out = np.asarray(values, dtype=np.int64)
        if not np.array_equal(out, np.asarray(values)):
            raise ValueError("Identifier outside the signed 64-bit range")
        return out
    arr = np.asarray(values, dtype=np.float64)
    if not np.isfinite(arr).all() or (np.abs(arr) > 2**53).any() or (arr != np.floor(arr)).any():
        raise ValueError("Identifiers must be exact integers")
    return arr.astype(np.int64)


def to_day(x) -> np.ndarray:
    return pd.to_datetime(pd.Series(np.asarray(x))).to_numpy().astype("datetime64[D]")


def is_true(flag) -> np.ndarray:
    s = pd.Series(flag)
    if pd.api.types.is_bool_dtype(s):
        return s.fillna(False).to_numpy(bool)
    if pd.api.types.is_numeric_dtype(s):
        return s.fillna(0).to_numpy() != 0
    return s.astype(str).str.strip().str.lower().isin(["1", "1.0", "true", "t"]).to_numpy()


def block_to_float(df: pd.DataFrame) -> np.ndarray:
    """Columns -> float64 matrix, non-numeric entries -> NaN."""
    out = np.empty((len(df), df.shape[1]), dtype=np.float64)
    for j, c in enumerate(df.columns):
        col = df[c]
        if not pd.api.types.is_numeric_dtype(col) or pd.api.types.is_bool_dtype(col):
            col = pd.to_numeric(col.astype("object"), errors="coerce")
        out[:, j] = col.to_numpy(dtype=np.float64, na_value=np.nan)
    return out


def ff12(sic: np.ndarray) -> np.ndarray:
    """Fama-French 12 industry code 0..11 from SIC (11 = Other)."""
    s = np.nan_to_num(np.asarray(sic, float), nan=-1).astype(int)
    out = np.full(s.shape, 11, dtype=np.int8)

    def rng(*pairs):
        m = np.zeros(s.shape, bool)
        for a, b in pairs:
            m |= (s >= a) & (s <= b)
        return m
    # assignment order mirrors the course benchmark (first match wins)
    rules = [
        (0, rng((100, 999), (2000, 2399), (2700, 2749), (2770, 2799), (3100, 3199), (3940, 3989))),
        (1, rng((2500, 2519), (3630, 3659), (3710, 3711), (3714, 3714), (3716, 3716), (3750, 3751),
                (3792, 3792), (3900, 3939), (3990, 3999))),
        (2, rng((2520, 2589), (2600, 2699), (2750, 2769), (3000, 3099), (3200, 3569), (3580, 3629),
                (3700, 3709), (3712, 3713), (3715, 3715), (3717, 3749), (3752, 3791), (3793, 3799),
                (3830, 3839), (3860, 3899))),
        (3, rng((1200, 1399), (2900, 2999))),
        (4, rng((2800, 2829), (2840, 2899))),
        (5, rng((3570, 3579), (3660, 3692), (3694, 3699), (3810, 3829), (7370, 7379))),
        (6, rng((4800, 4899))),
        (7, rng((4900, 4949))),
        (8, rng((5000, 5999), (7200, 7299), (7600, 7699))),
        (9, rng((2830, 2839), (3693, 3693), (3840, 3859), (8000, 8099))),
        (10, rng((6000, 6999))),
    ]
    done = np.zeros(s.shape, bool)
    for code, m in rules:
        sel = m & ~done
        out[sel] = code
        done |= m
    return out


def rank_by_month(arr: np.ndarray, mstart: np.ndarray, mend: np.ndarray) -> np.ndarray:
    """Within-month percentile ranks in (0, 1], average ties, NaN kept. Stored as
    float16 (resolution ~5e-4, finer than the 255 histogram bins of the tree model)
    so the full panel fits in a laptop's memory."""
    out = np.empty(arr.shape, dtype=np.float16)
    for lo, hi in zip(mstart, mend):
        out[lo:hi] = pd.DataFrame(arr[lo:hi]).rank(pct=True, method="average",
                                                   na_option="keep").to_numpy(np.float16)
    return out


def build_panel(meta: pd.DataFrame, feats: list[str],
                get_block: Callable[[list[str]], pd.DataFrame], block_size: int = 20) -> dict:
    """meta: id, eom, ret_exc_lead1m, ctff_test, sic (original row order).
    get_block(cols) returns those feature columns in the same original row order."""
    ids_raw = integer_ids(meta[ID_COL])
    eom = to_day(meta[DATE_COL])
    if np.isnat(eom).any():
        raise ValueError("Monthly dates cannot be missing")
    ok = np.ones(len(eom), dtype=bool)
    order = np.lexsort((ids_raw, eom.astype(np.int64)))
    order = order[ok[order]]
    # Reject ambiguous duplicates instead of keeping an input-order-dependent row.
    e_s, i_s = eom[order], ids_raw[order]
    dup = np.zeros(len(order), bool)
    dup[1:] = (e_s[1:] == e_s[:-1]) & (i_s[1:] == i_s[:-1])
    if dup.any():
        raise ValueError("Duplicate monthly security/date keys")

    ids = ids_raw[order]
    ids = ids.astype(np.int64) if np.all(ids == np.round(ids)) else ids
    eom = eom[order]
    ret = pd.to_numeric(meta[RET_COL], errors="coerce").to_numpy(np.float64)[order]
    test = (is_true(meta[TEST_COL].to_numpy()) if TEST_COL in meta.columns
            else np.zeros(len(meta), bool))[order]
    sic = (pd.to_numeric(meta["sic"], errors="coerce").to_numpy(np.float64)[order]
           if "sic" in meta.columns else np.full(len(order), np.nan))

    months, mstart = np.unique(eom, return_index=True)
    mend = np.append(mstart[1:], len(eom))
    month_ix = np.repeat(np.arange(len(months)), mend - mstart).astype(np.int32)

    n, F = len(order), len(feats)
    R = np.empty((n, F), dtype=np.float16)
    for j in range(0, F, block_size):
        cols = feats[j:j + block_size]
        blk = block_to_float(get_block(cols))[order]
        R[:, j:j + len(cols)] = rank_by_month(blk, mstart, mend)
        del blk

    # target: arctan of the within-month z-score of the lead excess return
    g = pd.Series(ret).groupby(month_ix)
    sd = g.transform("std").to_numpy()
    z = (ret - g.transform("mean").to_numpy()) / np.where(sd > 0, sd, np.nan)
    y = np.arctan(z).astype(np.float32)

    uid = np.unique(ids)
    code = np.searchsorted(uid, ids).astype(np.int32)
    return dict(feats=feats, ids=ids, eom=eom, ret=ret, test=test, ind=ff12(sic),
                months=months, mstart=mstart, mend=mend, month_ix=month_ix,
                R=R, y=y, uid=uid, code=code)


def panel_from_frames(chars: pd.DataFrame, features: pd.DataFrame) -> dict:
    feats = sorted(set(features["features"].dropna().astype(str)))
    if not feats or not set(feats).issubset(chars.columns):
        raise ValueError("Official feature list is missing or inconsistent")
    if set(feats) & {ID_COL, DATE_COL, RET_COL, TEST_COL}:
        raise ValueError("Keys, flags and targets cannot be features")
    meta_cols = [c for c in (ID_COL, DATE_COL, RET_COL, TEST_COL, "sic") if c in chars.columns]
    return build_panel(chars[meta_cols], feats, lambda cols: chars[cols])


def build_daily(ids, dates, rets, panel: dict, first_day: np.datetime64) -> dict:
    """Long daily returns restricted to panel ids and dates >= first_day,
    sorted by (day, id code)."""
    uid = panel["uid"]
    ids = np.asarray(ids)
    pos = np.clip(np.searchsorted(uid, ids), 0, len(uid) - 1)
    keep = (uid[pos] == ids) & (dates >= first_day) & np.isfinite(rets)
    c, d, r = pos[keep].astype(np.int32), dates[keep], rets[keep].astype(np.float32)
    # The historical calendar must not depend on stocks appearing later in the
    # monthly panel. All supplied past dates are known, even for non-held stocks.
    days = np.unique(dates[dates >= first_day])
    d_ix = np.searchsorted(days, d)
    order = np.lexsort((c, d_ix))
    c, d_ix, r = c[order], d_ix[order].astype(np.int32), r[order]
    dup = np.zeros(len(c), bool)
    dup[1:] = (c[1:] == c[:-1]) & (d_ix[1:] == d_ix[:-1])
    if dup.any():
        raise ValueError("Duplicate daily security/date keys")
    day_start = np.searchsorted(d_ix, np.arange(len(days) + 1)).astype(np.int64)
    return dict(dates=days, d=d_ix, c=c, r=r, day_start=day_start)


def daily_from_frame(daily_ret: pd.DataFrame, panel: dict, first_day) -> dict:
    cols = {c.lower(): c for c in daily_ret.columns}
    ids = integer_ids(daily_ret[cols["id"]])
    dates = to_day(daily_ret[cols["date"]])
    if np.isnat(dates).any():
        raise ValueError("Daily dates cannot be missing")
    rets = pd.to_numeric(daily_ret[cols["ret_exc"]], errors="coerce").to_numpy(np.float64)
    return build_daily(ids, dates, rets, panel, first_day)


def window_matrix(daily: dict, values: np.ndarray, d_lo: int, d_hi: int,
                  codes: np.ndarray, n_uid: int) -> np.ndarray:
    """Dense (days x names) matrix of `values` for days [d_lo, d_hi), NaN if absent."""
    lo, hi = daily["day_start"][d_lo], daily["day_start"][d_hi]
    colmap = np.full(n_uid, -1, dtype=np.int64)
    colmap[codes] = np.arange(len(codes))
    cc = colmap[daily["c"][lo:hi]]
    ok = cc >= 0
    M = np.full((d_hi - d_lo, len(codes)), np.nan, dtype=np.float32)
    M[daily["d"][lo:hi][ok] - d_lo, cc[ok]] = values[lo:hi][ok]
    return M


# ---------------------------------------------------------------------------------
# 1-2. Alpha models
# ---------------------------------------------------------------------------------
def make_gbm(n_train: int):
    leaf = int(min(GBM_MIN_LEAF, max(20, n_train // 50)))
    if GBM_FACTORY is not None:
        return GBM_FACTORY(min_samples_leaf=leaf, **GBM_PARAMS)
    from sklearn.ensemble import HistGradientBoostingRegressor
    return HistGradientBoostingRegressor(min_samples_leaf=leaf, **GBM_PARAMS)


def centred(R: np.ndarray, rows: np.ndarray) -> np.ndarray:
    X = R[rows].astype(np.float64) - 0.5
    X[np.isnan(X)] = 0.0
    return X


def gram(R, rows, y=None, chunk=50_000):
    F = R.shape[1]
    G, b = np.zeros((F, F)), np.zeros(F)
    for i in range(0, len(rows), chunk):
        X = centred(R, rows[i:i + chunk])
        G += X.T @ X
        if y is not None:
            b += X.T @ y[rows[i:i + chunk]].astype(np.float64)
    return G, b


def month_rank_ic(pred: np.ndarray, target: np.ndarray, mon: np.ndarray) -> float:
    """Mean over months of the Spearman correlation between pred and target."""
    d = pd.DataFrame({"m": mon, "p": pred, "y": target}).dropna()
    if d.empty:
        return -np.inf
    g = d.groupby("m")
    rp, ry = g["p"].rank(), g["y"].rank()
    rp = rp - rp.groupby(d["m"]).transform("mean")
    ry = ry - ry.groupby(d["m"]).transform("mean")
    num = (rp * ry).groupby(d["m"]).sum()
    den = np.sqrt((rp ** 2).groupby(d["m"]).sum() * (ry ** 2).groupby(d["m"]).sum())
    ic = (num / den.where(den > 0)).dropna()
    return float(ic.mean()) if len(ic) else -np.inf


def fit_ridge(P: dict, tr_rows: np.ndarray):
    """Ridge on centred ranks; penalty picked by rank IC on the last training months."""
    R, y, mix = P["R"], P["y"], P["month_ix"]
    tr_months = np.unique(mix[tr_rows])
    n_val = min(RIDGE_VAL_MONTHS, len(tr_months) // 3)
    G, b = gram(R, tr_rows, y)
    n_all = len(tr_rows)
    lam_best, info = RIDGE_GRID[2], {}
    if n_val >= 6:
        val_mask = mix[tr_rows] >= tr_months[-n_val]
        va_rows, co_rows = tr_rows[val_mask], tr_rows[~val_mask]
        if len(co_rows) > R.shape[1]:
            Gv, bv = gram(R, va_rows, y)
            Gc, bc = G - Gv, b - bv
            best = -np.inf
            for lam in RIDGE_GRID:
                beta = np.linalg.solve(Gc / len(co_rows) + lam / 12.0 * np.eye(len(bc)),
                                       bc / len(co_rows))
                pred = np.concatenate([centred(R, va_rows[i:i + 50_000]) @ beta
                                       for i in range(0, len(va_rows), 50_000)])
                ic = month_rank_ic(pred, y[va_rows], mix[va_rows])
                info[lam] = ic
                if ic > best + 1e-12:
                    best, lam_best = ic, lam
    beta = np.linalg.solve(G / n_all + lam_best / 12.0 * np.eye(len(b)), b / n_all)
    return beta, lam_best, info


def zscore_by_month(x: np.ndarray, mix: np.ndarray) -> np.ndarray:
    s = pd.Series(x)
    g = s.groupby(mix)
    sd = g.transform("std").to_numpy()
    z = (x - g.transform("mean").to_numpy()) / np.where(sd > 0, sd, np.nan)
    return z


def alpha_block(P: dict, t0: int, rows_pred: np.ndarray, diag: dict) -> tuple[np.ndarray, np.ndarray]:
    R, y, ms = P["R"], P["y"], P["mstart"]
    lo = ms[max(0, t0 - TRAIN_WINDOW)]
    tr = np.arange(lo, ms[t0])
    tr = tr[np.isfinite(y[tr])]
    if len(tr) > MAX_TRAIN_ROWS:
        tr = tr[-MAX_TRAIN_ROWS:]
    n_out = len(rows_pred)
    if len(tr) < 100:
        return np.full(n_out, np.nan), np.full(n_out, np.nan)

    X = np.empty((len(tr), R.shape[1]), dtype=np.float64)
    for i in range(0, len(tr), 100_000):
        X[i:i + 100_000] = R[tr[i:i + 100_000]]
    model = make_gbm(len(tr))
    model.fit(X, y[tr].astype(np.float64))
    del X
    gbm = model.predict(R[rows_pred].astype(np.float64))

    beta, lam, info = fit_ridge(P, tr)
    rdg = centred(R, rows_pred) @ beta
    diag.setdefault("alpha_fits", []).append(dict(
        refit_month=str(P["months"][t0]), n_train=int(len(tr)),
        train_from=str(P["months"][P["month_ix"][tr[0]]]),
        ridge_lambda=lam, ridge_val_ic={str(k): v for k, v in info.items()}))
    return gbm, rdg


# ---------------------------------------------------------------------------------
# 3. Risk model B: hierarchical-cluster Barra model
# ---------------------------------------------------------------------------------
def average_linkage(D: np.ndarray, k: int) -> np.ndarray:
    """Agglomerative clustering with average linkage (UPGMA) on distance matrix D,
    stopped at k clusters. Deterministic: ties broken by lowest (row, col) index."""
    m = D.shape[0]
    labels = np.arange(m)
    if k >= m:
        return labels
    Dm = D.astype(np.float64).copy()
    np.fill_diagonal(Dm, np.inf)
    size = np.ones(m)
    active = np.ones(m, bool)
    n_clusters = m
    while n_clusters > k:
        flat = int(np.argmin(Dm))
        i, j = divmod(flat, m)
        if i > j:
            i, j = j, i
        new = (Dm[i] * size[i] + Dm[j] * size[j]) / (size[i] + size[j])
        Dm[i, :], Dm[:, i] = new, new
        Dm[i, i] = np.inf
        Dm[j, :], Dm[:, j] = np.inf, np.inf
        size[i] += size[j]
        active[j] = False
        labels[labels == j] = i
        n_clusters -= 1
    _, lab = np.unique(labels, return_inverse=True)
    return lab


def cluster_map(P: dict, t0: int) -> tuple[np.ndarray, list]:
    """Aggregation matrix A (F x K): cluster exposure = centred ranks @ A."""
    R, ms, me = P["R"], P["mstart"], P["mend"]
    rows = np.arange(ms[max(0, t0 - CLUSTER_MONTHS)], me[t0 - 1])
    F = R.shape[1]
    cover = np.zeros(F)
    for i in range(0, len(rows), 50_000):
        cover += np.isfinite(R[rows[i:i + 50_000]]).sum(0)
    cover /= max(len(rows), 1)
    G, _ = gram(R, rows)
    var = np.diag(G).copy()
    valid = np.flatnonzero((cover >= CLUSTER_MIN_COVERAGE) & (var > 1e-12))
    A = np.zeros((F, 0))
    if len(valid) == 0:
        return A, []
    sd = np.sqrt(var[valid])
    C = G[np.ix_(valid, valid)] / np.outer(sd, sd)
    C = np.clip(C, -1, 1)
    k = min(N_CLUSTERS, len(valid))
    lab = average_linkage(1.0 - np.abs(C), k)
    K = lab.max() + 1
    A = np.zeros((F, K))
    members = []
    for c in range(K):
        idx = np.flatnonzero(lab == c)
        sub = C[np.ix_(idx, idx)]
        if len(idx) > 1:
            _, vec = np.linalg.eigh(sub)
            v = vec[:, -1]
            if v.sum() < 0 or (abs(v.sum()) < 1e-12 and v[np.flatnonzero(v)[0]] < 0):
                v = -v
            sign = np.where(v >= 0, 1.0, -1.0)
        else:
            sign = np.ones(1)
        A[valid[idx], c] = sign / len(idx)
        members.append([("" if s > 0 else "-") + P["feats"][valid[i]] for i, s in zip(idx, sign)])
    return A, members


def month_exposures(P: dict, s: int, A: np.ndarray) -> np.ndarray:
    """Exposure matrix for month s: 12 industry dummies + K z-scored clusters."""
    lo, hi = P["mstart"][s], P["mend"][s]
    ind = np.zeros((hi - lo, 12))
    ind[np.arange(hi - lo), P["ind"][lo:hi]] = 1.0
    if A.shape[1] == 0:
        return ind
    Z = centred(P["R"], np.arange(lo, hi)) @ A
    sd = Z.std(axis=0)
    Z = np.where(sd > 1e-12, (Z - Z.mean(0)) / np.where(sd > 1e-12, sd, 1.0), 0.0)
    return np.hstack([ind, Z])


def factor_regressions(P: dict, daily: dict, A: np.ndarray, d_lo: int, d_hi: int):
    """Daily cross-sectional ridge regressions for days [d_lo, d_hi).
    Day d uses exposures of the last panel month-end strictly before d."""
    months, ms, me = P["months"], P["mstart"], P["mend"]
    dates = daily["dates"]
    K = 12 + A.shape[1]
    Fr = np.full((d_hi - d_lo, K), np.nan)
    resid = np.full(len(daily["r"]), np.nan, dtype=np.float32)
    s_of_day = np.searchsorted(months, dates[d_lo:d_hi], side="left") - 1
    cache = {}
    for k, s in enumerate(s_of_day):
        if s < 0:
            continue
        if s not in cache:
            cache.clear()
            cache[s] = month_exposures(P, s, A)
        X_s = cache[s]
        a, b = daily["day_start"][d_lo + k], daily["day_start"][d_lo + k + 1]
        c = daily["c"][a:b]
        pc = P["code"][ms[s]:me[s]]
        pos = np.clip(np.searchsorted(pc, c), 0, len(pc) - 1)
        hit = pc[pos] == c
        if hit.sum() < K + 5:
            continue
        X, yv = X_s[pos[hit]], daily["r"][a:b][hit].astype(np.float64)
        n = len(yv)
        f = np.linalg.solve(X.T @ X + REG_RIDGE * n * np.eye(K), X.T @ yv)
        Fr[k] = f
        idx = np.arange(a, b)[hit]
        resid[idx] = (yv - X @ f).astype(np.float32)
    return Fr, resid


def ewma_weights(n: int, hl: float) -> np.ndarray:
    return 0.5 ** (np.arange(n - 1, -1, -1) / hl)


def factor_cov(Fr: np.ndarray) -> np.ndarray | None:
    ok = np.all(np.isfinite(Fr), axis=1)
    pos = np.flatnonzero(ok)
    if len(pos) < 60:
        return None
    X = Fr[ok]
    age = (len(Fr) - 1 - pos).astype(float)

    def wcov(hl):
        w = 0.5 ** (age / hl)
        w = w / w.sum()
        mu = w @ X
        Xc = X - mu
        C = (Xc * w[:, None]).T @ Xc
        return C / max(1.0 - np.sum(w ** 2), 1e-12)
    Ccor, Cvar = wcov(HL_COR), wcov(HL_VAR)
    sd_c = np.sqrt(np.clip(np.diag(Ccor), 0, None))
    with np.errstate(invalid="ignore", divide="ignore"):
        cor = Ccor / np.outer(sd_c, sd_c)
    cor[~np.isfinite(cor)] = 0.0
    np.fill_diagonal(cor, 1.0)
    sd_v = np.sqrt(np.clip(np.diag(Cvar), 0, None))
    return sd_v[:, None] * cor * sd_v[None, :] * DAYS_PER_MONTH


def specific_var(resid_win: np.ndarray) -> np.ndarray:
    T = resid_win.shape[0]
    w = ewma_weights(T, HL_SPEC)[:, None]
    obs = np.isfinite(resid_win)
    e2 = np.where(obs, resid_win.astype(np.float64) ** 2, 0.0)
    num, den = (w * e2).sum(0), (w * obs).sum(0)
    D = np.where(den > 0, num / np.where(den > 0, den, 1), np.nan) * DAYS_PER_MONTH
    need = min(SPEC_MIN_OBS, int(0.8 * T))
    D[obs.sum(0) < need] = np.nan
    good = np.isfinite(D) & (D > 0)
    if not good.any():
        return np.full(len(D), np.nan)
    floor = np.quantile(D[good], SPEC_FLOOR_Q)
    D = np.where(good, np.maximum(D, floor), np.median(D[good]))
    return D


# ---------------------------------------------------------------------------------
# 4. Portfolio operators
# ---------------------------------------------------------------------------------
def unit_gross(w: np.ndarray) -> np.ndarray:
    g = np.abs(w).sum()
    return w / g if g > 1e-15 else np.zeros_like(w)


def tan_minus_gmv_dense(MU: np.ndarray, S: np.ndarray) -> np.ndarray:
    """Columns of MU -> tangency minus GMV under S + kappa*mean(diag)*I, unit gross."""
    n = S.shape[0]
    lam = KAPPA * float(np.trace(S, dtype=np.float64)) / n
    A = S.astype(np.float64) + lam * np.eye(n)
    sol = np.linalg.solve(A, np.column_stack([MU, np.ones(n)]))
    b = sol[:, -1]
    W = sol[:, :-1] - np.outer(b, sol[:, :-1].sum(0) / b.sum())
    return np.column_stack([unit_gross(W[:, j]) for j in range(W.shape[1])])


def tan_minus_gmv_factor(MU: np.ndarray, X: np.ndarray, Om: np.ndarray, D: np.ndarray) -> np.ndarray:
    """Same operator under Sigma = X Om X' + diag(D), solved with Woodbury
    (form that does not require Om to be invertible)."""
    n, K = X.shape
    tot = D + np.einsum("ij,jk,ik->i", X, Om, X)
    Dl = D + KAPPA * tot.mean()
    RHS = np.column_stack([MU, np.ones(n)])
    DiR = RHS / Dl[:, None]
    DiX = X / Dl[:, None]
    M = np.eye(K) + Om @ (X.T @ DiX)
    sol = DiR - DiX @ np.linalg.solve(M, Om @ (X.T @ DiR))
    b = sol[:, -1]
    W = sol[:, :-1] - np.outer(b, sol[:, :-1].sum(0) / b.sum())
    return np.column_stack([unit_gross(W[:, j]) for j in range(W.shape[1])])


def ledoit_wolf(X: np.ndarray) -> np.ndarray:
    """Ledoit-Wolf (2004) shrinkage toward a scaled identity. X: (T, n) demeaned."""
    T, n = X.shape
    S = (X.T @ X) / T
    tr = float(np.trace(S, dtype=np.float64))
    mu = tr / n
    nS2 = float(np.square(S, dtype=np.float64).sum())
    d2 = nS2 - 2.0 * mu * tr + n * mu ** 2
    rs = np.square(X, dtype=np.float64).sum(1)
    b2 = float(np.square(rs).sum()) / T ** 2 - nS2 / 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)
    return S


def beta_neutral(w: np.ndarray, beta: np.ndarray) -> np.ndarray:
    h = w != 0
    if h.sum() < 3:
        return w
    B = np.column_stack([beta[h], np.ones(h.sum())])
    coef, *_ = np.linalg.lstsq(B, w[h], rcond=None)
    out = np.zeros_like(w)
    out[h] = w[h] - B @ coef
    return unit_gross(out)


def rank_fallback(mu: np.ndarray) -> np.ndarray:
    ok = np.isfinite(mu)
    w = np.zeros(len(mu))
    if ok.sum() >= 2:
        r = pd.Series(mu[ok]).rank().to_numpy()
        w[ok] = r - r.mean()
    return unit_gross(w)


# ---------------------------------------------------------------------------------
# Strategy
# ---------------------------------------------------------------------------------
def variant_names() -> list[str]:
    return [f"{m}|{b}|{o}" for m in MU_VARIANTS for b in BOOK_VARIANTS for o in OVERLAY_VARIANTS]


def sharpe(r: np.ndarray) -> float:
    r = r[np.isfinite(r)]
    if len(r) < 12 or r.std(ddof=1) <= 0:
        return -np.inf
    return float(r.mean() / r.std(ddof=1) * np.sqrt(12))


def run_strategy(P: dict, daily: dict, keep_all: bool = False) -> dict:
    np.random.seed(SEED)
    months, ms, me = P["months"], P["mstart"], P["mend"]
    nm = len(months)
    test_m = np.unique(P["month_ix"][P["test"]])
    T0 = int(test_m[0]) if len(test_m) else nm
    first = max(MIN_TRAIN_MONTHS, T0 - VAL_MONTHS) if len(test_m) else MIN_TRAIN_MONTHS
    first = min(first, nm - 1)
    anchor = first
    trade = np.arange(first, nm)
    log(f"panel {len(P['ids']):,} rows x {len(P['feats'])} features, {nm} months "
        f"{months[0]}..{months[-1]}; first test month {months[min(T0, nm - 1)]}; "
        f"trading from {months[first]}")

    n_uid = len(P["uid"])
    dates = daily["dates"]
    rows_trade = np.arange(ms[first], me[nm - 1])
    base = ms[first]
    mu = {v: np.full(len(rows_trade), np.nan) for v in MU_VARIANTS}
    raw_gbm = np.full(len(rows_trade), np.nan)
    raw_rdg = np.full(len(rows_trade), np.nan)
    combos = [f"{m}|{b}" for m in MU_VARIANTS for b in BOOK_VARIANTS]
    W = {c: np.zeros(len(rows_trade), dtype=np.float32) for c in combos}
    sig = {c: np.full(nm, np.nan) for c in combos}
    diag: dict = {"clusters": [], "month_info": []}

    for t0 in range(anchor, nm, REFIT_EVERY):
        t1 = min(t0 + REFIT_EVERY, nm)
        rows_b = np.arange(ms[t0], me[t1 - 1])
        log(f"block {months[t0]}..{months[t1 - 1]}: fitting alpha")
        g, r = alpha_block(P, t0, rows_b, diag)
        raw_gbm[rows_b - base], raw_rdg[rows_b - base] = g, r
        mix = P["month_ix"][rows_b]
        zg, zr = zscore_by_month(g, mix), zscore_by_month(r, mix)
        mu["gbm"][rows_b - base] = zg
        bl = np.where(np.isfinite(zr), BLEND_W_GBM * zg + (1 - BLEND_W_GBM) * zr, zg)
        mu["blend"][rows_b - base] = zscore_by_month(bl, mix)

        # risk model B for this block
        A, members = cluster_map(P, t0)
        diag["clusters"].append(dict(refit_month=str(months[t0]), members=members))
        end_last = int(np.searchsorted(dates, months[t1 - 1], side="right"))
        end_first = int(np.searchsorted(dates, months[t0], side="right"))
        d_lo = max(0, end_first - COV_OBS)
        log(f"  factor regressions on {end_last - d_lo} days, K={12 + A.shape[1]}")
        Fr, resid = factor_regressions(P, daily, A, d_lo, end_last)

        for t in range(t0, t1):
            lo, hi = ms[t], me[t]
            n_t = hi - lo
            codes = P["code"][lo:hi]
            sl = slice(lo - base, hi - base)
            MU = np.column_stack([mu[v][sl] for v in MU_VARIANTS])
            okmu = np.all(np.isfinite(MU), axis=1)
            end = int(np.searchsorted(dates, months[t], side="right"))
            info = dict(eom=str(months[t]), n=int(n_t))

            # betas and risk model A
            Xd = window_matrix(daily, daily["r"], max(0, end - LW_WINDOW), end, codes, n_uid)
            obs = np.isfinite(Xd)
            Xb = Xd[-BETA_WINDOW:]
            ob = np.isfinite(Xb)
            mkt = np.nanmean(np.where(ob, Xb, np.nan), axis=1) if Xb.size else np.zeros(0)
            beta = np.ones(n_t)
            if Xb.shape[0] >= BETA_MIN_OBS:
                mk = np.nan_to_num(mkt - np.nanmean(mkt))
                Rb = np.where(ob, Xb, 0.0) - np.where(ob, Xb, 0.0).mean(0)
                vm = float(mk @ mk)
                if vm > 0:
                    beta = (Rb.T @ mk) / vm
                beta[ob.sum(0) < BETA_MIN_OBS] = np.nan
                med = np.nanmedian(beta) if np.isfinite(beta).any() else 1.0
                beta = np.where(np.isfinite(beta), beta, med)
            enough = (obs.sum(0) >= min(LW_MIN_OBS, max(20, Xd.shape[0] // 3))) & okmu
            wA = None
            if enough.sum() >= MIN_NAMES:
                Xa = np.where(obs[:, enough], Xd[:, enough], np.float32(0.0))
                Xa -= Xa.mean(0, keepdims=True)
                S = ledoit_wolf(Xa * np.float32(np.sqrt(DAYS_PER_MONTH)))
                wA = np.zeros((n_t, len(MU_VARIANTS)))
                wA[enough] = tan_minus_gmv_dense(MU[enough], S)
                info["nA"] = int(enough.sum())

            # risk model B
            wB, XB, Om, DB = None, None, None, None
            k_end = end - d_lo
            Om = factor_cov(Fr[max(0, k_end - COV_OBS):k_end]) if k_end > 0 else None
            if Om is not None:
                Rw = window_matrix(daily, resid, max(0, end - SPEC_WINDOW), end, codes, n_uid)
                DB = specific_var(Rw)
                XB = month_exposures(P, t, A)
                okB = okmu & np.isfinite(DB)
                if okB.sum() >= MIN_NAMES:
                    wB = np.zeros((n_t, len(MU_VARIANTS)))
                    wB[okB] = tan_minus_gmv_factor(MU[okB], XB[okB], Om, DB[okB])
                    info["nB"] = int(okB.sum())

            for j, mv in enumerate(MU_VARIANTS):
                for bk in BOOK_VARIANTS:
                    parts = []
                    if "A" in bk and wA is not None:
                        parts.append(wA[:, j])
                    if "B" in bk and wB is not None:
                        parts.append(wB[:, j])
                    if not parts:  # requested book unavailable: use the other / fallback
                        parts = [x[:, j] for x in (wA, wB) if x is not None]
                    w = (np.mean(parts, axis=0) if parts else rank_fallback(MU[:, j]))
                    w = beta_neutral(unit_gross(w), beta)
                    c = f"{mv}|{bk}"
                    W[c][sl] = w
                    if Om is not None and DB is not None:
                        Dfill = np.where(np.isfinite(DB), DB, np.nanmedian(DB) if np.isfinite(DB).any() else 0)
                        e = XB.T @ w
                        sig[c][t] = float(np.sqrt(max(e @ Om @ e + np.sum(w ** 2 * Dfill), 0)))
            diag["month_info"].append(info)
        log(f"  block done ({t1 - t0} months)")

    # ---- returns, overlay, selection -----------------------------------------
    ret_t = np.nan_to_num(P["ret"][rows_trade])
    mix_t = P["month_ix"][rows_trade]
    has_ret = pd.Series(np.isfinite(P["ret"][rows_trade])).groupby(mix_t).any().to_numpy()
    mlist = np.arange(first, nm)
    mult = {}
    for c in combos:
        s = sig[c][first:]
        m = np.ones(len(s))
        hist = []
        for i, v in enumerate(s):
            if np.isfinite(v) and v > 0 and len(hist) >= OVERLAY_MIN_HIST:
                m[i] = np.clip(np.median(hist) / v, *OVERLAY_CLIP)
            if np.isfinite(v) and v > 0:
                hist.append(v)
        mult[c] = m
    rets = {}
    for c in combos:
        base_r = pd.Series(W[c] * ret_t).groupby(mix_t).sum().to_numpy()
        base_r = np.where(has_ret, base_r, np.nan)
        rets[f"{c}|raw"] = base_r
        rets[f"{c}|vol"] = base_r * mult[c]
    val = (mlist < T0)
    table = {v: sharpe(rets[v][val]) for v in variant_names()}
    best = max(variant_names(), key=lambda v: (table[v], v == DEFAULT_VARIANT))
    if not len(test_m) or not np.isfinite(table[best]):
        best = DEFAULT_VARIANT
    log("validation Sharpe ratios: " + ", ".join(f"{k}={v:.2f}" for k, v in table.items()))
    log(f"selected configuration: {best}")

    mvb, bk, ov = best.split("|")
    w_final = W[f"{mvb}|{bk}"].astype(np.float64)
    if ov == "vol":
        w_final = w_final * mult[f"{mvb}|{bk}"][mix_t - first]
    out = dict(selected=best, val_sharpe=table, first=first, T0=T0,
               rows=rows_trade, w=w_final, months=mlist, returns=rets)
    if keep_all:
        out.update(W=W, mult=mult, mu=mu, raw_gbm=raw_gbm, raw_rdg=raw_rdg,
                   sig=sig, diag=diag)
    return out


def needed_first_day(P: dict) -> np.datetime64:
    nm = len(P["months"])
    test_m = np.unique(P["month_ix"][P["test"]])
    T0 = int(test_m[0]) if len(test_m) else nm
    first = min(max(MIN_TRAIN_MONTHS, T0 - VAL_MONTHS) if len(test_m) else MIN_TRAIN_MONTHS, nm - 1)
    # ~ (COV_OBS + LW_WINDOW buffer) trading days -> calendar days with margin
    return P["months"][first] - np.timedelta64(int(COV_OBS * 1.5) + 30, "D")


def weights_frame(P: dict, res: dict, test_only: bool = True) -> pd.DataFrame:
    rows = res["rows"]
    out = pd.DataFrame({"id": P["ids"][rows].astype(np.int64),
                        "eom": pd.to_datetime(P["eom"][rows]),
                        "w": np.nan_to_num(res["w"], nan=0.0, posinf=0.0, neginf=0.0)})
    # Return every requested row, including zero holdings during warm-up.
    requested = P["test"] if test_only and P["test"].any() else np.ones(len(P["ids"]), bool)
    keys = pd.DataFrame({"id": P["ids"][requested], "eom": pd.to_datetime(P["eom"][requested])})
    out = keys.merge(out, on=["id", "eom"], how="left", validate="one_to_one")
    out["w"] = out["w"].fillna(0.0)
    return out[["id", "eom", "w"]].reset_index(drop=True)


def main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    np.random.seed(SEED)
    if chars.empty:
        raise ValueError("Monthly input cannot be empty")
    log("building panel")
    P = panel_from_frames(chars, features)
    daily = daily_from_frame(daily_ret, P, needed_first_day(P))
    log(f"daily returns: {len(daily['r']):,} rows, {len(daily['dates']):,} days")
    # Four native threads avoid oversubscription while retaining deterministic fits.
    with threadpool_limits(limits=4), warnings.catch_warnings(), np.errstate(all="ignore"):
        warnings.simplefilter("ignore", RuntimeWarning)
        res = run_strategy(P, daily)
    out = weights_frame(P, res)
    log(f"returning {len(out):,} rows; mean gross/month "
        f"{out.groupby('eom')['w'].apply(lambda x: x.abs().sum()).mean():.3f}")
    return out