"""Hierarchical Committee with a Self-Auditing Portfolio: CTF submission (MGT 924, PS3).

Version: ashley/M8 (parent ashley/M7). Theme: "carry the forecast layer's thesis into the portfolio".
Pre-registered in experiments/configs/ashley/M8_preregistration.md. M7's forecast layer is frozen
as tested (variant vHE: pooled vintages, size-group corrections by partial pooling, evidence
weights, disagreement shrinkage, industry neutralization). The portfolio adds:
    P1 uncertainty-aware risk: Sigma* = 21 Sigma + Omega, where Omega is the variance of the mean
       of the committee members (vintages x base/corrected), in Grinold-calibrated return units
    P2 partially pooled covariance: Ledoit-Wolf shrinkage toward a size-block constant-correlation
       target
    P3 evidence-based exposure timing: monthly scale from the fundamental-law ratio
       alpha'w / w'Sigma* w relative to its trailing median, shrunk halfway, clipped to [0.5, 2]
Variants: v0 (= M7 vHE), vU, vK, vT, vUK, vUKT. Placebo variants "<variant>_P<seed>" use random
size blocks in P2 (design check only). CONFIG["variant"] selects the submitted variant by the
pre-registered rule using 1975-1989 data only.

Rules this file follows (jkpfactors.com/ctf/rules):
- `main(chars, features, daily_ret)` returns id (int), eom (date), w (float), with no missing
  values (Rules 11-12).
- Every quantity used at month t is known at t (Rule 1). Every vintage is trained on months
  strictly before its refit month. Refit calendars are anchored at the start of the data. Daily
  windows end at t.
- Only provided data. No feature or stock is chosen by hand. The FF12 map is the fixed public
  Fama-French definition (Rules 2-3).
- Fixed seed and thread count, deterministic algorithms (Rule 18). No network access, no shell,
  no dynamic code execution, no file I/O (Rules 10, 15).
"""
from __future__ import annotations

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

CONFIG = {
    "seed": 42,
    "n_threads": 12,              # Ashley allowed 12 for M3 (2026-10-03); weights are thread-invariant
    "min_train_months": 36,
    "refit_every": 36,
    "lgbm": {
        "n_estimators": 200,
        "learning_rate": 0.05,
        "num_leaves": 255,
        "min_child_samples": 2000,
        "max_bin": 63,
    },
    "cov_window_days": 756,
    "min_daily_obs": 250,
    "min_names": 50,
    "vintage_offsets": [0, 12, 24],   # months; vintage ensemble (E)
    "tau2_floor": 0.05,               # disagreement shrinkage (D): floor on signal-variance share
    "trust_window": 60,               # months of out-of-sample IC behind the trust weights (T)
    "ic_prior": 0.05,                 # Grinold calibration: prior IC, weighted as 12 months of evidence
    "variant": "vUKT",                # selected by the pre-registered rule on 1975-1989
}
DESIGN_VARIANTS = ["v0", "vU", "vK", "vT", "vUK", "vUKT"]
FORECAST_VARIANT = "vHE"              # M7's frozen forecast layer
GROUPS = ("small", "large", "mega")   # CTF size groups; any other label falls back to the pooled view
MIN_FORECAST_SPREAD = 1e-6   # below this, forecasts are treated as uninformative (hold nothing)
TARGET = "ret_exc_lead1m"

# Fama-French 12 industries (Ken French Data Library, Siccodes12.txt; verified 2026-10-02).
FF12_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)],
}


# ----------------------------------------------------------------------------------------------
# Stage 0: load and normalise
# ----------------------------------------------------------------------------------------------
def to_int_ids(values) -> np.ndarray:
    """Security ids as int64. Accepts ints, floats with integral values, or Decimal objects."""
    arr = np.asarray(values)
    if arr.dtype == object:
        arr = arr.astype(np.int64)
    return arr.astype(np.int64)


def is_test_row(flag: pd.Series) -> np.ndarray:
    """ctff_test arrives as bool, 0/1, or '0'/'1' depending on the source."""
    if pd.api.types.is_bool_dtype(flag):
        return flag.fillna(False).to_numpy()
    if pd.api.types.is_numeric_dtype(flag):
        return flag.fillna(0).ne(0).to_numpy()
    return flag.astype(str).str.strip().str.lower().isin(["1", "1.0", "true", "t"]).to_numpy()


