"""Markowitz-ML: an industry- and beta-neutral, risk-optimised ML portfolio for the JKP CTF.

Pipeline (every step uses only information available at the portfolio date):
  1. Features   Cross-sectional rank transform of all characteristics to [-0.5, 0.5],
                missing values set to the median (0).
  2. Returns    Equal-weight ensemble of (a) ridge regression on the ranks, (b) random-
                feature ridge regression and (c) gradient-boosted trees. Refit every 12
                months on an expanding window; ridge penalties picked inside the training
                window by time-series validation; tree settings fixed in advance.
                Target: next-month excess return demeaned within 2-digit SIC industry.
  3. Risk       (i) Barra-style factor model from daily returns: industry factors + style
                factors from hierarchical clustering of the characteristics, Newey-West
                EWMA factor covariance, EWMA specific variances.
                (ii) Ledoit-Wolf shrunk sample covariance of trailing daily returns with a
                Tikhonov ridge.
  4. Portfolio  Mean-variance weights under each risk model, subject to zero net weight in
                every 2-digit SIC industry and zero market beta; the two books are combined
                50/50 at equal ex-ante risk and scaled to a constant ex-ante volatility.

Borrowed from public CTF entries (disclosed per the competition's guidance): the
Ledoit-Wolf/Tikhonov daily-covariance book and the arctan target for the trees
(Geertsema), the idea of holding one forecast in two risk-model books and making the
book beta-neutral (Keshri). All were validated on 1962-1989 before adoption.

Data wrangling is done in ibis (DuckDB backend); linear algebra in numpy.

CTF Admin Modifications (2026-10-07):
--------------------------------------
1. Changed the two conversions of the sic column to integer from cast to
   try_cast (in ff12 and in build_panel).
   Reason: the full dataset stores sic as text and some rows hold an empty
   string, which made the conversion stop the run with "Could not convert
   string '' to INT32". With try_cast an empty value becomes missing, which
   the existing code already maps to industry 12 (Other) and to sic2 = -1.
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field

import ibis
import ibis.expr.types as ir
import numpy as np
import pyarrow as pa
from scipy.cluster.hierarchy import fcluster, linkage
from scipy.spatial.distance import squareform
from scipy.stats import rankdata

SEED = 20261003


@dataclass
class Config:
    refit_every: int = 12  # months between model refits
    min_train_months: int = 120  # months of training data before the first prediction
    val_months: int = 60  # trailing part of the training window used to pick lambdas
    ridge_lambdas: tuple = tuple(10.0 ** np.arange(-5, 3.01, 0.5))  # x mean eigenvalue
    rf_features: int = 4096  # number of random features (sin + cos pairs)
    rf_gamma: float = 1.0  # std of the random projection (x'w) across stocks
    rf_lambdas: tuple = tuple(10.0 ** np.arange(-5, 3.01, 0.5))
    models: tuple = ("ridge", "rf", "gbt")  # ensemble members
    gbt_params: dict = field(default_factory=lambda: dict(
        max_iter=300, learning_rate=0.05, max_leaf_nodes=255, min_samples_leaf=2000,
        l2_regularization=0.0, early_stopping=False))
    n_clusters: int = 13  # style factors in the risk model (as in JKP 2023)
    cov_halflife: int = 504  # days, factor covariance
    var_halflife: int = 84  # days, factor and specific variances
    nw_lags: int = 5
    beta_halflife: int = 252  # days, trailing market beta (EWMA)
    beta_neutral: bool = True  # add zero-beta to the neutrality constraints
    second_book: bool = True  # combine with a Ledoit-Wolf sample-covariance book
    lw_window: int = 756  # days of daily returns for the sample covariance
    lw_kappa: float = 1.0  # Tikhonov ridge, multiple of the mean eigenvalue
    vol_overlay: bool = False  # Moreira-Muir scaling by trailing realised volatility
    overlay_window: int = 24
    overlay_clip: tuple = (0.75, 1.5)  # Newey-West lags (days) for the factor covariance
    min_spec_obs: int = 21  # days of residuals required for a stock's own specific var
    target_vol: float = 0.10  # annualised ex-ante volatility of the final portfolio
    industry_neutral: bool = True
    neutral_groups: str = "sic2"  # "ff12" or "sic2" (2-digit SIC) for the constraint
    target: str = "raw_sic2"  # prediction target, see make_target
    spec_shrink: float = 0.0  # shrink log specific variance toward size-decile mean
    mu_scaling: str = "raw"  # "raw" or "vol": Grinold mu = IC * sigma_spec * z
    extra: dict = field(default_factory=dict)


# ----------------------------------------------------------------------------------------
# 1. Data preparation (ibis)
# ----------------------------------------------------------------------------------------

def ff12(sic: ir.StringValue) -> ir.IntegerValue:
    """Fama-French 12 industry from a SIC code (1-12, 12 = Other)."""
    # CTF-FIX: empty sic strings become missing instead of stopping the run
    s = sic.try_cast("int32")
    ranges = {
        1: [(100, 999), (2000, 2399), (2700, 2749), (2770, 2799), (3100, 3199), (3940, 3989)],
        2: [(2500, 2519), (2590, 2599), (3630, 3659), (3710, 3711), (3714, 3714), (3716, 3716),
            (3750, 3751), (3792, 3792), (3900, 3939), (3990, 3999)],
        3: [(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)],
        4: [(1200, 1399), (2900, 2999)],
        5: [(2800, 2829), (2840, 2899)],
        6: [(3570, 3579), (3660, 3692), (3694, 3699), (3810, 3829), (7370, 7379)],
        7: [(4800, 4899)],
        8: [(4900, 4949)],
        9: [(5000, 5999), (7200, 7299), (7600, 7699)],
        10: [(2830, 2839), (3693, 3693), (3840, 3859), (8000, 8099)],
        11: [(6000, 6999)],
    }
    case = ibis.cases(
        *[(s.between(lo, hi), ind) for ind, rs in ranges.items() for lo, hi in rs],
        else_=12,
    )
    return case.cast("int8")


@dataclass
class Panel:
    """Stock-month panel sorted by (eom, id), held as numpy arrays."""

    ids: np.ndarray  # int64
    eom: np.ndarray  # datetime64[D]
    month: np.ndarray  # int32 month index (0 = first month in data)
    months: np.ndarray  # datetime64[D] eom value for each month index
    start: np.ndarray  # row offset of each month (len = n_months + 1)
    X: np.ndarray  # float32 ranks in [-0.5, 0.5], n x p
    y: np.ndarray  # float32 next-month excess return (nan if unknown)
    ind: np.ndarray  # int8 FF12 industry, 0-based
    sic2: np.ndarray  # int16 2-digit SIC (-1 if missing)
    test: np.ndarray  # bool ctff_test
    me: np.ndarray  # float32 market equity
    features: list


def build_panel(chars: ir.Table, features: list[str]) -> Panel:
    """Sort, map industries, and rank-transform characteristics month by month."""
    tbl = (
        chars.mutate(
            id=chars.id.cast("int64"),
            ind=ff12(chars.sic) - 1,
            # CTF-FIX: empty sic strings become missing instead of stopping the run
            sic2=(chars.sic.try_cast("int32") // 100).fill_null(-1).cast("int16"),
            test=chars.ctff_test.cast("string").lower().isin(["true", "t", "1"]),
            me_raw=chars.market_equity,
        )
        .order_by(["eom", "id"])
    )
    arr = tbl.select("id", "eom", "ind", "sic2", "test", "ret_exc_lead1m", "me_raw").to_pyarrow()
    ids = arr["id"].to_numpy()
    eom = arr["eom"].to_numpy().astype("datetime64[D]")
    months, month = np.unique(eom, return_inverse=True)
    start = np.searchsorted(month, np.arange(len(months) + 1))

    # fetch the characteristics in column blocks to bound peak memory
    X = np.empty((len(ids), len(features)), dtype=np.float32)
    for j0 in range(0, len(features), 50):
        block = features[j0:j0 + 50]
        cols = tbl.select(*block).to_pyarrow()
        for j, f in enumerate(block):
            X[:, j0 + j] = cols[f].to_numpy(zero_copy_only=False).astype(np.float32)
        del cols
    for m in range(len(months)):
        sl = slice(start[m], start[m + 1])
        X[sl] = rank_block(X[sl])

    return Panel(
        ids=ids,
        eom=eom,
        month=month.astype(np.int32),
        months=months,
        start=start,
        X=X,
        y=arr["ret_exc_lead1m"].to_numpy(zero_copy_only=False).astype(np.float32),
        ind=arr["ind"].to_numpy().astype(np.int8),
        sic2=arr["sic2"].to_numpy().astype(np.int16),
        test=arr["test"].to_numpy(zero_copy_only=False).astype(bool),
        me=arr["me_raw"].to_numpy(zero_copy_only=False).astype(np.float32),
        features=list(features),
    )


def rank_block(x: np.ndarray) -> np.ndarray:
    """Average-tie ranks per column scaled to [-0.5, 0.5]; missing -> 0."""
    r = rankdata(x, axis=0, nan_policy="omit", method="average")
    n = np.sum(~np.isnan(x), axis=0)
    with np.errstate(invalid="ignore", divide="ignore"):
        out = (r - 1) / np.maximum(n - 1, 1) - 0.5
    out[np.isnan(out)] = 0.0
    return out.astype(np.float32)


def industry_demean(v: np.ndarray, p: Panel, groups: str = "ff12") -> np.ndarray:
    """Subtract the month-industry mean (ignoring nans)."""
    g = p.sic2.astype(np.int64) + 1 if groups == "sic2" else p.ind.astype(np.int64)
    key = p.month.astype(np.int64) * 1024 + g
    ok = ~np.isnan(v)
    _, inv = np.unique(key, return_inverse=True)
    s = np.bincount(inv[ok], weights=v[ok], minlength=inv.max() + 1)
    c = np.bincount(inv[ok], minlength=inv.max() + 1)
    return (v - (s / np.maximum(c, 1))[inv]).astype(np.float32)


# ----------------------------------------------------------------------------------------
# 2. Return prediction
# ----------------------------------------------------------------------------------------

class RandomFeatures:
    """Random Fourier features z(x) = [sin(g x W), cos(g x W)] / sqrt(P/2)."""

    def __init__(self, p: int, n_features: int, gamma: float, seed: int = SEED):
        rng = np.random.default_rng(seed)
        # ranks have variance 1/12, so x'w has standard deviation ~ gamma
        sd = gamma / np.sqrt(p / 12)
        self.W = (rng.standard_normal((p, n_features // 2)) * sd).astype(np.float32)
        self.scale = np.float32(1.0 / np.sqrt(n_features // 2))

    def __call__(self, X: np.ndarray) -> np.ndarray:
        H = X @ self.W
        return np.hstack([np.sin(H), np.cos(H)]) * self.scale


def monthly_moments(Z_of_month, months, y_of_month):
    """Per-month Gram matrices Z'Z and Z'y (float64)."""
    G, b, n = [], [], []
    for m in months:
        Z = Z_of_month(m).astype(np.float64)
        y = y_of_month(m).astype(np.float64)
        G.append(Z.T @ Z)
        b.append(Z.T @ y)
        n.append(len(y))
    return G, b, n


def ridge_path(G: np.ndarray, b: np.ndarray, n: int, lambdas) -> np.ndarray:
    """Ridge coefficients for each lambda; k x p.

    The penalty is lambda times the average eigenvalue of G/n, so the same grid applies
    to any feature set regardless of scale.
    """
    evals, evecs = np.linalg.eigh(G / n)
    scale = max(evals.mean(), 1e-12)
    proj = evecs.T @ (b / n)
    return np.stack([evecs @ (proj / (evals + lam * scale)) for lam in lambdas])


def ridge_fit(G: np.ndarray, b: np.ndarray, n: int, lam: float) -> np.ndarray:
    """Single-lambda ridge with the same penalty scaling as ridge_path."""
    A = G / n
    scale = max(np.trace(A) / len(A), 1e-12)
    return np.linalg.solve(A + lam * scale * np.eye(len(A)), b / n)


def ls_sharpe(pred: np.ndarray, y: np.ndarray, groups: list[slice]) -> float:
    """Sharpe ratio of the monthly prediction-weighted long-short portfolio."""
    rets = []
    for sl in groups:
        s = pred[sl] - pred[sl].mean()
        a = np.abs(s).sum()
        rets.append(0.0 if a == 0 else s @ y[sl] / a)
    rets = np.array(rets)
    return rets.mean() / rets.std() if rets.std() > 0 else 0.0


# ----------------------------------------------------------------------------------------
# 3. Risk model
# ----------------------------------------------------------------------------------------

def cluster_features(X: np.ndarray, rows: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
    """Hierarchical clustering (average linkage, 1 - |corr|) of the characteristics.

    Returns the cluster label of every feature (-1 for features with no cross-sectional
    variation in the sample) and its sign relative to the cluster's first principal
    component, so that cluster scores average aligned ranks.
    """
    Xs = X[rows].astype(np.float64)
    valid = np.flatnonzero(Xs.std(axis=0) > 1e-6)
    C = np.nan_to_num(np.corrcoef(Xs[:, valid], rowvar=False))
    np.fill_diagonal(C, 1.0)
    D = np.clip(1 - np.abs(C), 0, None)
    np.fill_diagonal(D, 0)
    lab = fcluster(linkage(squareform(D, checks=False), "average"), k, "maxclust") - 1
    labels = np.full(X.shape[1], -1)
    labels[valid] = lab
    signs = np.zeros(X.shape[1])
    signs[valid] = 1.0
    for c in range(lab.max() + 1):
        loc = np.flatnonzero(lab == c)
        if len(loc) > 1:
            v = np.linalg.eigh(C[np.ix_(loc, loc)])[1][:, -1]
            sg = np.sign(v * np.sign(v.sum()))
            sg[sg == 0] = 1
            signs[valid[loc]] = sg
    return labels, signs


def exposures(X: np.ndarray, ind: np.ndarray, labels: np.ndarray, signs: np.ndarray,
              k: int, n_ind: int = 12) -> np.ndarray:
    """Industry dummies followed by k standardised cluster scores."""
    H = np.zeros((len(ind), n_ind), dtype=np.float32)
    H[np.arange(len(ind)), ind] = 1
    S = np.zeros((len(ind), k), dtype=np.float32)
    for c in range(k):
        idx = labels == c
        if idx.any():
            S[:, c] = (X[:, idx] * signs[idx]).mean(axis=1)
    sd = S.std(0)
    S = (S - S.mean(0)) / np.where(sd > 1e-8, sd, 1)
    return np.hstack([H, S])


@dataclass
class RiskState:
    """Risk model at one month-end, aligned to that month's panel rows."""

    B: np.ndarray  # n x K exposures
    F: np.ndarray  # K x K monthly factor covariance
    d: np.ndarray  # n monthly specific variances


