#!/usr/bin/env python
"""
SPRINT A — Replication of the 2024 UMDAC Datathon CHAMPION's pipeline
(Team Hokkien Mee Is Red: Sim Sze Yu, John Ong Ming Hom, Lei Wing Teng)

Champion deck steps (champion-slides-text.txt)   -> this script
----------------------------------------------------------------
1. Data cleaning: linear interpolation nulls/zeros -> clean_data()
2. Feature engineering: 11 technical indicators    -> engineer_features()
   (EMA12/26, MACD, RSI, Bollinger, VWAP, ATR, OBV,
    Stochastic, Volatility, MFI, Whale Ratio, Momentum) + on-chain cols
3. Correlation matrix + feature importance elbow   -> plots
4. 60h-context sliding-window thinking w/ lags     -> add_lags()
5. XGBoost regressor, why: non-linear/regularized  -> train_model()
   (RandomizedSearchCV over TimeSeriesSplit)
6. Target: next-hour price change                  -> make_target()
7. Optimize buy/sell thresholds (theirs: ±0.110)   -> sweep_thresholds()
8. Backtest w/ fees: Return, Sharpe, DD, WinRate,  -> backtest()
   Profit Factor (theirs: 227.93% ret, 1.66 Sharpe)

Outputs -> sprint-a/output/
Run: ~/datathon-2026/.venv/bin/python sprint-a/sprint_a_pipeline.py
"""
import json
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from xgboost import XGBRegressor
from sklearn.model_selection import RandomizedSearchCV, TimeSeriesSplit

BASE = Path(__file__).resolve().parent.parent          # ~/datathon-2026
DATA = BASE / "champions/2024-champion-hokkien-mee-is-red/datasets/UM_datathon_2024.csv"
OUT = Path(__file__).resolve().parent / "output"
OUT.mkdir(exist_ok=True)
FEE = 0.0002          # 0.02% per position change (assumption; champion reported fees too)
RNG = 42


# ---------------------------------------------------------------- 1. CLEAN
def clean_data() -> pd.DataFrame:
    df = pd.read_csv(DATA, parse_dates=["time"]).sort_values("time").reset_index(drop=True)
    # Champion: "linear interpolation to handle null or zero value"
    metric_cols = [c for c in df.columns if c != "time"]
    df[metric_cols] = df[metric_cols].replace(0, np.nan)
    df[metric_cols] = df[metric_cols].interpolate(method="linear", limit_direction="both")
    df = df.dropna(subset=["close_price"])
    print(f"[clean] rows={len(df)}  {df.time.min()} -> {df.time.max()}")
    return df


# ---------------------------------------------------- 2. FEATURE ENGINEERING
def rsi(close: pd.Series, n: int = 14) -> pd.Series:
    d = close.diff()
    up = d.clip(lower=0).ewm(alpha=1 / n, adjust=False).mean()
    dn = (-d.clip(upper=0)).ewm(alpha=1 / n, adjust=False).mean()
    return 100 - 100 / (1 + up / dn)


def engineer_features(df: pd.DataFrame) -> pd.DataFrame:
    x = pd.DataFrame(index=df.index)
    c, h, l, v, lo = df.close_price, df.high_price, df.low_price, df.volume, df.low_price

    x["ret_1h"] = c.pct_change()
    x["ema_12"] = c.ewm(span=12, adjust=False).mean() / c - 1
    x["ema_26"] = c.ewm(span=26, adjust=False).mean() / c - 1
    macd = c.ewm(span=12, adjust=False).mean() - c.ewm(span=26, adjust=False).mean()
    x["macd"] = macd / c
    x["macd_signal"] = macd.ewm(span=9, adjust=False).mean() / c
    x["rsi_14"] = rsi(c, 14)
    ma, sd = c.rolling(20).mean(), c.rolling(20).std()
    x["boll_upper"] = (ma + 2 * sd) / c - 1
    x["boll_lower"] = (ma - 2 * sd) / c - 1
    tp = (h + l + c) / 3
    x["vwap_24"] = (tp * v).rolling(24).sum() / v.rolling(24).sum() / c - 1
    tr = pd.concat([h - l, (h - c.shift()).abs(), (l - c.shift()).abs()], axis=1).max(axis=1)
    x["atr_14"] = tr.rolling(14).mean() / c
    x["obv"] = (np.sign(c.diff()).fillna(0) * v).cumsum()
    x["obv_z"] = (x["obv"] - x["obv"].rolling(168).mean()) / x["obv"].rolling(168).std()
    ll, hh = l.rolling(14).min(), h.rolling(14).max()
    x["stoch_14"] = 100 * (c - ll) / (hh - ll)
    x["volatility_24"] = x["ret_1h"].rolling(24).std()
    mfv = tp * v * np.sign(tp.diff()).fillna(0)
    x["mfi_14"] = mfv.rolling(14).sum() / (v.rolling(14).sum() * (tp.diff().abs().fillna(0)).rolling(14).sum() + 1e-9)
    x["whale_ratio"] = df.exchange_whale_ratio
    x["momentum_12"] = c.pct_change(12)
    # on-chain / market microstructure straight from the provided dataset
    for col in ["open_interest", "estimated_leverage_ratio", "taker_buy_sell_ratio",
                "coinbase_premium_gap", "miner_supply_ratio",
                "long_liquidations_usd", "short_liquidations_usd", "difficulty"]:
        x[col] = df[col] / (df[col].rolling(168).mean() + 1e-9)   # 1-week relative level

    # Champion: "features with time lags are highly influential"
    for lag in (1, 6, 24, 60):                                    # incl. their 60h horizon
        x[f"ret_lag_{lag}"] = x["ret_1h"].shift(lag)
        x[f"oi_lag_{lag}"] = x["open_interest"].shift(lag)
    return x.replace([np.inf, -np.inf], np.nan)