def ff12(sic: np.ndarray) -> np.ndarray:
    """Fama-French 12-industry code (1..12) for each SIC code; missing SIC maps to 12 (Other)."""
    out = np.full(len(sic), 12, dtype=np.int8)
    for ind, ranges in FF12_RANGES.items():
        for lo, hi in ranges:
            out[(sic >= lo) & (sic <= hi)] = ind
    return out


def prepare_panel(chars: pd.DataFrame, feats: list[str]) -> dict:
    """Sorted panel arrays. Rows are ordered by (eom, id) so every month is one contiguous slice."""
    ids = to_int_ids(chars["id"].to_numpy())
    eom = pd.to_datetime(chars["eom"]).to_numpy().astype("datetime64[D]")
    order = np.lexsort((ids, eom))
    months, month_ix = np.unique(eom[order], return_inverse=True)
    starts = np.searchsorted(month_ix, np.arange(len(months)), side="left")
    ends = np.searchsorted(month_ix, np.arange(len(months)), side="right")
    X = np.empty((len(order), len(feats)), dtype=np.float32)
    for j, f in enumerate(feats):
        X[:, j] = chars[f].to_numpy(dtype=np.float32)[order]
    test = (is_test_row(chars["ctff_test"])[order] if "ctff_test" in chars.columns
            else np.zeros(len(order), dtype=bool))
    sic = pd.to_numeric(chars["sic"], errors="coerce").to_numpy()[order]
    grp = np.full(len(order), -1, dtype=np.int8)
    if "size_grp" in chars.columns:
        label = chars["size_grp"].astype(str).str.strip().str.lower().to_numpy()[order]
        for j, g in enumerate(GROUPS):
            grp[label == g] = j
    return {
        "grp": grp,
        "ids": ids[order],
        "eom_out": chars["eom"].to_numpy()[order],
        "months": months, "month_ix": month_ix, "starts": starts, "ends": ends,
        "X": X,
        "r": chars[TARGET].to_numpy(dtype=np.float64)[order],
        "ind": ff12(np.nan_to_num(sic, nan=-1)),
        "test": test,
    }


# ----------------------------------------------------------------------------------------------
# Stage 1: features and target (unchanged from M0)
# ----------------------------------------------------------------------------------------------
def rank_within_month(P: dict) -> None:
    """Replace X in place by within-month percentile ranks on (0, 1]. Missing stays missing."""
    X = P["X"]
    for lo, hi in zip(P["starts"], P["ends"]):
        X[lo:hi] = pd.DataFrame(X[lo:hi]).rank(axis=0, pct=True, method="average",
                                                na_option="keep").to_numpy(dtype=np.float32)


def training_target(P: dict) -> np.ndarray:
    """arctan of the next-month return standardised within its month (Geertsema's transform)."""
    g = pd.DataFrame({"m": P["month_ix"], "r": P["r"]}).groupby("m")["r"]
    z = (P["r"] - g.transform("mean").to_numpy()) / g.transform("std").to_numpy()
    return np.arctan(z).astype(np.float32)


# ----------------------------------------------------------------------------------------------
# Stages 4-5: expert fit and forecasts (unchanged from M0)
# ----------------------------------------------------------------------------------------------
def output_months(P: dict, cfg: dict) -> np.ndarray:
    """Month indices that receive weights: the ctff_test months, or every month after warm-up."""
    if P["test"].any():
        out = np.unique(P["month_ix"][P["test"]])
    else:
        out = np.arange(cfg["min_train_months"], len(P["months"]))
    out = out[out >= cfg["min_train_months"]]
    assert len(out), "not enough months to train before the first output month"
    return out


def forecast(P: dict, y: np.ndarray, out: np.ndarray, cfg: dict, log=print) -> np.ndarray:
    """Expanding-window LightGBM, refit on a calendar anchored at the start of the data."""
    grid = np.arange(cfg["min_train_months"], len(P["months"]), cfg["refit_every"])
    model_for = grid[np.searchsorted(grid, out, side="right") - 1]
    mu = np.full(len(P["ids"]), np.nan, dtype=np.float32)
    params = {**cfg["lgbm"], "random_state": cfg["seed"], "n_jobs": cfg["n_threads"],
              "deterministic": True, "force_row_wise": True, "verbose": -1}
    for g in np.unique(model_for):
        fit_rows = int(P["starts"][g])
        keep = np.isfinite(y[:fit_rows])
        model = lgb.LGBMRegressor(**params)
        model.fit(P["X"][:fit_rows][keep], y[:fit_rows][keep])
        for t in out[model_for == g]:
            lo, hi = int(P["starts"][t]), int(P["ends"][t])
            mu[lo:hi] = model.booster_.predict(P["X"][lo:hi])
        log(f"  refit {P['months'][g]}  trained on {int(keep.sum()):,} rows  "
            f"-> {int((model_for == g).sum())} months")
    return mu