# ----------------------------------------------------------------------------------------
# 4. Portfolio construction
# ----------------------------------------------------------------------------------------

def sigma_inv(rs: RiskState, V: np.ndarray) -> np.ndarray:
    """Sigma^-1 V via Woodbury, Sigma = B F B' + diag(d)."""
    Dinv = 1.0 / rs.d
    DV = V * Dinv[:, None] if V.ndim == 2 else V * Dinv
    BtDB = rs.B.T @ (rs.B * Dinv[:, None])
    K = rs.F.shape[0]
    M = np.eye(K) + rs.F @ BtDB
    inner = np.linalg.solve(M, rs.F @ (rs.B.T @ DV))
    return DV - (rs.B @ inner) * (Dinv[:, None] if V.ndim == 2 else Dinv)


def port_vol(rs: RiskState, w: np.ndarray) -> float:
    f = rs.B.T @ w
    return float(np.sqrt(f @ rs.F @ f + np.sum(rs.d * w * w)))


def constraint_matrix(ind: np.ndarray, beta: np.ndarray | None) -> np.ndarray:
    """Industry dummies (zero net weight per industry), plus market beta if given."""
    present = np.unique(ind)
    H = np.zeros((len(ind), len(present)))
    H[np.arange(len(ind)), np.searchsorted(present, ind)] = 1
    return H if beta is None else np.hstack([H, beta[:, None]])


