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

SEED = 42
ALPHAS = 10.0 ** np.arange(-4, 3)
VALID_YEARS = 10
MIN_TRAIN_YEARS = 5
BLOCK = 50
CHUNK = 100_000
VOL_MONTHS = 12
MIN_DAYS = 60
N_STYLES = 13
MIN_INDUSTRY = 10
RISK_YEARS = 10
HL_VOL, HL_CORR, HL_SPEC = 84, 504, 84
VT_DAYS = 63
TARGET_VOL = 0.05
N_THREADS = 8
MAX_TREES = 1000
EARLY_STOP = 50
XGB_PARAMS = {"objective": "reg:squarederror", "tree_method": "hist", "max_depth": 3, "learning_rate": 0.05,
              "min_child_weight": 1000, "subsample": 0.5, "colsample_bytree": 0.5, "max_bin": 64,
              "seed": SEED, "nthread": N_THREADS}


def to_bool(s):
    if s.dtype == bool:
        return s.to_numpy()
    return s.astype(str).str.strip().str.lower().isin(["1", "1.0", "true", "t", "yes"]).to_numpy()


def neutral_ranks(values, month, industry):
    v = pd.DataFrame(values)
    g = v.groupby(month)
    x = ((g.rank(method="average") - 0.5) / g.transform("count") - 0.5).fillna(0.0)
    x = x - x.groupby([month, industry]).transform("mean")
    return x.to_numpy(np.float32)


def neutral_target(ret, month, industry):
    s = pd.Series(ret)
    g = s.groupby(month)
    y = (g.rank(method="average") - 0.5) / g.transform("count") - 0.5
    y = y - y.groupby([month, industry]).transform("mean")
    return y.to_numpy(np.float64)


def gram(X, y, rows):
    p = X.shape[1]
    xtx, xty = np.zeros((p, p)), np.zeros(p)
    for i in range(0, len(rows), CHUNK):
        r = rows[i:i + CHUNK]
        xc = X[r].astype(np.float64)
        xtx += xc.T @ xc
        xty += xc.T @ y[r]
    return xtx, xty


def window(grams, years):
    return sum(grams[t][0] for t in years), sum(grams[t][1] for t in years)


def ridge(xtx, xty, alphas):
    d, v = np.linalg.eigh(xtx)
    lam = np.asarray(alphas) * d.mean()
    return v @ ((v.T @ xty)[:, None] / (d[:, None] + lam[None, :]))


def predict(X, rows, B):
    out = np.empty((len(rows), B.shape[1]))
    for i in range(0, len(rows), CHUNK):
        r = rows[i:i + CHUNK]
        out[i:i + len(r)] = X[r].astype(np.float64) @ B
    return out


def monthly_ic(month, pred, ret):
    ok = ~np.isnan(ret)
    m = month[ok]
    a = pd.DataFrame(pred[ok]).groupby(m).rank()
    b = pd.Series(ret[ok]).groupby(m).rank()
    a = a - a.groupby(m).transform("mean")
    b = b - b.groupby(m).transform("mean")
    cov = a.mul(b, axis=0).groupby(m).sum()
    va = (a ** 2).groupby(m).sum()
    vb = (b ** 2).groupby(m).sum()
    return cov.div(np.sqrt(va.mul(vb, axis=0)))


def split(t, first):
    h = t - first
    if h >= VALID_YEARS + MIN_TRAIN_YEARS:
        return range(first, t - VALID_YEARS), range(t - VALID_YEARS, t)
    if h >= 2:
        n = max(1, int(round(0.3 * h)))
        return range(first, t - n), range(t - n, t)
    return None


def ridge_predictions(X, ret, month, fy, grams, years):
    first = min(grams)
    pred = np.full(len(fy), np.nan)
    for t in years:
        s = split(t, first)
        if s is None:
            continue
        valid = np.flatnonzero(np.isin(fy, list(s[1])))
        ic = monthly_ic(month[valid], predict(X, valid, ridge(*window(grams, s[0]), ALPHAS)), ret[valid])
        best = int(np.argmax((ic.mean() / ic.std()).to_numpy()))
        rows = np.flatnonzero(fy == t)
        pred[rows] = predict(X, rows, ridge(*window(grams, range(first, t)), [ALPHAS[best]]))[:, 0]
    return pred