def industry_neutral(mu: np.ndarray, ind: np.ndarray) -> np.ndarray:
    """Subtract each industry's mean forecast (variant B)."""
    out = mu.astype(np.float64).copy()
    for g in np.unique(ind):
        sel = ind == g
        out[sel] -= out[sel].mean()
    return out


# ----------------------------------------------------------------------------------------------
# Stage 6: covariance from daily returns
# ----------------------------------------------------------------------------------------------
def daily_matrix(daily_ret: pd.DataFrame, panel_ids: np.ndarray, first_date) -> dict:
    """Dense (days x ids) float32 matrix of daily excess returns, NaN where a stock did not trade."""
    rid = to_int_ids(daily_ret["id"].to_numpy())
    rdate = pd.to_datetime(daily_ret["date"]).to_numpy().astype("datetime64[D]")
    keep = np.isin(rid, panel_ids) & (rdate >= first_date)
    dates, d_ix = np.unique(rdate[keep], return_inverse=True)
    uids, i_ix = np.unique(rid[keep], return_inverse=True)
    M = np.full((len(dates), len(uids)), np.nan, dtype=np.float32)
    M[d_ix, i_ix] = daily_ret["ret_exc"].to_numpy(dtype=np.float32)[keep]
    return {"dates": dates, "ids": uids, "M": M}


def lw_constant_correlation(X: np.ndarray) -> tuple[np.ndarray, float]:
    """Ledoit & Wolf (2004) shrinkage toward constant correlation (M0, unchanged)."""
    T, n = X.shape
    S = X.T @ X / T
    var = np.diag(S).copy()
    sd = np.sqrt(np.maximum(var, 1e-20))
    rbar = float(((S / np.outer(sd, sd)).sum() - n) / (n * (n - 1)))
    F = rbar * np.outer(sd, sd)
    np.fill_diagonal(F, var)
    Y = X * X
    pi_mat = Y.T @ Y / T - S * S
    theta = (Y * X).T @ X / T - var[:, None] * S
    off = theta * (sd[None, :] / sd[:, None])
    np.fill_diagonal(off, 0.0)
    pi_hat = float(pi_mat.sum())
    rho_hat = float(np.trace(pi_mat) + rbar * off.sum())
    gamma_hat = float(((F - S) ** 2).sum())
    delta = min(max((pi_hat - rho_hat) / gamma_hat / T, 0.0), 1.0) if gamma_hat > 0 else 1.0
    return delta * F + (1.0 - delta) * S, delta


def month_window(D: dict, ids: np.ndarray, month_end, cfg: dict):
    """Demeaned trailing daily returns for one month's stocks; stocks short of history are not held."""
    end = int(np.searchsorted(D["dates"], month_end, side="right"))
    lo = max(0, end - cfg["cov_window_days"])
    col = np.clip(np.searchsorted(D["ids"], ids), 0, len(D["ids"]) - 1)
    has = D["ids"][col] == ids
    W = D["M"][lo:end][:, col[has]].astype(np.float64)
    obs = np.isfinite(W)
    enough = obs.sum(axis=0) >= cfg["min_daily_obs"]
    usable = np.zeros(len(ids), dtype=bool)
    usable[np.flatnonzero(has)[enough]] = True
    W = np.where(obs, W, 0.0)[:, enough]
    W -= W.mean(axis=0, keepdims=True)
    return usable, W


# ----------------------------------------------------------------------------------------------
# Stage 7: portfolios and the vC blend
# ----------------------------------------------------------------------------------------------
def tangency_minus_gmv(mu: np.ndarray, Sigma: np.ndarray) -> np.ndarray:
    """Dollar-neutral mean-variance weights at unit gross exposure (M0 rule, ridge off).

    Uninformative forecasts (spread below MIN_FORECAST_SPREAD) give zero weights, not 0/0.
    """
    n = len(mu)
    if not np.isfinite(mu).all() or np.std(mu) < MIN_FORECAST_SPREAD:
        return np.zeros(n)
    sol = np.linalg.solve(Sigma, np.column_stack([mu.astype(np.float64), np.ones(n)]))
    a, b = sol[:, 0], sol[:, 1]
    w = a - (a.sum() / b.sum()) * b
    gross = np.abs(w).sum()
    return w / gross if gross > 0 and np.isfinite(gross) else np.zeros(n)



