"""Markowitz-ML with a clustered Barra risk model: market-, beta- and industry-neutral.

Pipeline (every step uses only information dated <= the formation month t):

 1. Features. The 402 CTF characteristics are ranked cross-sectionally each month and mapped
    to [-0.5, 0.5] (missing -> 0, the median). Engineered features follow fixed rules applied to
    every input, with no hand-picking: returns over the last 1/2/3/5/10 trading days (from
    daily_ret), the share of missing characteristics, and an industry-relative version (rank
    minus FF12 industry-month mean) of every one of these 408 columns. 816 inputs in total.
 2. Target. arctan of the within-month z-score of ret_exc_lead1m: the model learns the
    cross-section (market neutral by construction) and fat tails cannot dominate the loss.
 3. Alpha. Ensemble of LightGBM regressors trained on an expanding window of all rows with
    eom < t (their labels are realised by t), refit every 12 months.
 4. Risk. Barra-style model Sigma = B F B' + D: 12 FF12 industry dummies + K style factors
    (hierarchical clusters of the characteristics, re-estimated every 60 months on the
    preceding 120 months), daily WLS factor returns, EWMA factor covariance (correlation
    half-life 504d, variance 84d), EWMA specific variance (84d).
 5. Portfolio. Book 1 (50%): w ∝ Sigma^{-1}(mu - A lambda) with A'w = 0 for
    A = [beta, industry dummies] (beta- and industry-neutral, hence dollar-neutral).
    Book 2 (50%): Ledoit-Wolf shrunk 756-day daily sample covariance plus ridge, tangency minus
    GMV (dollar-neutral). Each book is scaled to 10% ex-ante vol, blended, and the blend is
    rescaled to 10% annualised ex-ante vol under the factor model.

Run locally:  python model.py <dir with ctff_*.parquet>   (writes output.csv)

CTF Admin Modifications (2026-10-07):
--------------------------------------
1. Added requirements.txt pinning lightgbm==4.7.0.
   Reason: the script imports lightgbm, which is not pre-installed, and no dependency file was
   uploaded. The version is the current release, not necessarily the one the author used.
"""
from __future__ import annotations

import os
import time

import numpy as np
import pandas as pd
from scipy.cluster.hierarchy import fcluster, linkage
from scipy.spatial.distance import squareform

SEED = 42
N_THREADS = int(os.environ.get("CTF_THREADS", "32"))
REFIT_EVERY = 12
MIN_TRAIN_MONTHS = 36
LGB_PARAMS = dict(n_estimators=1000, learning_rate=0.01, num_leaves=255, min_child_samples=2000,
                  colsample_bytree=0.5, subsample=0.5, subsample_freq=1,
                  random_state=SEED, n_jobs=N_THREADS, deterministic=True, force_row_wise=True,
                  verbose=-1)
LGB_SEEDS = (42,)

K_STYLE = 40
CLUSTER_EVERY, CLUSTER_LOOKBACK = 60, 120
HL_COR, HL_VAR, HL_SPEC, COV_OBS, MIN_SPEC_OBS = 504, 84, 84, 2520, 63
RIDGE = 1e-4
TARGET_VOL = 0.10
HORIZONS = (1, 2, 3, 5, 10)
LW_SHARE, LW_DAYS, LW_MIN_OBS, LW_KAPPA = 0.50, 756, 250, 1.0


# ============================================================================ helpers
def log(msg):
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


def is_true(s: pd.Series) -> np.ndarray:
    if pd.api.types.is_bool_dtype(s):
        return s.fillna(False).to_numpy()
    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 cs_rank(X: np.ndarray, starts, ends) -> np.ndarray:
    """Within-month percentile rank mapped to [-0.5, 0.5]; missing -> 0."""
    R = np.zeros(X.shape, dtype=np.float32)
    for lo, hi in zip(starts, ends):
        block = pd.DataFrame(X[lo:hi])
        r = block.rank(method="average", pct=True).to_numpy()
        n_obs = np.maximum(block.notna().sum(axis=0).to_numpy(), 1)
        R[lo:hi] = np.nan_to_num(r - 0.5 / n_obs - 0.5, nan=0.0)
    return R