def xgb_fit_predict(X, y, tr, va, full, rows):
    dtr = xgb.QuantileDMatrix(X[tr], label=y[tr], max_bin=64)
    dva = xgb.QuantileDMatrix(X[va], label=y[va], ref=dtr, max_bin=64)
    bst = xgb.train(XGB_PARAMS, dtr, num_boost_round=MAX_TREES, evals=[(dva, "valid")],
                    early_stopping_rounds=EARLY_STOP, verbose_eval=False)
    n = bst.best_iteration + 1
    del dtr, dva
    dfull = xgb.QuantileDMatrix(X[full], label=y[full], max_bin=64)
    bst = xgb.train(XGB_PARAMS, dfull, num_boost_round=n)
    del dfull
    return bst.inplace_predict(X[rows]).astype(np.float64), n


def xgb_predictions(X, y, fy, years):
    first = int(fy.min())
    pred = np.full(len(fy), np.nan)
    for t in years:
        s = split(t, first)
        if s is None:
            continue
        clock = time.time()
        a, b, c = np.searchsorted(fy, [s[1][0], t, t + 1])
        if np.isnan(y[:b]).any():
            ok = ~np.isnan(y)
            tr, va, full = np.flatnonzero(ok[:a]), a + np.flatnonzero(ok[a:b]), np.flatnonzero(ok[:b])
        else:
            tr, va, full = slice(0, a), slice(a, b), slice(0, b)
        pred[b:c], n = xgb_fit_predict(X, y, tr, va, full, slice(b, c))
        print(f"refit {t}  trees {n}  seconds {time.time() - clock:.0f}", flush=True)
    return pred


def zscore(p, m):
    s = pd.Series(p)
    g = s.groupby(m)
    return ((s - g.transform("mean")) / g.transform("std")).to_numpy()


def daily_arrays(daily, keep_ids, start_day, end_day):
    did = daily["id"].to_numpy().astype(np.int64)
    day = pd.to_datetime(daily["date"]).to_numpy().astype("datetime64[D]").astype(np.int64)
    dr = daily["ret_exc"].to_numpy(np.float64)
    keep = (day >= start_day) & (day <= end_day) & ~np.isnan(dr) & np.isin(did, keep_ids)
    order = np.lexsort((did[keep], day[keep]))
    did, day, dr = did[keep][order], day[keep][order], dr[keep][order]
    return did, day, day.astype("datetime64[D]").astype("datetime64[M]").astype(np.int64), dr


def stock_vol(did, mon, dr, ids, code):
    daily = pd.DataFrame({"m": mon, "id": did, "r2": dr ** 2})
    agg = daily.groupby(["m", "id"])["r2"].agg(["sum", "count"])
    months = agg.index.get_level_values("m")
    full = np.arange(months.min(), months.max() + 1)
    s2 = agg["sum"].unstack(fill_value=0.0).reindex(full, fill_value=0.0).rolling(VOL_MONTHS, min_periods=1).sum()
    n = agg["count"].unstack(fill_value=0).reindex(full, fill_value=0).rolling(VOL_MONTHS, min_periods=1).sum()
    vol = np.sqrt(s2 / n).where(n >= MIN_DAYS).to_numpy()
    cols = s2.columns.to_numpy()
    ci = np.clip(np.searchsorted(cols, ids), 0, len(cols) - 1)
    mi = np.clip(np.searchsorted(full, code), 0, len(full) - 1)
    ok = (cols[ci] == ids) & (full[mi] == code)
    out = np.full(len(ids), np.nan)
    out[ok] = vol[mi[ok], ci[ok]]
    return out


def style_loadings(xtx, k):
    sd = np.sqrt(np.diag(xtx))
    live = np.flatnonzero(sd > 0)
    corr = xtx[np.ix_(live, live)] / np.outer(sd[live], sd[live])
    dist = np.clip(1.0 - np.abs(corr), 0.0, None)
    np.fill_diagonal(dist, 0.0)
    labels = fcluster(linkage(squareform(dist, checks=False), method="average"), k, criterion="maxclust")
    L = np.zeros((len(sd), k))
    for c in range(k):
        j = np.flatnonzero(labels == c + 1)
        if len(j) == 0:
            continue
        v = np.linalg.eigh(corr[np.ix_(j, j)])[1][:, -1]
        v = v * np.sign(v[np.argmax(np.abs(v))])
        L[live[j], c] = np.where(v >= 0, 1.0, -1.0) / (len(j) * sd[live[j]])
    return L