# ----------------------------------------------------------------------------------------------
# M3 additions: vintage ensemble (E), disagreement shrinkage (D), missing-day correction (Z)
# ----------------------------------------------------------------------------------------------
def forecast_vintages(P: dict, y: np.ndarray, months: np.ndarray, cfg: dict, log=print) -> dict:
    """One LightGBM forecast array per refit-calendar offset (NaN where a vintage has no model yet).

    Vintage k refits every `refit_every` months starting at month index min_train_months + k,
    each time on all months strictly before the refit month.
    """
    params = {**cfg["lgbm"], "random_state": cfg["seed"], "n_jobs": cfg["n_threads"],
              "deterministic": True, "force_row_wise": True, "verbose": -1}
    out = {}
    for k in cfg["vintage_offsets"]:
        mu = np.full(len(P["ids"]), np.nan, dtype=np.float32)
        grid = np.arange(cfg["min_train_months"] + k, len(P["months"]), cfg["refit_every"])
        mk = months[months >= grid[0]] if len(grid) else months[:0]
        if len(mk):
            model_for = grid[np.searchsorted(grid, mk, side="right") - 1]
            for g in np.unique(model_for):
                fit_rows = int(P["starts"][g])
                keep = np.isfinite(y[:fit_rows])
                model = lgb.LGBMRegressor(**params)
                model.fit(P["X"][:fit_rows][keep], y[:fit_rows][keep])
                for t in mk[model_for == g]:
                    lo, hi = int(P["starts"][t]), int(P["ends"][t])
                    mu[lo:hi] = model.booster_.predict(P["X"][lo:hi])
                log(f"  vintage +{k:>2d}m refit {P['months'][g]}  trained on {int(keep.sum()):,} rows")
        out[k] = mu
    return out


def zscore(x: np.ndarray) -> np.ndarray:
    sd = x.std()
    return (x - x.mean()) / sd if sd > 0 else np.zeros_like(x)


def ensemble_forecast(stack: list, shrink: bool, floor: float) -> tuple[np.ndarray, float]:
    """Average of within-month z-scored vintage forecasts; optionally shrink by disagreement.

    Returns the forecast and the average shrink factor (1.0 when no shrinkage is applied).
    """
    Z = np.vstack([zscore(np.asarray(s, dtype=np.float64)) for s in stack])
    K = Z.shape[0]
    m = Z.mean(axis=0)
    if not shrink or K < 2:
        return m, 1.0
    noise = Z.var(axis=0, ddof=1) / K                     # variance of the mean across vintages
    total = m.var()
    tau2 = max(total - noise.mean(), floor * total)        # cross-sectional signal variance
    factor = tau2 / (tau2 + noise) if tau2 > 0 else np.zeros_like(m)
    return m * factor, float(np.mean(factor))


def month_window_both(D: dict, ids: np.ndarray, month_end, cfg: dict):
    """Demeaned trailing daily returns, as in M2.1, plus the missing-day corrected version (Z).

    The correction rescales each column so its variance equals the variance over observed days.
    Rescaling a column leaves its correlations with the other columns unchanged.
    """
    end = int(np.searchsorted(D["dates"], month_end, side="right"))
    lo = max(0, end - cfg["cov_window_days"])
    col = np.clip(np.searchsorted(D["ids"], ids), 0, len(D["ids"]) - 1)
    has = D["ids"][col] == ids
    W = D["M"][lo:end][:, col[has]].astype(np.float64)
    obs = np.isfinite(W)
    enough = obs.sum(axis=0) >= cfg["min_daily_obs"]
    usable = np.zeros(len(ids), dtype=bool)
    usable[np.flatnonzero(has)[enough]] = True
    W_obs = W[:, enough]
    W = np.where(obs, W, 0.0)[:, enough]
    W -= W.mean(axis=0, keepdims=True)
    var_zero = W.var(axis=0)
    var_obs = np.nanvar(W_obs, axis=0)
    scale = np.sqrt(np.where(var_zero > 0, var_obs / np.maximum(var_zero, 1e-300), 1.0))
    return usable, W, W * scale


# ----------------------------------------------------------------------------------------------
# M7 additions: group corrections by partial pooling (H) and evidence weights (E)
# ----------------------------------------------------------------------------------------------
# variant -> (extra trees E, lambda rule). E = 0 reproduces M3's vZED exactly.
SPEC = {"v0": (0, None), "vH25": (25, "full"), "vH100": (100, "full"), "vHE": (100, "evidence")}


