"""JKP Common Task Framework submission: peer-relative boosted forecasts + Barra-style risk model.

Authors: Andrew Qi, Marcus Kwauk, Boyce Yang, Harley Chen (Yale SOM, MGT 924).

Pipeline (re-estimated every month, using only information available at the formation date t):
  1. Peer groups: each month, the 402 characteristics are percentile-ranked within the month; stocks are clustered
     by Ward's method on the top 20 principal components of these ranks into K = round(n / 100) groups (clipped to
     [10, 60]). Uses that month's characteristics only (no returns, no other months).
  2. Features: the 402 raw characteristics plus 402 peer-relative characteristics (a stock's within-month rank
     minus the average rank of its peer group; missing where the raw value is missing). No feature is hand-picked.
  3. Forecast: LightGBM on the 804 features. Target = arctan of the within-month z-score of next-month excess
     return. Refit every January on every row whose return was realised by the previous 31 December (expanding
     window). Settings (1,000 trees, learning rate 0.01, 255 leaves, >= 2,000 rows per leaf) are borrowed from
     Geertsema's CTF entry.
  4. Risk model: daily cross-sectional ridge regressions of stock returns on 12 Fama-French industries and the 402
     ranked characteristics (standardised) give factor returns and residuals. Factor covariance = EWMA correlation
     (half-life 504 days) x EWMA volatilities (half-life 84 days) over 2,520 days; specific variance = EWMA
     (half-life 84) of the last 252 residuals, at least 200 observations. Only daily data up to t is used.
  5. Weights: w = A^-1 mu - (1'A^-1 mu / 1'A^-1 1) A^-1 1 (tangency minus global minimum variance, i.e. the
     dollar-neutral mean-variance portfolio), A = Sigma + kappa * tr(Sigma)/n * I with kappa = 1, scaled to unit
     gross exposure. Stocks without enough daily history for a specific variance get weight 0.

Borrowed components (disclosed in the methodology PDF): the LightGBM settings, the tangency-minus-GMV mapping and
kappa (Geertsema's CTF submission); the Barra settings and the FF12 industry map (CTF Markowitz-ML /
Minimum-Variance benchmarks). Deterministic: fixed seed, LightGBM deterministic mode with a fixed thread count,
deterministic Ward clustering.
"""
import time

import lightgbm as lgb
from scipy.cluster.hierarchy import fcluster, linkage
import numpy as np
import pandas as pd

LGBM_PARAMS = dict(n_estimators=1000, learning_rate=0.01, num_leaves=255, min_child_samples=2000, random_state=42,
                   n_jobs=16, deterministic=True, force_row_wise=True, verbose=-1)
N_PCS, STOCKS_PER_GROUP, K_MIN, K_MAX = 20, 100, 10, 60     # peer groups: fixed in advance, not tuned
LGBM_MIN_ROWS = 1000                 # skip a refit year with less training data than this (tiny validation samples)
KAPPA = 1.0                          # extra ridge on the covariance in the portfolio step
DAYS_PER_MONTH = 21
FCT_RIDGE = 1e-4                     # ridge in the daily cross-sectional factor regressions
FCT_OBS = 2520                       # days of factor returns for the factor covariance
HL_COR, HL_VAR, HL_SPEC = 504, 84, 84
SPEC_WINDOW, SPEC_MIN_OBS = 252, 200
ID_COLS = ["id", "eom", "eom_ret", "sic", "ret_exc_lead1m", "ctff_test"]


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


# --------------------------------------------------------------------------------------------------------------
# 1. Features
# --------------------------------------------------------------------------------------------------------------
def prank(x, month_ix):
    """Within-month percentile rank in [-0.5, 0.5]; exact zeros -> -0.5; missing -> 0."""
    s = pd.Series(x)
    g = s.groupby(month_ix, sort=False)
    r = (g.rank(method="max") / g.transform("count")).to_numpy(dtype=np.float32)
    r[x == 0] = 0.0
    r = r - np.float32(0.5)
    r[np.isnan(r)] = 0.0
    return r


# --------------------------------------------------------------------------------------------------------------
# 2. Barra-style risk model
# --------------------------------------------------------------------------------------------------------------
IND_NAMES = ["BusEq", "Chems", "Durbl", "Enrgy", "Hlth", "Manuf", "Money", "NoDur", "Other", "Shops", "Telcm", "Utils"]