def markowitz(mu: np.ndarray, rs: RiskState, ind: np.ndarray, cfg: Config, n_ind: int = 12,
              beta: np.ndarray | None = None):
    """Mean-variance weights, optionally with zero net weight in every industry (and
    zero market beta)."""
    w = sigma_inv(rs, mu)
    if cfg.industry_neutral:
        H = constraint_matrix(ind, beta)
        SiH = sigma_inv(rs, H)
        lam = np.linalg.lstsq(H.T @ SiH, H.T @ w, rcond=None)[0]
        w = w - SiH @ lam
    vol = port_vol(rs, w) * np.sqrt(12)
    return w * (cfg.target_vol / vol) if vol > 0 else w


# ----------------------------------------------------------------------------------------
# 3b. Risk model estimation from daily returns
# ----------------------------------------------------------------------------------------

def align_daily(daily: ir.Table, p: Panel) -> dict[str, np.ndarray]:
    """Attach to each daily return the panel row of the same stock at the previous eom.

    Exposures known at the end of month m-1 explain daily returns during month m.
    """
    keys = ibis.memtable(pa.table({
        "id": pa.array(p.ids),
        "eom_prev": pa.array(p.eom),
        "row": pa.array(np.arange(len(p.ids), dtype=np.int32)),
    }))
    d = daily.mutate(
        id=daily.id.cast("int64"),
        eom_prev=(daily.date.truncate("M") - ibis.interval(days=1)).cast("date"),
    )
    j = (
        d.join(keys, ["id", "eom_prev"])
        .filter(lambda t: t.ret_exc.notnull())
        .select("date", "row", "ret_exc")
        .order_by(["date", "row"])
    )
    arr = j.to_pyarrow()
    return {
        "date": arr["date"].to_numpy().astype("datetime64[D]"),
        "row": arr["row"].to_numpy().astype(np.int32),
        "ret": arr["ret_exc"].to_numpy().astype(np.float32),
    }