def make_target(df: pd.DataFrame) -> pd.Series:
    # Champion: "y (target) = next hour price change"
    return df.close_price.pct_change().shift(-1)


# ------------------------------------------------------------- 3-4. DATASET
def build_dataset():
    df = clean_data()
    X = engineer_features(df)
    y = make_target(df)
    data = pd.concat([X, y.rename("target")], axis=1).dropna()
    feat_cols = [c for c in data.columns if c != "target"]
    split = int(len(data) * 0.8)                                  # champion: 80/20
    Xtr, Xte = data[feat_cols][:split], data[feat_cols][split:]
    ytr, yte = data.target[:split], data.target[split:]
    print(f"[dataset] features={len(feat_cols)}  train={len(Xtr)}  test={len(Xte)}")
    return df, data, feat_cols, Xtr, Xte, ytr, yte


# ------------------------------------------------------------ 5. XGBOOST
def train_model(Xtr, ytr):
    param_dist = {
        "n_estimators": [300, 600], "max_depth": [4, 6, 8],
        "learning_rate": [0.01, 0.03, 0.05], "subsample": [0.7, 0.9],
        "colsample_bytree": [0.7, 0.9], "min_child_weight": [1, 5],
    }
    base = XGBRegressor(objective="reg:squarederror", random_state=RNG, n_jobs=8, tree_method="hist")
    search = RandomizedSearchCV(base, param_dist, n_iter=6, cv=TimeSeriesSplit(3),
                                scoring="neg_mean_absolute_error", random_state=RNG, verbose=0)
    search.fit(Xtr, ytr)
    print(f"[model] best params: {search.best_params_}")
    print(f"[model] CV MAE: {-search.best_score_:.6f} (next-hour return)")
    return search.best_estimator_


# --------------------------------------------- 7. THRESHOLDS + 8. BACKTEST
def run_positions(pred: np.ndarray, buy_thr: float):
    sell_thr = -buy_thr                                           # champion: symmetric ±0.110
    pos = np.where(pred > buy_thr, 1.0, np.where(pred < sell_thr, -1.0, 0.0))
    return pos


def backtest(pos: np.ndarray, ret: np.ndarray):
    trades = np.diff(pos, prepend=0.0) != 0
    strat = pos * ret - FEE * np.abs(np.diff(pos, prepend=0.0))
    equity = 1 + strat.cumsum()
    peak = np.maximum.accumulate(equity)
    dd = (equity / peak - 1).min()
    active = strat[pos != 0]
    win = (active > 0).mean() if len(active) else 0.0
    gross_win, gross_loss = strat[strat > 0].sum(), -strat[strat < 0].sum()
    n_hours = len(strat)
    return {
        "total_return_pct": round(100 * (equity[-1] - 1), 2),
        "sharpe_annualized": round(float(strat.mean() / (strat.std() + 1e-12) * np.sqrt(24 * 365)), 2),
        "max_drawdown_pct": round(100 * dd, 2),
        "win_rate_pct": round(100 * win, 2),
        "num_trades": int(trades.sum()),
        "profit_factor": round(float(gross_win / (gross_loss + 1e-12)), 3),
        "hours": n_hours,
    }