def parse_variant(v: str) -> tuple[str, int]:
    """'vHE' -> ('vHE', 0); 'vHE_P07' -> ('vHE', 7). Seed 0 means the real size groups."""
    base, _, tag = v.partition("_P")
    assert base in SPEC, f"unknown variant {v}"
    return base, (int(tag) if tag else 0)


def placebo_groups(P: dict, seed: int) -> np.ndarray:
    """Random groups with the real group counts: each month's labels are shuffled across its stocks.

    Months are drawn in order from one seeded generator, so truncating the data at a later month
    leaves every earlier month's draw unchanged.
    """
    rng = np.random.default_rng(seed)
    g = P["grp"].copy()
    for lo, hi in zip(P["starts"], P["ends"]):
        g[lo:hi] = rng.permutation(g[lo:hi])
    return g


def forecast_hierarchy(P, y, months, cfg, group_sets: dict, extras: list, log=print):
    """Base vintages (M3) plus, per group set and E, forecasts with group corrections.

    At each refit the base LightGBM is fitted exactly as in M3. For each size group, a copy keeps
    training from the base booster for E extra trees on that group's rows (same training-row rule
    and settings). Groups too small to split get no correction and keep the base forecast.
    Returns base[k] and corr[(seed, E)][k], each a forecast array over all rows.
    """
    params = {**cfg["lgbm"], "random_state": cfg["seed"], "n_jobs": cfg["n_threads"],
              "deterministic": True, "force_row_wise": True, "verbose": -1}
    base, corr = {}, {(s, e): {} for s in group_sets for e in extras}
    for k in cfg["vintage_offsets"]:
        mu = np.full(len(P["ids"]), np.nan, dtype=np.float32)
        cm = {key: np.full(len(P["ids"]), np.nan, dtype=np.float32) for key in corr}
        grid = np.arange(cfg["min_train_months"] + k, len(P["months"]), cfg["refit_every"])
        mk = months[months >= grid[0]] if len(grid) else months[:0]
        if len(mk):
            model_for = grid[np.searchsorted(grid, mk, side="right") - 1]
            for g in np.unique(model_for):
                fit_rows = int(P["starts"][g])
                keep = np.isfinite(y[:fit_rows])
                model = lgb.LGBMRegressor(**params)
                model.fit(P["X"][:fit_rows][keep], y[:fit_rows][keep])
                served = mk[model_for == g]
                rows = np.concatenate([np.arange(P["starts"][t], P["ends"][t]) for t in served])
                mu[rows] = model.booster_.predict(P["X"][rows])
                for key in corr:
                    cm[key][rows] = mu[rows]
                for seed, grp in group_sets.items():
                    for j in range(len(GROUPS)):
                        gk = (grp[:fit_rows] == j) & keep
                        sel = rows[grp[rows] == j]
                        if gk.sum() < 2 * cfg["lgbm"]["min_child_samples"] or not len(sel):
                            continue
                        for e in extras:
                            m2 = lgb.LGBMRegressor(**{**params, "n_estimators": e})
                            m2.fit(P["X"][:fit_rows][gk], y[:fit_rows][gk], init_model=model.booster_)
                            cm[(seed, e)][sel] = m2.booster_.predict(P["X"][sel])
                log(f"  vintage +{k:>2d}m refit {P['months'][g]}  trained on {int(keep.sum()):,} rows"
                    f"  + corrections for {len(group_sets)} group set(s)")
        base[k] = mu
        for key in corr:
            corr[key][k] = cm[key]
    return base, corr


def spearman(a: np.ndarray, b: np.ndarray) -> float:
    if len(a) < 3:
        return np.nan
    ra = pd.Series(a).rank().to_numpy()
    rb = pd.Series(b).rank().to_numpy()
    if ra.std() == 0 or rb.std() == 0:
        return np.nan
    return float(np.corrcoef(ra, rb)[0, 1])


def ic_history(P, base, corrected, grp, months, cfg) -> np.ndarray:
    """Monthly rank IC per size group: array (n_months, 2, n_groups); view 0 corrected, 1 base.

    Each view is the average of its within-month z-scored vintages over rows where every
    available base vintage exists and the target is realised. A month's IC is only read at
    later months (see trust_weights).
    """
    ic = np.full((len(P["months"]), 2, len(GROUPS)), np.nan)
    for t in months:
        lo, hi = int(P["starts"][t]), int(P["ends"][t])
        avail = [k for k in cfg["vintage_offsets"] if np.isfinite(base[k][lo:hi]).any()]
        if not avail:
            continue
        rows = np.ones(hi - lo, dtype=bool)
        for k in avail:
            rows &= np.isfinite(base[k][lo:hi])
        rows &= np.isfinite(P["r"][lo:hi])
        if rows.sum() < 3:
            continue
        vb = np.mean([zscore(base[k][lo:hi][rows].astype(np.float64)) for k in avail], axis=0)
        vc = np.mean([zscore(corrected[k][lo:hi][rows].astype(np.float64)) for k in avail], axis=0)
        g, r = grp[lo:hi][rows], P["r"][lo:hi][rows]
        for j in range(len(GROUPS)):
            sel = g == j
            ic[t, 0, j] = spearman(vc[sel], r[sel])
            ic[t, 1, j] = spearman(vb[sel], r[sel])
    return ic