def estimate_risk(p: Panel, dly: dict, B: np.ndarray, cfg: Config):
    """Daily cross-sectional WLS factor regressions + EWMA covariances.

    Returns, for every month-end, the monthly factor covariance F[m] and, for every panel
    row, the monthly specific variance d (both using daily data up to that eom only).
    """
    n_m, K = len(p.months), B.shape[1]
    ids_u, id_idx = np.unique(p.ids, return_inverse=True)
    a_cov = 0.5 ** (1 / cfg.cov_halflife)
    a_var = 0.5 ** (1 / cfg.var_halflife)

    L = cfg.nw_lags
    S_cov = np.zeros((L + 1, K, K)); w_cov = 0.0  # EWMA of f_t f_{t-l}', l = 0..L
    S_var = np.zeros((L + 1, K)); w_var = 0.0
    f_hist = np.zeros((L, K))  # most recent factor returns, f_hist[0] = yesterday
    spec = np.zeros(len(ids_u)); spec_w = np.zeros(len(ids_u)); spec_n = np.zeros(len(ids_u))
    a_beta = 0.5 ** (1 / cfg.beta_halflife)
    s_rm = np.zeros(len(ids_u)); s_mm = np.zeros(len(ids_u))  # EWMA of r_i*m and m*m
    beta_row = np.full(len(p.ids), 1.0)
    last_day = np.full(len(ids_u), -1, dtype=np.int64)

    F = np.full((n_m, K, K), np.nan)
    d_row = np.full(len(p.ids), np.nan)
    me = p.me.astype(np.float64).copy()
    for m in range(n_m):  # fill missing size with the same month's median (no look-ahead)
        sl = slice(p.start[m], p.start[m + 1])
        v = me[sl]
        v[np.isnan(v)] = np.nanmedian(v) if (~np.isnan(v)).any() else 1.0
    sqrt_me = np.sqrt(me.clip(1e-3))

    dates = dly["date"]
    bounds = np.flatnonzero(np.diff(dates.astype(np.int64))) + 1
    bounds = np.concatenate([[0], bounds, [len(dates)]])
    day_vals = dates[bounds[:-1]]
    # month index each day belongs to (the eom at or after the date)
    m_of_day = np.searchsorted(p.months, day_vals)

    def close_month(m):
        if w_cov > 0:
            # Newey-West (Bartlett) aggregation of daily autocovariances
            cov = S_cov[0] / w_cov
            var = S_var[0] / w_var
            for l in range(1, L + 1):
                b = 1 - l / (L + 1)
                G = S_cov[l] / w_cov
                cov = cov + b * (G + G.T)
                var = var + 2 * b * S_var[l] / w_var
            var = np.maximum(var, 0.25 * S_var[0] / w_var)
            ev, evec = np.linalg.eigh((cov + cov.T) / 2)  # keep PSD
            cov = (evec * np.maximum(ev, 1e-12)) @ evec.T
            sd = np.sqrt(np.diag(cov))
            corr = cov / np.outer(sd, sd).clip(1e-20)
            F[m] = 21 * corr * np.outer(np.sqrt(var), np.sqrt(var))
        sl = slice(p.start[m], p.start[m + 1])
        gi = id_idx[sl]
        ok = spec_n[gi] >= cfg.min_spec_obs
        v = np.full(len(gi), np.nan)
        v[ok] = spec[gi[ok]] / spec_w[gi[ok]]
        # fallback: median specific variance of the same industry, else cross-section
        if (~ok).any():
            ind = p.ind[sl]
            med_all = np.nanmedian(v) if ok.any() else 0.02 ** 2
            for g in np.unique(ind[~ok]):
                vals = v[(ind == g) & ok]
                v[(ind == g) & ~ok] = np.median(vals) if len(vals) else med_all
        d_row[sl] = 21 * v
        bt = np.where(s_mm[gi] > 0, s_rm[gi] / np.maximum(s_mm[gi], 1e-20), np.nan)
        bt[spec_n[gi] < cfg.min_spec_obs] = np.nan
        med = np.nanmedian(bt) if np.isfinite(bt).any() else 1.0
        beta_row[sl] = np.clip(np.where(np.isfinite(bt), bt, med), -1, 4)

    cur_m = 0
    for k in range(len(day_vals)):
        m = m_of_day[k]
        while cur_m < m and cur_m < n_m:  # all days up to eom[cur_m] processed
            close_month(cur_m)
            cur_m += 1
        if m >= n_m:
            break
        sl = slice(bounds[k], bounds[k + 1])
        rows, r = dly["row"][sl], dly["ret"][sl].astype(np.float64)
        if len(rows) < 3 * K:
            continue
        r = np.clip(r, -0.5, 1.0)
        Bd = B[rows].astype(np.float64)
        wts = sqrt_me[rows]
        sw = np.sqrt(wts)
        f, *_ = np.linalg.lstsq(Bd * sw[:, None], r * sw, rcond=None)
        e = r - Bd @ f
        lagged = np.vstack([f, f_hist])  # f_t, f_{t-1}, ..., f_{t-L}
        S_cov = a_cov * S_cov + f[None, :, None] * lagged[:, None, :]
        S_var = a_var * S_var + f[None, :] * lagged
        w_cov = a_cov * w_cov + 1
        w_var = a_var * w_var + 1
        f_hist = lagged[:L]
        gi = id_idx[rows]
        # decay each stock's EWMA by the number of trading days since its last update
        gap = np.where(last_day[gi] >= 0, k - last_day[gi], 1)
        dec = a_var ** gap
        mkt = r.mean()  # equal-weighted market return of the day
        db = a_beta ** gap
        s_rm[gi] = db * s_rm[gi] + r * mkt
        s_mm[gi] = db * s_mm[gi] + mkt * mkt
        spec[gi] = dec * spec[gi] + e * e
        spec_w[gi] = dec * spec_w[gi] + 1
        spec_n[gi] += 1
        last_day[gi] = k
    while cur_m < n_m:
        close_month(cur_m)
        cur_m += 1
    return F, d_row, beta_row