def ff12(sic: np.ndarray) -> np.ndarray:
    s = np.nan_to_num(sic, nan=-1).astype(int)
    rules = [
        (0, [(100, 999), (2000, 2399), (2700, 2749), (2770, 2799), (3100, 3199), (3940, 3989)]),
        (1, [(2500, 2519), (3630, 3659), (3710, 3711), (3714, 3714), (3716, 3716), (3750, 3751),
             (3792, 3792), (3900, 3939), (3990, 3999)]),
        (2, [(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, [(1200, 1399), (2900, 2999)]),
        (4, [(2800, 2829), (2840, 2899)]),
        (5, [(3570, 3579), (3660, 3692), (3694, 3699), (3810, 3829), (7370, 7379)]),
        (6, [(4800, 4899)]),
        (7, [(4900, 4949)]),
        (8, [(5000, 5999), (7200, 7299), (7600, 7699)]),
        (9, [(2830, 2839), (3693, 3693), (3840, 3859), (8000, 8099)]),
        (10, [(6000, 6999)]),
    ]
    out = np.full(len(s), 11)
    done = np.zeros(len(s), bool)
    for code, rngs in rules:
        m = np.zeros(len(s), bool)
        for lo, hi in rngs:
            m |= (s >= lo) & (s <= hi)
        m &= ~done
        out[m] = code
        done |= m
    return out


def ewma_weights(n, hl):
    w = 0.5 ** (np.arange(n)[::-1] / hl)
    return w / w.sum()


# ============================================================================ data
class Data:
    def __init__(self, chars, features, daily_ret):
        feats = sorted(features["features"].astype(str).tolist())   # canonical order (Rule 18)
        df = chars[["id", "eom", "sic", "size_grp", "ret_exc_lead1m", "ctff_test"] + feats].copy()
        df["id"] = pd.to_numeric(df["id"]).astype(np.int64)
        df["eom"] = pd.to_datetime(df["eom"]).dt.normalize()
        df = df.sort_values(["eom", "id"], kind="stable").reset_index(drop=True)
        self.feats = feats
        self.ids = df["id"].to_numpy()
        self.eom = df["eom"].to_numpy()
        self.months, self.m = np.unique(self.eom, return_inverse=True)
        n_m = len(self.months)
        self.starts = np.searchsorted(self.m, np.arange(n_m), side="left")
        self.ends = np.searchsorted(self.m, np.arange(n_m), side="right")
        self.y = pd.to_numeric(df["ret_exc_lead1m"], errors="coerce").to_numpy(np.float64)
        self.test = is_true(df["ctff_test"])
        self.size = df["size_grp"].astype(str).to_numpy()
        self.ind = ff12(pd.to_numeric(df["sic"], errors="coerce").to_numpy())
        Xraw = np.array(df[feats].apply(pd.to_numeric, errors="coerce").to_numpy(np.float32), copy=True)
        Xraw[~np.isfinite(Xraw)] = np.nan
        self.me = Xraw[:, feats.index("market_equity")].astype(np.float64)
        del df

        # dense daily matrix for panel ids
        d = daily_ret[["id", "date", "ret_exc"]]
        rid = pd.to_numeric(d["id"]).to_numpy().astype(np.int64)
        keep = np.isin(rid, np.unique(self.ids))
        rdate = pd.to_datetime(d["date"]).to_numpy()[keep].astype("datetime64[D]")
        rret = pd.to_numeric(d["ret_exc"]).to_numpy()[keep].astype(np.float32)
        self.dates, di = np.unique(rdate, return_inverse=True)
        self.dids, ii = np.unique(rid[keep], return_inverse=True)
        self.M = np.full((len(self.dates), len(self.dids)), np.nan, dtype=np.float32)
        self.M[di, ii] = rret
        self.col = np.searchsorted(self.dids, self.ids).clip(0, max(len(self.dids) - 1, 0))
        self.has = (self.dids[self.col] == self.ids) if len(self.dids) else np.zeros(len(self.ids), bool)
        self.d_end = np.searchsorted(self.dates, self.months.astype("datetime64[D]"), side="right")
        log(f"panel {len(self.ids):,} rows x {len(feats)} feats, {n_m} months; "
            f"daily {self.M.shape}")

        Xr = cs_rank(Xraw, self.starts, self.ends)
        self.X = np.hstack([Xr, self._engineered(Xraw, Xr)])
        self.Xr = self.X[:, :len(feats)]
        del Xr
        del Xraw
        log(f"features ready {self.X.shape}")

    def rows(self, t):
        return np.arange(self.starts[t], self.ends[t])

    def _engineered(self, Xraw, Xr):
        """Rule-based features (no hand-picked anomalies, Rule 2):
        R1 cumulative excess return over the last h trading days for every h in HORIZONS,
        R2 share of missing characteristics, R3 industry-relative version (rank minus FF12
        industry-month mean rank) of every base column: 402 characteristics, R1 and R2."""
        n = len(self.ids)
        H = max(HORIZONS)
        R1 = np.full((n, len(HORIZONS)), np.nan, dtype=np.float32)
        for t in range(len(self.months)):
            rs = self.rows(t)
            ok = self.has[rs]
            end = self.d_end[t]
            if end < H or not ok.any():
                continue
            lr = np.log1p(self.M[end - H:end][:, self.col[rs][ok]].astype(np.float64))
            for j, h in enumerate(HORIZONS):
                b = lr[-h:]
                v = np.where(np.isfinite(b).sum(0) >= max(1, h - 1), np.nansum(b, 0), np.nan)
                R1[rs[ok], j] = np.expm1(v)
        R2 = np.isnan(Xraw).mean(1, keepdims=True).astype(np.float32)
        own = cs_rank(np.hstack([R1, R2]), self.starts, self.ends)
        base = np.hstack([Xr, own])
        key = self.m.astype(np.int64) * 100 + self.ind
        order = np.argsort(key, kind="stable")
        starts = np.r_[0, np.flatnonzero(np.diff(key[order])) + 1]
        counts = np.diff(np.r_[starts, len(order)])
        sums = np.add.reduceat(base[order].astype(np.float64), starts, axis=0)
        means = (sums / counts[:, None]).astype(np.float32)
        grp = np.empty(n, np.int64)
        grp[order] = np.repeat(np.arange(len(starts)), counts)
        rel = base - means[grp]
        return np.hstack([own, rel])


# ============================================================================ alpha
def target(D: Data) -> np.ndarray:
    g = pd.Series(D.y).groupby(D.m)
    z = ((D.y - g.transform("mean")) / g.transform("std")).to_numpy()
    return np.arctan(z).astype(np.float32)


def alpha_predictions(D: Data, traded: list[int]) -> np.ndarray:
    import lightgbm as lgb
    y = target(D)
    pred = np.full(len(D.ids), np.nan, dtype=np.float32)
    models = None
    for k, t in enumerate(traded):
        if k % REFIT_EVERY == 0:
            hi = int(D.starts[t])                       # rows with eom < t
            ok = np.isfinite(y[:hi])
            tic = time.time()
            models = []
            for s in LGB_SEEDS:
                mdl = lgb.LGBMRegressor(**dict(LGB_PARAMS, random_state=s))
                mdl.fit(D.X[:hi][ok], y[:hi][ok])
                models.append(mdl)
            log(f"alpha refit {str(D.months[t])[:7]} on {ok.sum():,} rows ({time.time()-tic:.0f}s)")
        rs = D.rows(t)
        p = np.mean([mdl.booster_.predict(D.X[rs]) for mdl in models], axis=0)
        pred[rs] = ((p - p.mean()) / (p.std() + 1e-12)).astype(np.float32)
    return pred


# ============================================================================ risk
def cluster_features(D: Data, t0: int, k: int):
    lo = max(0, t0 - CLUSTER_LOOKBACK)
    ms = np.arange(lo, max(t0, lo + 1), 6)
    p = D.Xr.shape[1]
    C = np.zeros((p, p))
    for m in ms:
        Z = D.Xr[D.rows(m)].astype(np.float64)
        Z = Z - Z.mean(0)
        sd = Z.std(0)
        sd[sd == 0] = np.inf
        Z = Z / sd
        C += Z.T @ Z / len(Z)
    C /= len(ms)
    np.fill_diagonal(C, 1.0)
    dist = np.clip(1 - np.abs(C), 0, 2)
    np.fill_diagonal(dist, 0)
    lab = fcluster(linkage(squareform(dist, checks=False), "average"), k, criterion="maxclust") - 1
    sign = np.ones(p)
    for c in np.unique(lab):
        j = np.flatnonzero(lab == c)
        if len(j) > 1:
            _, evec = np.linalg.eigh(C[np.ix_(j, j)])
            v = evec[:, -1]
            sign[j] = np.sign(v + 1e-12) * np.sign(v[np.argmax(np.abs(v))])
    return lab, sign


def exposures(D: Data, rs, lab, sign):
    Z = D.Xr[rs].astype(np.float64) * sign
    S = np.zeros((len(rs), K_STYLE))
    for c in range(K_STYLE):
        j = lab == c
        if j.any():
            S[:, c] = Z[:, j].mean(1)
    sd = S.std(0)
    S = (S - S.mean(0)) / np.where(sd > 0, sd, 1)
    I = np.zeros((len(rs), 12))
    I[np.arange(len(rs)), D.ind[rs]] = 1.0
    return np.column_stack([I, S])


class Risk:
    def __init__(self, D: Data):
        self.D = D
        self.period = None

    def _build(self, t0, t1):
        D = self.D
        n_m = len(D.months)
        lab, sign = cluster_features(D, t0, K_STYLE)
        first_day = max(0, D.d_end[t0] - COV_OBS - 30)
        first_m = max(0, int(np.searchsorted(D.d_end, first_day, side="right")) - 1)
        last_day = D.d_end[min(t1, n_m - 1)]
        days, fr, rr, rc, rv = [], [], [], [], []
        for s in range(first_m, min(t1, n_m - 1)):
            rs = D.rows(s)
            rs = rs[D.has[rs]]
            if len(rs) < 30:
                continue
            B = exposures(D, rs, lab, sign)
            me = D.me[rs]
            med = np.nanmedian(me) if np.isfinite(me).any() else 1.0
            wts = np.sqrt(np.nan_to_num(me, nan=med).clip(1e-6))
            cols = D.col[rs]
            for d in range(D.d_end[s], min(D.d_end[s + 1], last_day)):
                r = D.M[d, cols].astype(np.float64)
                ok = np.isfinite(r)
                if ok.sum() < 30:
                    continue
                Bo, wo, ro = B[ok], wts[ok], r[ok]
                A = Bo.T @ (wo[:, None] * Bo)
                A[np.diag_indices_from(A)] += RIDGE * np.trace(A) / len(A)
                f = np.linalg.solve(A, Bo.T @ (wo * ro))
                days.append(d)
                fr.append(f)
                rr.append(np.full(ok.sum(), d, dtype=np.int64))
                rc.append(cols[ok])
                rv.append((ro - Bo @ f).astype(np.float32))
        days = np.array(days, dtype=np.int64)
        E = np.full((len(days), len(D.dids)), np.nan, dtype=np.float32)
        if len(days):
            E[np.searchsorted(days, np.concatenate(rr)), np.concatenate(rc)] = np.concatenate(rv)
        self.period = dict(t0=t0, t1=t1, lab=lab, sign=sign, days=days, FR=np.array(fr), E=E)

    def month(self, t, anchor):
        if self.period is None or not (self.period["t0"] <= t < self.period["t1"]):
            t0 = anchor + ((t - anchor) // CLUSTER_EVERY) * CLUSTER_EVERY
            self._build(t0, t0 + CLUSTER_EVERY)
        D, p = self.D, self.period
        sel = np.flatnonzero(p["days"] < D.d_end[t])[-COV_OBS:]
        FR = p["FR"][sel]
        wc, wv = ewma_weights(len(sel), HL_COR), ewma_weights(len(sel), HL_VAR)
        Fc = FR - (wc[:, None] * FR).sum(0)
        C = (wc[:, None] * Fc).T @ Fc
        sd_c = np.sqrt(np.maximum(np.diag(C), 1e-20))
        Fv = FR - (wv[:, None] * FR).sum(0)
        sd_v = np.sqrt((wv[:, None] * Fv ** 2).sum(0))
        F = C / np.outer(sd_c, sd_c) * np.outer(sd_v, sd_v) * 21

        rs = D.rows(t)
        B = exposures(D, rs, p["lab"], p["sign"])
        Eh = p["E"][sel[-252:]][:, D.col[rs]]
        Eh = np.where(D.has[rs][None, :], Eh, np.nan)
        ok = np.isfinite(Eh)
        ws = ewma_weights(len(Eh), HL_SPEC)[:, None] * ok
        spec = (ws * np.nan_to_num(Eh) ** 2).sum(0) / np.maximum(ws.sum(0), 1e-12) * 21
        good = ok.sum(0) >= MIN_SPEC_OBS
        Dv = spec.copy()
        base = np.median(spec[good]) if good.any() else 0.01
        for g in np.unique(D.size[rs]):
            m = D.size[rs] == g
            Dv[m & ~good] = 1.5 * (np.median(spec[m & good]) if (m & good).any() else base)
        Dv = np.maximum(Dv, np.quantile(Dv[good], 0.05) if good.any() else 1e-4)
        return B, F, Dv


def woodbury_solve(B, F, Dv, V):
    Di = 1.0 / Dv
    DV = Di[:, None] * V
    inner = np.linalg.solve(np.linalg.pinv(F) + B.T @ (Di[:, None] * B), B.T @ DV)
    return DV - Di[:, None] * (B @ inner)


def port_var(B, F, Dv, w):
    f = B.T @ w
    return float(f @ F @ f + (w * w) @ Dv)


def markowitz_weights(B, F, Dv, mu, ind, me):
    n = len(mu)
    wm = np.nan_to_num(me, nan=0.0)
    wm = wm / wm.sum() if wm.sum() > 0 else np.ones(n) / n
    cov_m = B @ (F @ (B.T @ wm)) + Dv * wm
    beta = cov_m / port_var(B, F, Dv, wm)
    I = np.zeros((n, 12))
    I[np.arange(n), ind] = 1.0
    A = np.column_stack([beta, I[:, I.sum(0) > 0]])
    S = woodbury_solve(B, F, Dv, np.column_stack([mu, A]))
    Smu, SA = S[:, 0], S[:, 1:]
    lam = np.linalg.lstsq(A.T @ SA, A.T @ Smu, rcond=None)[0]
    w = Smu - SA @ lam
    return w * (TARGET_VOL / np.sqrt(12)) / np.sqrt(max(port_var(B, F, Dv, w), 1e-18))


def lw_weights(D: Data, t, mu):
    """Ledoit-Wolf daily covariance (756d ending at eom_t) + ridge; tangency minus GMV."""
    rs = D.rows(t)
    w = np.zeros(len(rs))
    ok = D.has[rs]
    end = D.d_end[t]
    X = D.M[max(0, end - LW_DAYS):end][:, D.col[rs][ok]].astype(np.float64)
    obs = np.isfinite(X)
    enough = obs.sum(0) >= min(LW_MIN_OBS, max(len(X) // 3, 20))
    idx = np.flatnonzero(ok)[enough]
    if len(idx) < 50:
        return w
    X = np.where(obs, X, 0.0)[:, enough]
    X = (X - X.mean(0)) * np.sqrt(21)
    T, n = X.shape
    S = X.T @ X / T
    m = np.trace(S) / n
    d2 = (S ** 2).sum() - 2 * m * np.trace(S) + n * m ** 2
    b2 = ((X ** 2).sum(1) ** 2).sum() / T ** 2 - (S ** 2).sum() / T
    delta = min(max(b2, 0.0), d2) / d2 if d2 > 0 else 1.0
    S *= (1 - delta)
    S[np.diag_indices(n)] += delta * m + LW_KAPPA * np.trace(S) / n
    sol = np.linalg.solve(S, np.column_stack([mu[idx], np.ones(n)]))
    a, b = sol[:, 0], sol[:, 1]
    wl = a - a.sum() / b.sum() * b
    w[idx] = wl * (TARGET_VOL / np.sqrt(12)) / np.sqrt(max(wl @ S @ wl, 1e-18))
    return w


# ============================================================================ main
def main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    np.random.seed(SEED)
    t_start = time.time()
    D = Data(chars, features, daily_ret)

    test_months = np.unique(D.m[D.test])
    first = int(test_months[0]) if len(test_months) else MIN_TRAIN_MONTHS
    first = max(first, MIN_TRAIN_MONTHS)
    traded = list(range(first, len(D.months)))
    assert traded, "not enough months to train"
    log(f"trading {len(traded)} months {str(D.months[traded[0]])[:10]} .. {str(D.months[traded[-1]])[:10]}")

    mu = alpha_predictions(D, traded)
    risk = Risk(D)
    W = np.zeros(len(D.ids))
    for t in traded:
        rs = D.rows(t)
        B, F, Dv = risk.month(t, anchor=traded[0])
        mut = np.nan_to_num(mu[rs])
        w = (1 - LW_SHARE) * markowitz_weights(B, F, Dv, mut, D.ind[rs], D.me[rs])
        if LW_SHARE > 0:
            w = w + LW_SHARE * lw_weights(D, t, mut)
        W[rs] = w * (TARGET_VOL / np.sqrt(12)) / np.sqrt(max(port_var(B, F, Dv, w), 1e-18))
        if (t - traded[0]) % 60 == 0:
            log(f"portfolio {str(D.months[t])[:7]}: n={len(rs)} gross={np.abs(W[rs]).sum():.2f}")

    out = pd.DataFrame({"id": D.ids, "eom": pd.to_datetime(D.eom).date, "w": W})
    keep = np.zeros(len(D.ids), bool)
    keep[np.concatenate([D.rows(t) for t in traded])] = True
    if D.test.any():
        keep &= D.test
    out = out[keep].reset_index(drop=True)
    out["id"] = out["id"].astype(np.int64)
    out["w"] = out["w"].astype(np.float64)
    assert len(out) and out.notna().all().all()
    log(f"done: {len(out):,} rows in {(time.time() - t_start) / 60:.1f} min")
    return out[["id", "eom", "w"]]


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

    src = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(".")
    chars = pd.read_parquet(src / "ctff_chars.parquet")
    features = pd.read_parquet(src / "ctff_features.parquet")
    daily_ret = pd.read_parquet(src / "ctff_daily_ret.parquet")
    pf = main(chars, features, daily_ret)
    pf.to_csv("output.csv", index=False)