def trust_weights(ic: np.ndarray, t: int, window: int) -> np.ndarray:
    """Weight on view 0 per size group at month t, from ICs of months strictly before t (M6 rule).

    The target of month m is realised at m + 1, so months m <= t - 1 are known at t. With fewer
    than `window` months of IC, or no positive IC, the weight is 1/2.
    """
    w = np.full(len(GROUPS), 0.5)
    for j in range(len(GROUPS)):
        past = ic[:t, :, j]
        past = past[np.isfinite(past).all(axis=1)][-window:]
        if len(past) < window:
            continue
        a, b = (max(float(x), 0.0) for x in past.mean(axis=0))
        if a + b > 0:
            w[j] = 0.5 * a / (a + b) + 0.25
    return w


# ----------------------------------------------------------------------------------------------
# M8 additions: portfolio layer (P1 uncertainty-aware risk, P2 block target, P3 exposure timing)
# ----------------------------------------------------------------------------------------------
# portfolio variant -> (P1, P2, P3)
PSPEC = {"v0": (False, False, False), "vU": (True, False, False), "vK": (False, True, False),
         "vT": (False, False, True), "vUK": (True, True, False), "vUKT": (True, True, True)}


def parse_portfolio_variant(v: str) -> tuple[str, int]:
    """'vUKT' -> ('vUKT', 0); 'vUKT_P07' -> ('vUKT', 7). The seed only changes P2's blocks."""
    base, _, tag = v.partition("_P")
    assert base in PSPEC, f"unknown variant {v}"
    return base, (int(tag) if tag else 0)


def lw_block_correlation(X: np.ndarray, blocks: np.ndarray) -> tuple[np.ndarray, float]:
    """Ledoit-Wolf shrinkage toward a block constant-correlation target (P2).

    The target correlation of stocks i and j is the mean sample correlation over all pairs in
    their block pair. The intensity follows Ledoit and Wolf (2004) with the average correlation
    replaced pair by pair by the block value; with one block it equals lw_constant_correlation.
    """
    T, n = X.shape
    S = X.T @ X / T
    var = np.diag(S).copy()
    sd = np.sqrt(np.maximum(var, 1e-20))
    C = S / np.outer(sd, sd)
    labels, lab = np.unique(blocks, return_inverse=True)
    G = np.zeros((n, len(labels)))
    G[np.arange(n), lab] = 1.0
    cnt = G.sum(axis=0)
    sums = G.T @ C @ G - np.diag(cnt)               # remove the unit diagonal of C
    pairs = np.outer(cnt, cnt) - np.diag(cnt)
    rbar_all = float((C.sum() - n) / (n * (n - 1)))
    Rb = np.where(pairs > 0, sums / np.maximum(pairs, 1), rbar_all)
    R = Rb[lab][:, lab]
    F = R * np.outer(sd, sd)
    np.fill_diagonal(F, var)
    Y = X * X
    pi_mat = Y.T @ Y / T - S * S
    theta = (Y * X).T @ X / T - var[:, None] * S
    off = theta * (sd[None, :] / sd[:, None])
    np.fill_diagonal(off, 0.0)
    pi_hat = float(pi_mat.sum())
    rho_hat = float(np.trace(pi_mat) + (R * off).sum())
    gamma_hat = float(((F - S) ** 2).sum())
    delta = min(max((pi_hat - rho_hat) / gamma_hat / T, 0.0), 1.0) if gamma_hat > 0 else 1.0
    return delta * F + (1.0 - delta) * S, delta