# ----------------------------------------------------------------------------------------
# 2b. Expanding-window training and out-of-sample prediction
# ----------------------------------------------------------------------------------------

class Moments:
    """Running sums of Z'Z, Z'y and n over months added so far."""

    def __init__(self, dim: int):
        self.G = np.zeros((dim, dim)); self.b = np.zeros(dim); self.n = 0; self.upto = 0

    def add_until(self, m_end: int, Z_of, y_of):
        for m in range(self.upto, m_end):
            Z, y = Z_of(m), y_of(m)
            if len(y):
                self.G += (Z.T @ Z).astype(np.float64)
                self.b += (Z.T @ y).astype(np.float64)
                self.n += len(y)
        self.upto = max(self.upto, m_end)


def make_target(p: Panel, kind: str = "raw") -> np.ndarray:
    """Industry-demeaned next-month return (raw) or its centred cross-sectional rank.

    kind: "raw" | "rank" (FF12 demeaning) or "raw_sic2" | "rank_sic2" (2-digit SIC)."""
    groups = "sic2" if kind.endswith("_sic2") else "ff12"
    kind = kind.removesuffix("_sic2")
    y = industry_demean(p.y, p, groups)
    if kind == "rank":
        out = np.full_like(y, np.nan)
        for m in range(len(p.months)):
            sl = slice(p.start[m], p.start[m + 1])
            v = y[sl]
            ok = ~np.isnan(v)
            if ok.sum() > 1:
                r = rankdata(v[ok]) / ok.sum() - 0.5
                out[sl][ok] = r
        y = industry_demean(out, p, groups)
    return y


def predict_gbt(p: Panel, cfg: Config, target: np.ndarray, log=print,
                stop: int | None = None) -> np.ndarray:
    """Gradient-boosted trees (sklearn HistGradientBoosting) on the ranks, fixed
    hyper-parameters, refit every cfg.refit_every months on all realised months.
    The target is arctan of the month-standardised target, which bounds the tails."""
    from sklearn.ensemble import HistGradientBoostingRegressor
    z = standardize_by_month(target, p)
    yt = np.arctan(z).astype(np.float32)
    n_m = len(p.months) if stop is None else stop
    pred = np.full(len(p.ids), np.nan, dtype=np.float32)
    for T in range(cfg.min_train_months + 1, n_m, cfg.refit_every):
        rows = slice(0, p.start[T])
        ok = ~np.isnan(yt[rows])
        m = HistGradientBoostingRegressor(random_state=SEED, **cfg.gbt_params)
        m.fit(p.X[rows][ok], yt[rows][ok])
        hi = min(T + cfg.refit_every, n_m)
        pred[p.start[T]:p.start[hi]] = m.predict(p.X[p.start[T]:p.start[hi]])
        log(f"  gbt refit {str(p.months[T])[:7]}: {ok.sum():,} rows")
    return pred