def ff12_class(sic):
    s = pd.to_numeric(pd.Series(sic), errors="coerce").fillna(-1).astype(int).to_numpy()

    def r(*spans):
        m = np.zeros(len(s), bool)
        for sp in spans:
            lo, hi = (sp, sp) if isinstance(sp, int) else sp
            m |= (s >= lo) & (s <= hi)
        return m

    conds = [
        ("NoDur", r((100, 999), (2000, 2399), (2700, 2749), (2770, 2799), (3100, 3199), (3940, 3989))),
        ("Durbl", r((2500, 2519), (3630, 3659), (3710, 3711), 3714, 3716, (3750, 3751), 3792, (3900, 3939), (3990, 3999))),
        ("Manuf", r((2520, 2589), (2600, 2699), (2750, 2769), (3000, 3099), (3200, 3569), (3580, 3629), (3700, 3709),
                    (3712, 3713), 3715, (3717, 3749), (3752, 3791), (3793, 3799), (3830, 3839), (3860, 3899))),
        ("Enrgy", r((1200, 1399), (2900, 2999))),
        ("Chems", r((2800, 2829), (2840, 2899))),
        ("BusEq", r((3570, 3579), (3660, 3692), (3694, 3699), (3810, 3829), (7370, 7379))),
        ("Telcm", r((4800, 4899))),
        ("Utils", r((4900, 4949))),
        ("Shops", r((5000, 5999), (7200, 7299), (7600, 7699))),
        ("Hlth", r((2830, 2839), 3693, (3840, 3859), (8000, 8099))),
        ("Money", r((6000, 6999))),
    ]
    out = np.full(len(s), "Other", dtype=object)
    for name, m in reversed(conds):          # first matching rule wins
        out[m] = name
    return out


def standardize(Z):
    Z = Z - Z.mean(0, keepdims=True)
    sd = Z.std(0, keepdims=True)
    return np.where(sd > 1e-8, Z / np.where(sd > 1e-8, sd, 1), 0).astype(np.float32)


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


def factor_cov(F_win):
    T = F_win.shape[0]

    def wcov(hl):
        w = ewma_weights(T, hl)
        Fc = F_win - w @ F_win
        return (Fc * w[:, None]).T @ Fc / (1 - (w**2).sum())

    C = wcov(HL_COR)
    sd_c = np.sqrt(np.clip(np.diag(C), 1e-20, None))
    corr = C / np.outer(sd_c, sd_c)
    sd_v = np.sqrt(np.clip(np.diag(wcov(HL_VAR)), 1e-20, None))
    return corr * np.outer(sd_v, sd_v)


def specific_var(E_win):
    T = E_win.shape[0]
    w = (0.5 ** (np.arange(T)[::-1] / HL_SPEC))[:, None] * np.isfinite(E_win)
    v = np.nansum(np.where(np.isfinite(E_win), E_win, 0) ** 2 * w, 0) / np.clip(w.sum(0), 1e-12, None)
    v[np.isfinite(E_win).sum(0) < SPEC_MIN_OBS] = np.nan
    return v


def barra_tan_minus_gmv(mu, X, Sf, dvar, kappa):
    """Tangency minus GMV with Sigma = X Sf X' + diag(dvar) plus kappa ridge, via the Woodbury identity."""
    n = len(mu)
    tr = np.einsum("ij,jk,ik->", X, Sf, X) + dvar.sum()
    Dp = dvar + kappa * tr / n
    Sf = (Sf + Sf.T) / 2 + 1e-12 * np.trace(Sf) / len(Sf) * np.eye(len(Sf))
    Bm = np.column_stack([mu, np.ones(n)])
    DiB, DiX = Bm / Dp[:, None], X / Dp[:, None]
    inner = np.linalg.inv(Sf) + X.T @ DiX
    sol = DiB - DiX @ np.linalg.solve(inner, X.T @ DiB)
    a, b = sol[:, 0], sol[:, 1]
    w = a - (a.sum() / b.sum()) * b
    return w / np.abs(w).sum()


# --------------------------------------------------------------------------------------------------------------
# peer groups
def peer_groups(X):
    """X: (n_stocks, 402) ranked characteristics of one month -> peer-group label per stock (Ward on 20 PCs)."""
    n = X.shape[0]
    k = int(np.clip(round(n / STOCKS_PER_GROUP), K_MIN, K_MAX))
    if n <= k:
        return np.arange(n)
    Z = X - X.mean(0)
    U, s, _ = np.linalg.svd(Z, full_matrices=False)
    P = U[:, :N_PCS] * s[:N_PCS]
    return fcluster(linkage(P, "ward"), k, "maxclust") - 1