def base_ic_and_dispersion(P, base, months, cfg):
    """Monthly rank IC of the base view over all stocks (months in `months`), and the monthly
    cross-sectional standard deviation of realised returns (all months). Both are read only at
    later months (see grinold_scale)."""
    ic = np.full(len(P["months"]), np.nan)
    xs = np.full(len(P["months"]), np.nan)
    for t in range(len(P["months"])):
        r = P["r"][int(P["starts"][t]):int(P["ends"][t])]
        r = r[np.isfinite(r)]
        if len(r) > 2:
            xs[t] = r.std()
    for t in months:
        lo, hi = int(P["starts"][t]), int(P["ends"][t])
        avail = [k for k in cfg["vintage_offsets"] if np.isfinite(base[k][lo:hi]).any()]
        if not avail:
            continue
        rows = np.isfinite(P["r"][lo:hi])
        for k in avail:
            rows &= np.isfinite(base[k][lo:hi])
        if rows.sum() < 3:
            continue
        vb = np.mean([zscore(base[k][lo:hi][rows].astype(np.float64)) for k in avail], axis=0)
        ic[t] = spearman(vb, P["r"][lo:hi][rows])
    return ic, xs


def grinold_scale(ic, xs, t, cfg) -> float:
    """c_t = IC x cross-sectional return volatility, from months strictly before t.

    The target of month m is realised at m + 1, so months m <= t - 1 are known at t. The IC is
    the trailing mean over up to `trust_window` months, shrunk toward `ic_prior` with the weight
    of 12 months, and floored at zero.
    """
    past_ic = ic[:t][np.isfinite(ic[:t])][-cfg["trust_window"]:]
    n = len(past_ic)
    ic_bar = max((n * (past_ic.mean() if n else 0.0) + 12 * cfg["ic_prior"]) / (n + 12), 0.0)
    past_xs = xs[:t][np.isfinite(xs[:t])][-12:]
    return ic_bar * (float(past_xs.mean()) if len(past_xs) else 0.0)


def build_portfolios(P, base, corr, sets, months, variants, D, cfg, ic_all, xs_sd, log=print):
    """Weights and predicted monthly volatility for each portfolio variant, month by month.

    Every variant uses M7's frozen vHE forecast. v0 reproduces M7's portfolio exactly.
    """
    n_rows = len(P["ids"])
    W = {v: np.zeros(n_rows) for v in variants}
    PV = {v: np.zeros(n_rows) for v in variants}
    offs = cfg["vintage_offsets"]
    e, _ = SPEC[FORECAST_VARIANT]
    q_hist = {v: [] for v in variants}
    for i, t in enumerate(months):
        lo, hi = int(P["starts"][t]), int(P["ends"][t])
        usable, _, Xz = month_window_both(D, P["ids"][lo:hi], P["months"][t], cfg)
        if usable.sum() < cfg["min_names"]:
            continue
        ind = P["ind"][lo:hi][usable]
        avail = [k for k in offs if np.isfinite(base[k][lo:hi][usable]).all()]
        Sig, delta = lw_constant_correlation(Xz)
        if not avail:
            continue
        # M7's frozen forecast (vHE)
        g = sets[0]["grp"][lo:hi][usable]
        wg = trust_weights(sets[0]["ic"], t, cfg["trust_window"])
        lam = np.where(g >= 0, wg[np.clip(g, 0, None)], 0.0)
        stack, members = [], []
        for k in avail:
            b = base[k][lo:hi][usable].astype(np.float64)
            c = corr[(0, e)][k][lo:hi][usable].astype(np.float64)
            c = np.where(np.isfinite(c), c, b)
            stack.append(b + lam * (c - b))
            members += [zscore(b), zscore(c)]
        m, _ = ensemble_forecast(stack, True, cfg["tau2_floor"])
        mn = industry_neutral(m, ind)
        sd_m = mn.std()
        c_t = grinold_scale(ic_all, xs_sd, t, cfg)
        alpha = c_t * mn / sd_m if sd_m > 0 else np.zeros_like(mn)
        Zm = np.vstack(members)
        Dev = (c_t * (Zm - Zm.mean(axis=0))).T                         # N x M
        Mn = Zm.shape[0]
        Omega = Dev @ Dev.T / (Mn * (Mn - 1)) if Mn > 1 else np.zeros((len(mn), len(mn)))
        sig_blocks = {}
        for v in variants:
            name, seed = parse_portfolio_variant(v)
            p1, p2, p3 = PSPEC[name]
            if not (p1 or p2 or p3):
                w = tangency_minus_gmv(mn, Sig)                         # M7 exactly
                W[v][lo:hi][usable] = w
                PV[v][lo:hi][usable] = np.sqrt(max(21.0 * w @ Sig @ w, 0.0))
                continue
            if p2:
                if seed not in sig_blocks:
                    blk = (P["grp"] if seed == 0 else placebo_groups(P, seed))[lo:hi][usable]
                    sig_blocks[seed] = lw_block_correlation(Xz, blk)[0]
                S_d = sig_blocks[seed]
            else:
                S_d = Sig
            S_m = 21.0 * S_d + (Omega if p1 else 0.0)
            w = tangency_minus_gmv(alpha if p1 else mn, S_m)
            if p3:
                var_p = float(w @ S_m @ w)
                q = float(w @ alpha) / var_p if var_p > 0 else np.nan
                prior = [x for x in q_hist[v][-cfg["trust_window"]:] if np.isfinite(x)]
                med = float(np.median(prior)) if len(prior) >= 12 else np.nan
                rel = q / med if np.isfinite(q) and np.isfinite(med) and med > 0 else 1.0
                w = w * float(np.clip(0.5 * rel + 0.5, 0.5, 2.0))
                q_hist[v].append(q)
            W[v][lo:hi][usable] = w
            PV[v][lo:hi][usable] = np.sqrt(max(float(w @ S_m @ w), 0.0))
        if i % 36 == 0:
            log(f"  {P['months'][t]}  held {int(usable.sum()):,}/{hi - lo:,}  delta(Z) {delta:.3f}  "
                f"c_t {c_t:.4f}  lambda s/l/m {np.round(wg, 2).tolist()}"
                + "".join(f"  {v} gross {np.abs(W[v][lo:hi]).sum():.2f}" for v in variants
                          if PSPEC[parse_portfolio_variant(v)[0]][2]))
    return W, PV