def predict_oos(p: Panel, cfg: Config, model: str, target: np.ndarray,
                log=print) -> np.ndarray:
    """Out-of-sample predictions for every row, refit every cfg.refit_every months."""
    if model == "ridge":
        feat = lambda m: p.X[p.start[m]:p.start[m + 1]]
        dim, lambdas = p.X.shape[1], cfg.ridge_lambdas
    elif model == "rf":
        rf = RandomFeatures(p.X.shape[1], cfg.rf_features, cfg.rf_gamma)
        feat = lambda m: rf(p.X[p.start[m]:p.start[m + 1]])
        dim, lambdas = cfg.rf_features, cfg.rf_lambdas
    else:
        raise ValueError(model)

    def rows_ok(m):
        sl = slice(p.start[m], p.start[m + 1])
        return ~np.isnan(target[sl])

    Z_of = lambda m: feat(m)[rows_ok(m)]
    y_of = lambda m: target[p.start[m]:p.start[m + 1]][rows_ok(m)]

    n_m = len(p.months)
    pred = np.full(len(p.ids), np.nan, dtype=np.float32)
    full, lag = Moments(dim), Moments(dim)
    first = cfg.min_train_months + 1
    for T in range(first, n_m, cfg.refit_every):
        # Data with returns realised by eom[T]: months <= T-1
        full.add_until(T, Z_of, y_of)
        lag.add_until(T - cfg.val_months, Z_of, y_of)
        # pick lambda on the trailing validation window
        betas = ridge_path(lag.G, lag.b, lag.n, lambdas)
        vm = range(T - cfg.val_months, T)
        Zv = [Z_of(m) for m in vm]
        yv = np.concatenate([y_of(m) for m in vm])
        sizes = np.cumsum([0] + [len(z) for z in Zv])
        groups = [slice(sizes[i], sizes[i + 1]) for i in range(len(Zv))]
        Zv = np.vstack(Zv)
        scores = [ls_sharpe(Zv @ bb, yv, groups) for bb in betas]
        best = int(np.argmax(scores))
        beta = ridge_fit(full.G, full.b, full.n, lambdas[best])
        for m in range(T, min(T + cfg.refit_every, n_m)):
            sl = slice(p.start[m], p.start[m + 1])
            pred[sl] = feat(m) @ beta
        log(f"  {model} refit {str(p.months[T])[:7]}: lambda={lambdas[best]:.3g} "
            f"val SR={scores[best] * np.sqrt(12):.2f}")
    return pred


def standardize_by_month(v: np.ndarray, p: Panel) -> np.ndarray:
    out = np.full_like(v, np.nan)
    for m in range(len(p.months)):
        sl = slice(p.start[m], p.start[m + 1])
        x = v[sl]
        if np.isnan(x).all():
            continue
        s = np.nanstd(x)
        out[sl] = (x - np.nanmean(x)) / (s if s > 0 else 1)
    return out


# ----------------------------------------------------------------------------------------
# 5. End-to-end pipeline
# ----------------------------------------------------------------------------------------

def risk_exposures(p: Panel, cfg: Config, log=print) -> np.ndarray:
    """Exposures for every row. Clusters are re-estimated every refit_every months from
    the characteristics of the preceding 10 years (no return data involved). The first
    fit, after 10 years of data, also defines exposures for those earlier months, which
    only serve to warm up the EWMA risk estimates before any portfolio is formed."""
    K = 12 + cfg.n_clusters
    B = np.zeros((len(p.ids), K), dtype=np.float32)
    n_m = len(p.months)
    rng = np.random.default_rng(SEED)
    first = min(119, n_m - 1)
    for T in range(first, n_m, cfg.refit_every):
        lo = p.start[max(0, T - 119)]
        rows = np.arange(lo, p.start[T + 1])
        if len(rows) > 200_000:
            rows = np.sort(rng.choice(rows, 200_000, replace=False))
        labels, signs = cluster_features(p.X, rows, cfg.n_clusters)
        m_lo = 0 if T == first else T
        m_hi = min(T + cfg.refit_every, n_m)
        for m in range(m_lo, m_hi):
            sl = slice(p.start[m], p.start[m + 1])
            B[sl] = exposures(p.X[sl], p.ind[sl], labels, signs, cfg.n_clusters)
    return B


def combine_predictions(preds: dict[str, np.ndarray], p: Panel) -> np.ndarray:
    """Equal-weight average of month-standardised predictions."""
    z = [standardize_by_month(v, p) for v in preds.values()]
    return np.nanmean(np.stack(z), axis=0) if len(z) > 1 else z[0]