def exposures(Xm, ind, L, groups):
    s = Xm.astype(np.float64) @ L
    sd = s.std(axis=0)
    s = (s - s.mean(axis=0)) / np.where(sd > 0, sd, 1.0)
    return np.hstack([(ind[:, None] == groups[None, :]).astype(np.float64), s])


def solve_sigma(E, F, dvar, V):
    a = V / dvar[:, None]
    Ed = E / dvar[:, None]
    M = np.eye(E.shape[1]) + F @ (E.T @ Ed)
    return a - Ed @ np.linalg.solve(M, F @ (E.T @ a))


def markowitz(E, F, dvar, mu, A):
    S = solve_sigma(E, F, dvar, np.column_stack([mu, A]))
    lam = np.linalg.solve(A.T @ S[:, 1:], A.T @ S[:, 0])
    w = S[:, 0] - S[:, 1:] @ lam
    return 2.0 * w / np.abs(w).sum()


def simple_weights(pred, has, month, industry):
    x = np.where(has, pred, 0.0)
    x = x - pd.Series(x).groupby([month, industry]).transform("mean").to_numpy()
    g = pd.Series(np.abs(x)).groupby(month).transform("sum").to_numpy()
    return np.where(g > 0, 2.0 * x / g, 0.0)


def build_weights(X, ids, code, month, fy, industry, me, pred, has, grams, daily):
    first_code, first_year = int(code[has].min()), int(fy[has].min())
    rs = first_code - 12 * RISK_YEARS + 1
    start_day = np.datetime64(rs, "M").astype("datetime64[D]").astype(np.int64)
    end_day = month.max().astype("datetime64[D]").astype(np.int64)
    did, dday, dmon, dr = daily_arrays(daily, np.unique(ids[code >= rs - 1]), start_day, end_day)
    if len(did) == 0:
        return simple_weights(pred, has, month, industry)
    raw = stock_vol(did, dmon, dr, ids[has], code[has])
    sig = np.full(len(code), np.nan)
    sig[has] = np.where(np.isnan(raw), pd.Series(raw).groupby(code[has]).transform("median").to_numpy(), raw)
    L = style_loadings(window(grams, range(min(grams), first_year))[0], N_STYLES)
    pre = (code >= rs - 1) & (fy < first_year)
    counts = pd.crosstab(code[pre], industry[pre]).mean()
    big = counts.index[counts >= MIN_INDUSTRY].to_numpy()
    risk_ind = np.where(np.isin(industry, big), industry, -2)
    groups = np.append(np.sort(big), -2)
    months, starts = np.unique(code, return_index=True)
    ends = np.append(starts[1:], len(code))
    d_lo, d_hi = np.searchsorted(dmon, months), np.searchsorted(dmon, months, side="right")
    udays = np.unique(dday)
    uid_all = np.unique(ids[code >= rs - 1])
    K = len(groups) + N_STYLES
    lv, lc, ls = 0.5 ** (1 / HL_VOL), 0.5 ** (1 / HL_CORR), 0.5 ** (1 / HL_SPEC)
    cv, cc, wv, wc = np.zeros((K, K)), np.zeros((K, K)), 0.0, 0.0
    spec, spec_w, nobs = np.zeros(len(uid_all)), np.zeros(len(uid_all)), np.zeros(len(uid_all))
    w = np.zeros(len(code))
    prev = None
    for k in np.flatnonzero(months >= rs - 1):
        a, b = starts[k], ends[k]
        if prev is not None and d_hi[k] > d_lo[k]:
            E0, uid0, sw0 = prev
            sl = slice(d_lo[k], d_hi[k])
            R = (pd.DataFrame({"d": dday[sl], "id": did[sl], "r": dr[sl]})
                 .pivot_table(index="d", columns="id", values="r").reindex(columns=uid0).to_numpy())
            p0 = np.searchsorted(uid_all, uid0)
            for r in R:
                ok = ~np.isnan(r)
                f = np.linalg.lstsq(E0[ok] * sw0[ok, None], r[ok] * sw0[ok], rcond=None)[0]
                res = r[ok] - E0[ok] @ f
                cv = lv * cv + (1 - lv) * np.outer(f, f)
                cc = lc * cc + (1 - lc) * np.outer(f, f)
                wv, wc = lv * wv + (1 - lv), lc * wc + (1 - lc)
                p = p0[ok]
                spec[p] = ls * spec[p] + (1 - ls) * res ** 2
                spec_w[p] = ls * spec_w[p] + (1 - ls)
                nobs[p] += 1
        E = exposures(X[a:b], risk_ind[a:b], L, groups)
        prev = (E, ids[a:b], np.nan_to_num(me[a:b]) ** 0.25)
        if not has[a:b].any() or wv == 0:
            continue
        vf = np.sqrt(np.diag(cv) / wv)
        c = cc / wc
        dc = np.sqrt(np.diag(c))
        dc = np.where(dc > 0, dc, 1.0)
        F = 21 * np.outer(vf, vf) * c / np.outer(dc, dc)
        p = np.searchsorted(uid_all, ids[a:b])
        dvar = np.where(nobs[p] >= MIN_DAYS, 21 * spec[p] / np.maximum(spec_w[p], 1e-12), np.nan)
        dvar = np.where(np.isnan(dvar), np.nanmedian(dvar) if np.isfinite(dvar).any() else 1e-3, dvar)
        ind = industry[a:b]
        A = (ind[:, None] == np.unique(ind)[None, :]).astype(np.float64)
        wm = np.nan_to_num(markowitz(E, F, dvar, np.nan_to_num(pred[a:b] * sig[a:b]), A))
        lo = np.searchsorted(dday, udays[max(np.searchsorted(udays, dday[d_hi[k] - 1]) - VT_DAYS + 1, 0)])
        H = (pd.DataFrame({"d": dday[lo:d_hi[k]], "id": did[lo:d_hi[k]], "r": dr[lo:d_hi[k]]})
             .pivot_table(index="d", columns="id", values="r").reindex(columns=ids[a:b]).fillna(0.0).to_numpy())
        hist = np.sqrt(252) * (H @ wm).std(ddof=1) if len(H) > 1 else 0.0
        w[a:b] = wm * (TARGET_VOL / hist if hist > 0 else 1.0)
    return w