# ----------------------------------------------------------------------------------------------
# Entry point
# ----------------------------------------------------------------------------------------------
def build_weights(chars, features, daily_ret, cfg=CONFIG, kappas=None, out_range=None, log=print):
    """Run all stages. `kappas` names the portfolio variants to compute (design check); main uses
    CONFIG. Returns the panel, output month indices, a dict of weights per variant (plus
    '<variant>__pred_vol' arrays), and None (kept for the runner's interface).
    """
    np.random.seed(cfg["seed"])
    variants = [cfg["variant"]] if kappas is None else list(kappas)
    with threadpool_limits(limits=cfg["n_threads"]):
        feats = [str(f) for f in features["features"].tolist()]
        P = prepare_panel(chars, feats)
        log(f"panel {len(P['ids']):,} rows x {len(feats)} features, {len(P['months'])} months")
        rank_within_month(P)
        y = training_target(P)
        if out_range is None:
            out = output_months(P, cfg)
        else:
            lo_d, hi_d = (np.datetime64(d) for d in out_range)
            out = np.arange(cfg["min_train_months"], len(P["months"]))
            out = out[(P["months"][out] >= lo_d) & (P["months"][out] <= hi_d)]
        log(f"output {len(out)} months, {P['months'][out[0]]} .. {P['months'][out[-1]]}")
        # Forecasts also cover the trust window before the first output month, so the evidence
        # records exist from the start. Those months receive no weights.
        hist = np.arange(max(cfg["min_train_months"], int(out[0]) - cfg["trust_window"]), int(out[-1]) + 1)
        e, _ = SPEC[FORECAST_VARIANT]
        base, corr = forecast_hierarchy(P, y, hist, cfg, {0: P["grp"]}, [e], log)
        sets = {0: {"grp": P["grp"], "ic": ic_history(P, base, corr[(0, e)], P["grp"], hist, cfg)}}
        ic_all, xs_sd = base_ic_and_dispersion(P, base, hist, cfg)
        first_day = P["months"][out[0]] - np.timedelta64(int(cfg["cov_window_days"] * 1.6), "D")
        D = daily_matrix(daily_ret, np.unique(P["ids"]), first_day)
        log(f"daily matrix {D['M'].shape[0]:,} days x {D['M'].shape[1]:,} ids")
        W, PV = build_portfolios(P, base, corr, sets, out, variants, D, cfg, ic_all, xs_sd, log)
    result = {v: W[v] for v in variants}
    result.update({f"{v}__pred_vol": PV[v] for v in variants})
    return P, out, result, None


def main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    """Monthly portfolio weights for the ctff_test months. Columns: id, eom, w."""
    P, out, W, _ = build_weights(chars, features, daily_ret)
    rows = np.isin(P["month_ix"], out)
    result = pd.DataFrame({
        "id": P["ids"][rows].astype(np.int64),
        "eom": P["eom_out"][rows],
        "w": W[CONFIG["variant"]][rows],
    })
    assert len(result) > 0 and result.notna().all().all(), "invalid output"
    return result