# main
# --------------------------------------------------------------------------------------------------------------
def main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    t0 = time.perf_counter()
    np.random.seed(1)                        # no randomness is used; set for completeness (Rule 18)
    feats = features["features"].astype(str).tolist()

    # ---- 1. panel sorted by (eom, id) and percentile ranks ---------------------------------------------------
    df = chars[ID_COLS].copy()
    df["eom"] = pd.to_datetime(df["eom"])
    df["eom_ret"] = pd.to_datetime(df["eom_ret"])
    order = np.lexsort((df["id"].to_numpy(), df["eom"].to_numpy()))     # stable sort by (eom, id)
    df = df.iloc[order].reset_index(drop=True)
    month_ix = pd.factorize(df["eom"])[0]
    X_all = np.empty((len(df), len(feats)), dtype=np.float32)        # ranked: Barra exposures
    X_raw = np.empty((len(df), len(feats)), dtype=np.float32)        # raw: LightGBM inputs
    for j, f in enumerate(feats):
        x = chars[f].to_numpy(dtype=np.float64)[order]
        X_all[:, j] = prank(x, month_ix)
        X_raw[:, j] = x.astype(np.float32)
    ids_all = df["id"].to_numpy().astype(np.int64)
    eom_all = df["eom"].to_numpy()
    ret_year = df["eom_ret"].dt.year.to_numpy()
    test = df["ctff_test"].astype(str).str.strip().isin(["1", "1.0", "True", "true"]).to_numpy()
    ind_all = ff12_class(df["sic"])
    IND_MAT = np.stack([(ind_all == k) for k in IND_NAMES], axis=1).astype(np.float32)
    r = df["ret_exc_lead1m"]
    g = r.groupby(df["eom"])
    y_lgbm = np.arctan(((r - g.transform("mean")) / g.transform("std")).to_numpy(np.float32))
    df = df[["id", "eom"]]
    log(f"ranked {X_all.shape[1]} features for {X_all.shape[0]:,} stock-months", t0)

    # ---- 1b. peer-relative characteristics (month-by-month, that month's ranks only) ----------------------------
    X_lgbm = np.empty((len(df), 2 * len(feats)), dtype=np.float32)
    X_lgbm[:, :len(feats)] = X_raw
    m_starts = np.flatnonzero(np.r_[True, month_ix[1:] != month_ix[:-1]])
    m_ends = np.r_[m_starts[1:], len(month_ix)]
    for lo, hi in zip(m_starts, m_ends):
        X = X_all[lo:hi].astype(np.float64)
        lab = peer_groups(X)
        G = np.zeros((hi - lo, lab.max() + 1))
        G[np.arange(hi - lo), lab] = 1.0
        rel = (X - G @ ((G.T @ X) / G.sum(0)[:, None])).astype(np.float32)
        rel[np.isnan(X_raw[lo:hi])] = np.nan
        X_lgbm[lo:hi, len(feats):] = rel
    del X_raw
    log(f"peer-relative features for {len(m_starts)} months", t0)

    # ---- 2. LightGBM forecasts, refit every January ---------------------------------------------------------
    assert np.all(np.diff(ret_year) >= 0), "rows must be sorted by month"
    year_list = np.unique(ret_year)
    # formation months to produce weights for: the test months (all months with a training year if none flagged)
    months_all = np.unique(eom_all)
    target_months = np.unique(eom_all[test]) if test.any() else months_all[np.isin(
        months_all, eom_all[ret_year > year_list[0]])]
    target_years = np.unique(ret_year[np.isin(eom_all, target_months)])

    mu_all = np.zeros(len(X_all))
    for Y in target_years:
        n_train = int(np.searchsorted(ret_year, Y, "left"))       # every row whose return is realised by Y-1
        lo, hi = n_train, int(np.searchsorted(ret_year, Y, "right"))
        ok = np.isfinite(y_lgbm[:n_train])
        if ok.sum() < LGBM_MIN_ROWS:
            continue                         # not enough history: zero forecast (zero weights)
        model = lgb.LGBMRegressor(**LGBM_PARAMS)
        model.fit(X_lgbm[:n_train][ok], y_lgbm[:n_train][ok])
        mu_all[lo:hi] = model.booster_.predict(X_lgbm[lo:hi]).astype(np.float64)
        log(f"LightGBM for return year {Y}: trained on {int(ok.sum()):,} rows", t0)
    del X_lgbm

    # ---- 3. daily return matrix and daily factor regressions -------------------------------------------------
    d = daily_ret[["id", "date", "ret_exc"]]
    d = d[d["id"].astype(np.int64).isin(np.unique(ids_all))]
    d_dates, di = np.unique(pd.to_datetime(d["date"]).to_numpy().astype("datetime64[D]"), return_inverse=True)
    d_ids, ii = np.unique(d["id"].to_numpy().astype(np.int64), return_inverse=True)
    M = np.full((len(d_dates), len(d_ids)), np.nan, dtype=np.float32)
    M[di, ii] = d["ret_exc"].to_numpy(dtype=np.float32)
    del d, di, ii
    log(f"daily matrix {M.shape[0]:,} days x {M.shape[1]:,} stocks", t0)

    m_lo = np.searchsorted(eom_all, months_all, "left")
    m_hi = np.searchsorted(eom_all, months_all, "right")
    month_pos = {pd.Timestamp(m): k for k, m in enumerate(months_all)}

    def day_index(t):
        """Number of trading days up to and including date t."""
        return int(np.searchsorted(d_dates, np.datetime64(pd.Timestamp(t).date(), "D"), "right"))

    def exposures(k):
        lo, hi = m_lo[k], m_hi[k]
        return np.hstack([IND_MAT[lo:hi], standardize(X_all[lo:hi])])

    # factor returns / residuals for days in month k+1 use exposures known at the end of month k
    k_last = month_pos[pd.Timestamp(target_months[-1])] - 1
    d_start = day_index(months_all[0])
    d_end = day_index(months_all[k_last + 1]) if k_last + 1 < len(months_all) else len(d_dates)
    K = 12 + X_all.shape[1]
    F = np.zeros((max(d_end - d_start, 0), K))
    E = np.full((max(d_end - d_start, 0), len(d_ids)), np.nan, dtype=np.float32)
    for k in range(0, k_last + 1):
        a = day_index(months_all[k])
        b = day_index(months_all[k + 1]) if k + 1 < len(months_all) else len(d_dates)
        if b <= a:
            continue
        ids_k = ids_all[m_lo[k]:m_hi[k]]
        cols = np.clip(np.searchsorted(d_ids, ids_k), 0, len(d_ids) - 1)
        has = d_ids[cols] == ids_k
        R = M[a:b, cols[has]]
        full = np.isfinite(R).all(0)
        if full.sum() == 0:
            continue
        X = exposures(k)[has][full].astype(np.float64)
        R = R[:, full].astype(np.float64)
        n = X.shape[0]
        f = np.linalg.solve(X.T @ X / n + FCT_RIDGE * np.eye(K), X.T @ R.T / n)
        F[a - d_start:b - d_start] = f.T
        E[a - d_start:b - d_start, cols[has][full]] = (R - (X @ f).T).astype(np.float32)
    log(f"daily factor regressions for {k_last + 1} months", t0)

    # ---- 4. monthly portfolio weights ------------------------------------------------------------------------
    w_all = np.zeros(len(X_all))
    for t in target_months:
        k = month_pos[pd.Timestamp(t)]
        lo, hi = m_lo[k], m_hi[k]
        end = day_index(t) - d_start         # the last day used is the formation date t itself
        if end <= 1:
            continue
        Sf = factor_cov(F[max(0, end - FCT_OBS):end]) * DAYS_PER_MONTH
        ids_t = ids_all[lo:hi]
        cols = np.clip(np.searchsorted(d_ids, ids_t), 0, len(d_ids) - 1)
        has = d_ids[cols] == ids_t
        dvar = np.full(len(ids_t), np.nan)
        dvar[has] = specific_var(E[max(0, end - SPEC_WINDOW):end, cols[has]]) * DAYS_PER_MONTH
        use = np.isfinite(dvar)
        mu = mu_all[lo:hi]
        if use.sum() < 2 or not np.any(mu[use] != mu[use][0]):
            continue                         # not enough history (or no forecast dispersion): stay in cash
        X = exposures(k).astype(np.float64)
        keep = np.any(X != 0, axis=0)        # drop industries with no stocks this month
        w = barra_tan_minus_gmv(mu[use], X[use][:, keep], Sf[np.ix_(keep, keep)], dvar[use], KAPPA)
        if np.all(np.isfinite(w)):
            w_all[lo + np.flatnonzero(use)] = w
        else:
            print(f"warning: non-finite weights in {pd.Timestamp(t).date()}, month left in cash", flush=True)
        if pd.Timestamp(t).month == 12:
            log(f"weights through {pd.Timestamp(t).date()}: {int(use.sum())} of {hi - lo} stocks held", t0)

    out = pd.DataFrame({"id": ids_all, "eom": df["eom"].dt.date, "w": w_all})
    out = out[np.isin(eom_all, target_months)].reset_index(drop=True)
    log(f"done: {len(out):,} rows, {out['eom'].nunique()} months", t0)
    return out