def sweep_thresholds(pred_tr, ret_tr, grid=None):
    grid = grid if grid is not None else np.arange(0.0005, 0.0051, 0.00025)  # 0.05%..0.5%
    results = []
    for t in grid:
        m = backtest(run_positions(pred_tr, t), ret_tr)
        results.append((float(t), m["sharpe_annualized"], m["total_return_pct"]))
    best = max(results, key=lambda r: r[1])
    print(f"[thresholds] best train Sharpe {best[1]} at ±{best[0]*100:.3f}% (champion: ±0.110%)")
    return best[0], results


# ------------------------------------------------------------------ PLOTS
def make_plots(df, data, feat_cols, model, Xte, yte, pred_te, pos_te, sweep, buy_thr):
    t_idx = df.time.reindex(data.index)
    fig, ax = plt.subplots(2, 1, figsize=(12, 8), sharex=True)
    ax[0].plot(t_idx, data.target.cumsum(), label="Buy & Hold (cum return)")
    strat = pos_te * yte - FEE * np.abs(np.diff(pos_te, prepend=0))
    ax[1].plot(t_idx[-len(yte):], np.cumsum(strat), color="crimson", label="Strategy (test)")
    ax[1].plot(t_idx[-len(yte):], yte.cumsum(), color="grey", alpha=.6, label="Buy & Hold (test)")
    for a in ax: a.legend(); a.grid(alpha=.3)
    ax[0].set_title("Sprint A — champion replication (2024 UMDAC dataset)")
    ax[1].set_title(f"Out-of-sample: pos from ±{buy_thr*100:.2f}% thresholds")
    plt.tight_layout(); plt.savefig(OUT / "backtest_equity.png", dpi=120); plt.close()

    imp = pd.Series(model.feature_importances_, index=feat_cols).sort_values(ascending=True)
    fig, ax = plt.subplots(figsize=(8, 9))
    imp.tail(20).plot.barh(ax=ax); ax.set_title("XGBoost feature importance (top 20)")
    plt.tight_layout(); plt.savefig(OUT / "feature_importance.png", dpi=120); plt.close()

    cum = np.cumsum(np.sort(imp.values)[::-1])
    fig, ax = plt.subplots(figsize=(7, 4))
    ax.plot(range(1, len(cum) + 1), cum); ax.axvline(12, ls="--", c="r", label="champion's n=12")
    ax.set_xlabel("top-n features"); ax.set_ylabel("cumulative importance")
    ax.set_title("Feature-importance elbow curve"); ax.legend(); ax.grid(alpha=.3)
    plt.tight_layout(); plt.savefig(OUT / "elbow_curve.png", dpi=120); plt.close()

    fig, ax = plt.subplots(figsize=(7, 4))
    ax.plot([r[0] for r in sweep], [r[1] for r in sweep], marker=".")
    ax.axvline(buy_thr, ls="--", c="r"); ax.set_title("Threshold sweep (train Sharpe)")
    ax.set_xlabel("|threshold|"); ax.set_ylabel("Sharpe"); ax.grid(alpha=.3)
    plt.tight_layout(); plt.savefig(OUT / "threshold_sweep.png", dpi=120); plt.close()

    corr = data[feat_cols + ["target"]].corr()["target"].drop("target").sort_values(key=abs, ascending=False)
    corr.head(15).to_frame("corr_with_next_hour_change").to_csv(OUT / "top_correlations.csv")


# ------------------------------------------------------------------- MAIN
def main():
    df, data, feat_cols, Xtr, Xte, ytr, yte = build_dataset()
    model = train_model(Xtr, ytr)
    pred_tr, pred_te = model.predict(Xtr), model.predict(Xte)

    buy_thr, sweep = sweep_thresholds(pred_tr, ytr.values)
    pos_te = run_positions(pred_te, buy_thr)
    test = backtest(pos_te, yte.values)
    hold = backtest(np.ones(len(yte)), yte.values)
    test["buy_hold_return_pct"] = hold["total_return_pct"]
    test["threshold_pct"] = round(buy_thr * 100, 3)
    test["assumption_fee_per_trade"] = FEE

    print("\n=== SPRINT A — OUT-OF-SAMPLE RESULTS (honest, like the champion's forward test) ===")
    for k, v in test.items():
        print(f"  {k:28s}: {v}")
    (OUT / "metrics.json").write_text(json.dumps(test, indent=2))
    make_plots(df, data, feat_cols, model, Xte, yte, pred_te, pos_te, sweep, buy_thr)
    print(f"\n[done] artifacts in {OUT}")


if __name__ == "__main__":
    main()