def shrink_specific(d: np.ndarray, me: np.ndarray, s: float) -> np.ndarray:
    """Bayesian-style shrinkage of log specific variance toward its size-decile mean."""
    me = np.nan_to_num(me, nan=np.nanmedian(me))
    dec = np.minimum((rankdata(me) - 1) * 10 // len(me), 9).astype(int)
    ld = np.log(d)
    grp = np.bincount(dec, weights=ld, minlength=10) / np.maximum(np.bincount(dec, minlength=10), 1)
    return np.exp((1 - s) * ld + s * grp[dec])


class DailyMatrix:
    """Dense day x stock matrix of daily excess returns (NaN = not in the panel)."""

    def __init__(self, p: Panel, dly: dict):
        self.ids_u, id_idx = np.unique(p.ids, return_inverse=True)
        self.col_of_row = id_idx.astype(np.int32)
        self.days, day_idx = np.unique(dly["date"], return_inverse=True)
        self.M = np.full((len(self.days), len(self.ids_u)), np.nan, dtype=np.float32)
        self.M[day_idx, self.col_of_row[dly["row"]]] = np.clip(dly["ret"], -0.5, 1.0)

    def window(self, eom: np.datetime64, cols: np.ndarray, n_days: int, min_obs: int):
        """Trailing n_days returns up to and including eom; columns with >= min_obs."""
        end = int(np.searchsorted(self.days, eom, side="right"))
        X = self.M[max(0, end - n_days):end, cols]
        obs = np.isfinite(X)
        ok = obs.sum(0) >= min_obs
        X = np.where(obs, X, 0.0)[:, ok]
        return X - X.mean(0), ok


def ledoit_wolf(X: np.ndarray) -> np.ndarray:
    """Ledoit-Wolf (2004) linear shrinkage of the sample covariance toward a scaled identity."""
    T, n = X.shape
    X = X.astype(np.float64)
    S = X.T @ X / T
    mu = np.trace(S) / n
    d2 = np.sum(S * S) - 2 * mu * np.trace(S) + n * mu * mu
    rs = np.sum(X * X, axis=1)
    b2 = (np.sum(rs * rs) / T ** 2 - np.sum(S * S) / T)
    delta = min(max(b2, 0.0), d2) / d2 if d2 > 0 else 1.0
    S *= 1 - delta
    S[np.diag_indices(n)] += delta * mu
    return S


def sample_cov_book(mu: np.ndarray, groups: np.ndarray, beta: np.ndarray | None,
                    dm: DailyMatrix, cols: np.ndarray, eom, cfg: Config) -> np.ndarray:
    """Industry(-and-beta)-neutral mean-variance weights under a Tikhonov-regularised
    Ledoit-Wolf covariance of trailing daily returns; unit ex-ante volatility."""
    X, ok = dm.window(eom, cols, cfg.lw_window, 250)
    w = np.zeros(len(mu))
    if ok.sum() < 50:
        return w
    S = ledoit_wolf(X) * 21
    A = S + cfg.lw_kappa * np.trace(S) / len(S) * np.eye(len(S))
    H = constraint_matrix(groups[ok], None if beta is None else beta[ok])
    sol = np.linalg.solve(A, np.column_stack([mu[ok], H]))
    a, AH = sol[:, 0], sol[:, 1:]
    lam = np.linalg.lstsq(H.T @ AH, H.T @ a, rcond=None)[0]
    wo = a - AH @ lam
    v = np.sqrt(wo @ S @ wo)
    w[ok] = wo / v if v > 0 else 0
    return w


def vol_overlay(p: Panel, w: np.ndarray, cfg: Config) -> np.ndarray:
    """Moreira-Muir (2017) overlay: scale month t by the expanding mean of the book's
    trailing realised volatility over its current value, both from returns realised by
    eom[t] (months <= t-1), clipped to cfg.overlay_clip."""
    n_m = len(p.months)
    r = np.full(n_m, np.nan)
    for m in range(n_m):
        sl = slice(p.start[m], p.start[m + 1])
        if not np.isnan(w[sl]).all():
            r[m] = np.nansum(np.nan_to_num(w[sl]) * np.nan_to_num(p.y[sl]))
    out = w.copy()
    hist, vols = [], []
    for t in range(n_m):
        if t > 0 and not np.isnan(r[t - 1]):
            hist.append(r[t - 1])
        if len(hist) >= cfg.overlay_window:
            vols.append(np.std(hist[-cfg.overlay_window:], ddof=1))
            lev = np.clip(np.mean(vols) / vols[-1], *cfg.overlay_clip)
            sl = slice(p.start[t], p.start[t + 1])
            out[sl] = w[sl] * lev
    return out


def build_weights(p: Panel, mu: np.ndarray, B: np.ndarray, F: np.ndarray,
                  d: np.ndarray, cfg: Config, months=None, beta: np.ndarray | None = None,
                  dm: DailyMatrix | None = None) -> np.ndarray:
    """Markowitz weights for every row in the selected months (nan elsewhere)."""
    w = np.full(len(p.ids), np.nan)
    months = range(len(p.months)) if months is None else months
    for m in months:
        sl = slice(p.start[m], p.start[m + 1])
        mu_m = mu[sl]
        if np.isnan(mu_m).all() or np.isnan(F[m]).any():
            continue
        mu_m = np.nan_to_num(mu_m.astype(np.float64) - np.nanmean(mu_m))
        d_m = d[sl]
        if cfg.spec_shrink > 0:
            d_m = shrink_specific(d_m, p.me[sl], cfg.spec_shrink)
        rs = RiskState(B=B[sl].astype(np.float64), F=F[m], d=d_m)
        if cfg.mu_scaling == "vol":
            mu_m = mu_m * np.sqrt(rs.d)
        groups = p.sic2[sl] if cfg.neutral_groups == "sic2" else p.ind[sl]
        b_m = beta[sl] if (cfg.beta_neutral and beta is not None) else None
        wm = markowitz(mu_m, rs, groups, cfg, beta=b_m)
        if cfg.second_book and dm is not None:
            ws = sample_cov_book(mu_m, groups, b_m, dm, dm.col_of_row[sl], p.months[m], cfg)
            if np.abs(ws).sum() > 0:
                wm = 0.5 * wm / port_vol(rs, wm) + 0.5 * ws
                wm *= cfg.target_vol / np.sqrt(12) / port_vol(rs, wm)
        w[sl] = wm
    if cfg.vol_overlay:
        w = vol_overlay(p, w, cfg)
    return w


def run_pipeline(chars: ir.Table, features: list[str], daily: ir.Table, cfg: Config,
                 log=print) -> tuple[Panel, np.ndarray]:
    p = build_panel(chars, features)
    log(f"panel: {len(p.ids):,} rows, {len(p.months)} months, {p.X.shape[1]} features")
    target = make_target(p, cfg.target)
    preds = {m: (predict_gbt(p, cfg, target, log) if m == "gbt"
                 else predict_oos(p, cfg, m, target, log)) for m in cfg.models}
    mu = combine_predictions(preds, p)
    B = risk_exposures(p, cfg, log)
    dly = align_daily(daily, p)
    F, d, beta = estimate_risk(p, dly, B, cfg)
    dm = DailyMatrix(p, dly) if cfg.second_book else None
    del dly
    w = build_weights(p, mu, B, F, d, cfg, beta=beta, dm=dm)
    return p, w


def main(chars, features, daily_ret):
    """CTF entry point. Inputs arrive as pandas DataFrames from the harness; they are
    handed to DuckDB through ibis immediately and all work happens there and in numpy."""
    np.random.seed(SEED)
    os.environ.setdefault("OMP_NUM_THREADS", "32")
    cfg = Config()
    ibis.set_backend("duckdb")
    chars_t = ibis.memtable(chars)
    daily_t = ibis.memtable(daily_ret)
    feats = ibis.memtable(features).features.to_pyarrow().to_pylist()
    del chars, daily_ret

    p, w = run_pipeline(chars_t, feats, daily_t, cfg)
    keep = p.test.copy()
    w = np.where(np.isnan(w), 0.0, w)  # months without a model portfolio hold cash
    out = ibis.memtable(pa.table({
        "id": pa.array(p.ids[keep]),
        "eom": pa.array(p.eom[keep]),
        "w": pa.array(w[keep]),
    }))
    return out.to_pandas()


# ----------------------------------------------------------------------------------------
# 2c. Portfolio-ML: coefficients that maximise the Sharpe ratio of the final portfolio
# ----------------------------------------------------------------------------------------

def neutral_markowitz_map(rs: RiskState, ind: np.ndarray, V: np.ndarray) -> np.ndarray:
    """M V with M = Sigma^-1 - Sigma^-1 H (H' Sigma^-1 H)^-1 H' Sigma^-1 (symmetric)."""
    present = np.unique(ind)
    H = np.zeros((len(ind), len(present)))
    H[np.arange(len(ind)), np.searchsorted(present, ind)] = 1
    SiH = sigma_inv(rs, H)
    SiV = sigma_inv(rs, V)
    return SiV - SiH @ np.linalg.solve(H.T @ SiH, H.T @ SiV)


def managed_returns(p: Panel, B: np.ndarray, F: np.ndarray, d: np.ndarray,
                    feat=None, dim=None) -> np.ndarray:
    """Monthly returns of the characteristic-managed portfolios x_j' M r_{t+1}.

    Each month is normalised so that the portfolio built from a unit-variance signal has
    unit ex-ante variance, keeping months comparable across volatility regimes.
    """
    feat = feat or (lambda m: p.X[p.start[m]:p.start[m + 1]])
    dim = dim or p.X.shape[1]
    R = np.full((len(p.months), dim), np.nan)
    for m in range(len(p.months)):
        sl = slice(p.start[m], p.start[m + 1])
        y = p.y[sl]
        if np.isnan(F[m]).any() or np.isnan(y).all():
            continue
        rs = RiskState(B=B[sl].astype(np.float64), F=F[m], d=d[sl])
        r = np.nan_to_num(y.astype(np.float64))
        Mr = neutral_markowitz_map(rs, p.ind[sl], r)
        # scale: typical ex-ante variance of a neutral MV portfolio of one signal, z' M z
        z = np.random.default_rng(m).standard_normal(len(r))
        scale = np.sqrt(z @ neutral_markowitz_map(rs, p.ind[sl], z) / len(r))
        R[m] = feat(m).astype(np.float64).T @ Mr / (scale * len(r))
    return R


def predict_pfml(p: Panel, cfg: Config, R: np.ndarray, feat=None, lambdas=None,
                 log=print) -> np.ndarray:
    """Expanding-window ridge tangency portfolio of the managed portfolios; returns
    mu = x' beta for every row (the input to the Markowitz step)."""
    feat = feat or (lambda m: p.X[p.start[m]:p.start[m + 1]])
    lambdas = lambdas or cfg.ridge_lambdas
    n_m = len(p.months)
    ok = ~np.isnan(R).any(axis=1)
    pred = np.full(len(p.ids), np.nan, dtype=np.float32)
    first = cfg.min_train_months + 1
    for T in range(first, n_m, cfg.refit_every):
        tr = np.flatnonzero(ok[:T])  # months whose returns are realised by eom[T]
        if len(tr) < cfg.val_months + 24:
            continue
        fit, val = tr[:-cfg.val_months], tr[-cfg.val_months:]
        Rf = R[fit]
        betas = ridge_path(Rf.T @ Rf, Rf.sum(0), len(fit), lambdas)
        pv = R[val] @ betas.T  # validation returns for each lambda
        sr = pv.mean(0) / pv.std(0)
        best = int(np.argmax(sr))
        Ra = R[tr]
        beta = ridge_path(Ra.T @ Ra, Ra.sum(0), len(tr), [lambdas[best]])[0]
        for m in range(T, min(T + cfg.refit_every, n_m)):
            sl = slice(p.start[m], p.start[m + 1])
            pred[sl] = feat(m) @ beta
        log(f"  pfml refit {str(p.months[T])[:7]}: lambda={lambdas[best]:.3g} "
            f"val SR={sr[best] * np.sqrt(12):.2f}")
    return pred