def main(chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame) -> pd.DataFrame:
    np.random.seed(SEED)
    feats = sorted(set(features["features"].astype(str)) & set(chars.columns))
    eom = pd.to_datetime(chars["eom"]).to_numpy()
    code0 = eom.astype("datetime64[M]").astype(np.int64)
    id0 = chars["id"].to_numpy().astype(np.int64)
    order = np.lexsort((id0, code0))
    month, code, ids = eom[order], code0[order], id0[order]
    fy = 1970 + code // 12 + (code % 12 == 11)
    industry = (pd.to_numeric(chars["sic"], errors="coerce") // 100).fillna(-1).astype("int64").to_numpy()[order]
    me = (chars["market_equity"].to_numpy(np.float64)[order] if "market_equity" in chars.columns
          else np.ones(len(order)))
    ret = chars["ret_exc_lead1m"].to_numpy(np.float64)[order]
    test = to_bool(chars["ctff_test"])[order]
    X = np.empty((len(order), len(feats)), dtype=np.float32)
    for s in range(0, len(feats), BLOCK):
        cols = feats[s:s + BLOCK]
        X[:, s:s + len(cols)] = neutral_ranks(chars[cols].to_numpy(np.float64)[order], month, industry)
    y = neutral_target(ret, month, industry)
    years = [int(t) for t in np.unique(fy)]
    grams = {t: gram(X, y, np.flatnonzero((fy == t) & ~np.isnan(y))) for t in years}
    out_years = sorted({int(t) for t in fy[test]}) if test.any() else years
    zr = zscore(ridge_predictions(X, ret, month, fy, grams, out_years), month)
    zx = zscore(xgb_predictions(X, y, fy, out_years), month)
    pred = np.where(np.isnan(zx), zr, 0.5 * zr + 0.5 * zx)
    has = ~np.isnan(pred)
    w = (build_weights(X, ids, code, month, fy, industry, me, pred, has, grams, daily_ret)
         if has.any() else np.zeros(len(order)))
    keep = test if test.any() else has
    if not keep.any():
        keep = np.ones(len(order), dtype=bool)
    return pd.DataFrame({"id": ids[keep],
                         "eom": pd.DatetimeIndex(month[keep]).normalize().astype("datetime64[ns]"),
                         "w": np.nan_to_num(w[keep])}).reset_index(drop=True)