tools/tool_ml_train.py

"""
ML training pipeline: feature selection, cross-validation, model building,
and model serialisation for both stock and ETF pipelines.
"""
import os
import warnings
import numpy as np
import pandas as pd
import yfinance as yf
import joblib
import json
from datetime import datetime

from sklearn.ensemble          import (GradientBoostingClassifier,
                                       RandomForestClassifier,
                                       StackingClassifier)
from sklearn.linear_model      import LogisticRegression
from sklearn.preprocessing     import StandardScaler
from sklearn.model_selection   import TimeSeriesSplit
from sklearn.calibration       import CalibratedClassifierCV
from sklearn.feature_selection import SelectFromModel
from sklearn.metrics           import classification_report, balanced_accuracy_score

from tools.tool_ml_features import (
    DATA_DIR,
    ETF_MODEL_PATH, ETF_SCALER_PATH, ETF_SELECTOR_PATH, ETF_META_PATH,
    STOCK_MODEL_PATH, STOCK_SCALER_PATH, STOCK_SELECTOR_PATH, STOCK_META_PATH,
    _build_features, _build_labels,
)


# ── Feature selection ─────────────────────────────────────────────────────────

def _select_features(X: np.ndarray, y: np.ndarray,
                     feature_names: list) -> tuple:
    """
    Use a quick Random Forest to rank features by importance.
    Drop the bottom third — they add noise, not signal.

    This is called ONCE during training and saved to disk.
    At prediction time we just apply the saved selector.

    Why this helps:
      Financial data is noisy. With 17 features and only 1250 rows,
      weak features can overfit. Selecting the top 10-12 improves
      out-of-sample performance consistently.
    """
    selector_model = RandomForestClassifier(
        n_estimators=100, max_depth=4, random_state=42, n_jobs=-1
    )
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        selector_model.fit(X, y)

    selector = SelectFromModel(
        selector_model,
        threshold="0.5*mean",    # keep features above half average importance
        prefit=True,
    )

    selected_mask  = selector.get_support()
    selected_names = [n for n, s in zip(feature_names, selected_mask) if s]
    dropped_names  = [n for n, s in zip(feature_names, selected_mask) if not s]

    print(f"  Features kept:    {selected_names}")
    print(f"  Features dropped: {dropped_names}")

    return selector, selector.transform(X)


# ── Walk-forward validation ───────────────────────────────────────────────────

def _walk_forward_cv(X: np.ndarray, y: np.ndarray,
                     model, n_splits: int = 5,
                     gap: int = 10) -> list[float]:
    """
    Walk-forward cross-validation with a GAP between train and test.

    The gap is critical. Without it:
      Train ends on day 240, test starts on day 241.
      But our labels look 10 days forward — so day 240's label uses
      data from day 250, which overlaps with the test window.
      That's leakage. The gap prevents it.

    With gap=10:
      Train ends on day 240, test starts on day 251.
      Clean separation. Honest accuracy.

    Returns balanced accuracy per fold (not raw accuracy — balanced
    accuracy accounts for class imbalance so HOLD-heavy datasets
    don't inflate the score).
    """
    tscv   = TimeSeriesSplit(n_splits=n_splits, gap=gap)
    scores = []

    for fold, (tr_idx, te_idx) in enumerate(tscv.split(X)):
        X_tr, X_te = X[tr_idx], X[te_idx]
        y_tr, y_te = y[tr_idx], y[te_idx]

        # Need at least 2 classes in train to fit
        if len(np.unique(y_tr)) < 2:
            continue

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            model.fit(X_tr, y_tr)

        preds = model.predict(X_te)
        score = balanced_accuracy_score(y_te, preds)
        scores.append(score)
        print(f"  Fold {fold+1}: balanced accuracy = {score:.3f}")

    return scores


# ── Stacked model builder ─────────────────────────────────────────────────────

def _build_stacked_model(cv: int = 3):
    """
    Stack 3 diverse base models under a Logistic Regression meta-learner.

    Why these three?
      GradientBoosting — sequential tree boosting, best on tabular data
      RandomForest     — parallel bagging, decorrelated from GBM
      LogisticRegression — linear model, provides a different "view"

    Diversity matters. If all base models are GBM variants they'll
    make the same mistakes. GBM + RF + LR fail on different examples,
    so the meta-learner can learn from their disagreements.

    The meta-learner (also LogisticRegression) takes the probability
    outputs of all three and learns how to weight them optimally.
    cv=3 in StackingClassifier means it uses 3-fold CV internally
    to generate the meta-features without leakage.
    """
    base_estimators = [
        ("gbm", GradientBoostingClassifier(
            n_estimators=200,
            max_depth=3,           # shallow = resistant to overfit
            learning_rate=0.05,
            subsample=0.8,
            min_samples_leaf=20,
            random_state=42,
        )),
        ("rf", RandomForestClassifier(
            n_estimators=200,
            max_depth=5,
            min_samples_leaf=15,
            class_weight="balanced",
            random_state=42,
            n_jobs=-1,
        )),
        ("lr", LogisticRegression(
            C=0.1,                   # strong regularisation for small dataset
            class_weight="balanced",  # prevents collapsing to HOLD
            max_iter=1000,
            random_state=42,
        )),
    ]

    meta_learner = LogisticRegression(C=1.0, class_weight="balanced", max_iter=500, random_state=42)

    stacked = StackingClassifier(
        estimators      = base_estimators,
        final_estimator = meta_learner,
        cv              = cv,
        stack_method    = "predict_proba",
        n_jobs          = -1,
        passthrough     = False,
    )

    # Wrap in probability calibration
    # Isotonic regression corrects overconfident probabilities.
    calibrated = CalibratedClassifierCV(
        stacked,
        method = "isotonic",
        cv     = cv,
    )

    return calibrated


# ── Training ──────────────────────────────────────────────────────────────────

def train_model(training_ticker: str = "SPY",
                period: str           = "5y",
                forward_days: int     = 10,
                threshold: float      = 0.03,
                extra_tickers: list   = None,
                model_path: str       = None,
                scaler_path: str      = None,
                selector_path: str    = None,
                meta_path: str        = None) -> dict:
    """
    Full training pipeline with feature selection, walk-forward CV,
    stacked model, and calibration.
    """
    _model_path    = model_path    or STOCK_MODEL_PATH
    _scaler_path   = scaler_path   or STOCK_SCALER_PATH
    _selector_path = selector_path or STOCK_SELECTOR_PATH
    _meta_path     = meta_path     or STOCK_META_PATH

    if extra_tickers is None:
        extra_tickers = [
            "NVDA", "AAPL", "MSFT", "TSLA", "AMZN", "META", "GOOGL", "AMD",
            "QQQ", "VOO", "VTI",
            "SMH", "XLK", "XLF", "XLE", "ITA", "GLD",
            "VWRP.L", "WENS.L", "SGLN.L", "FTAL.L", "SEMI.L",
        ]

    os.makedirs(DATA_DIR, exist_ok=True)

    # ── Fetch VIX once for market-regime features ─────────────────────────────
    print("\nFetching VIX (market regime)...")
    try:
        vix_hist = yf.Ticker("^VIX").history(period=period)
        vix = vix_hist["Close"].copy() if not vix_hist.empty else None
        if vix is not None:
            vix.index = pd.to_datetime(vix.index).tz_localize(None)
            full_idx = pd.date_range(vix.index.min(), vix.index.max(), freq="D")
            vix = vix.reindex(full_idx).ffill()
    except Exception:
        vix = None

    # ── Pool multiple tickers for more training data ──────────────────────────
    all_features, all_labels = [], []
    for tkr in [training_ticker] + extra_tickers:
        print(f"\nFetching {tkr} ({period})...")
        hist = yf.Ticker(tkr).history(period=period)
        if hist.empty:
            print(f"  skipped (no data)")
            continue
        hist.index = pd.to_datetime(hist.index).tz_localize(None)

        feat = _build_features(hist, vix=vix)
        lbl  = _build_labels(hist, forward_days, threshold).reindex(feat.index).dropna()
        feat = feat.reindex(lbl.index)
        if len(lbl) < 50:
            print(f"  skipped (only {len(lbl)} rows)")
            continue

        b = (lbl == 1).sum(); s = (lbl == -1).sum(); h = (lbl == 0).sum()
        print(f"  {len(lbl)} rows | BUY:{b} HOLD:{h} SELL:{s}")
        all_features.append(feat)
        all_labels.append(lbl)

    features      = pd.concat(all_features)
    labels_series = pd.concat(all_labels)

    feature_names = features.columns.tolist()
    X_raw = features.values
    y     = labels_series.values

    n_buy  = (y ==  1).sum()
    n_sell = (y == -1).sum()
    n_hold = (y ==  0).sum()
    print(f"\nCombined dataset: {len(y)} rows | BUY:{n_buy} HOLD:{n_hold} SELL:{n_sell}")

    # Scale first (needed for LR in the stack)
    scaler   = StandardScaler()
    X_scaled = scaler.fit_transform(X_raw)

    # Feature selection
    print("\nSelecting features...")
    selector, X_selected = _select_features(X_scaled, y, feature_names)

    # Walk-forward CV on a single GBM (fast proxy for the stack)
    print("\nWalk-forward cross-validation...")
    proxy = GradientBoostingClassifier(
        n_estimators=100, max_depth=3, learning_rate=0.05,
        subsample=0.8, random_state=42
    )
    cv_scores = _walk_forward_cv(X_selected, y, proxy, n_splits=5, gap=forward_days)
    mean_cv   = float(np.mean(cv_scores))
    print(f"\nCV balanced accuracy: {mean_cv:.3f} ({mean_cv*100:.1f}%)")

    min_class = int(min(n_buy, n_sell, n_hold))
    if min_class >= 3:
        safe_cv = min(3, min_class)
        print(f"\nTraining stacked model (GBM + RF + LR, cv={safe_cv})...")
        model = _build_stacked_model(cv=safe_cv)
    else:
        print(f"\nDataset too small for stacking (min class={min_class}) — using simple GBM...")
        model = GradientBoostingClassifier(
            n_estimators=100, max_depth=3, learning_rate=0.05,
            subsample=0.8, random_state=42,
        )
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        model.fit(X_selected, y)

    # Per-class performance on full training set
    preds  = model.predict(X_selected)
    report = classification_report(y, preds, target_names=["SELL","HOLD","BUY"],
                                   output_dict=True, zero_division=0)

    # Save everything
    joblib.dump(model,    _model_path)
    joblib.dump(scaler,   _scaler_path)
    joblib.dump(selector, _selector_path)

    meta = {
        "trained_on":        training_ticker,
        "trained_at":        datetime.now().isoformat(),
        "n_samples":         len(y),
        "n_features_raw":    len(feature_names),
        "n_features_kept":   int(selector.get_support().sum()),
        "forward_days":      forward_days,
        "threshold":         threshold,
        "cv_balanced_acc":   round(mean_cv, 4),
        "cv_fold_scores":    [round(s, 4) for s in cv_scores],
        "class_distribution":{"BUY": int(n_buy), "HOLD": int(n_hold), "SELL": int(n_sell)},
        "per_class": {
            "BUY":  {"precision": round(report["BUY"]["precision"],  3),
                     "recall":    round(report["BUY"]["recall"],     3)},
            "SELL": {"precision": round(report["SELL"]["precision"], 3),
                     "recall":    round(report["SELL"]["recall"],    3)},
            "HOLD": {"precision": round(report["HOLD"]["precision"], 3),
                     "recall":    round(report["HOLD"]["recall"],    3)},
        }
    }
    with open(_meta_path, "w") as f:
        json.dump(meta, f, indent=2)

    print(f"\nDone. Model saved.")
    print(f"BUY  — precision: {meta['per_class']['BUY']['precision']:.0%}  recall: {meta['per_class']['BUY']['recall']:.0%}")
    print(f"SELL — precision: {meta['per_class']['SELL']['precision']:.0%}  recall: {meta['per_class']['SELL']['recall']:.0%}")
    print(f"HOLD — precision: {meta['per_class']['HOLD']['precision']:.0%}  recall: {meta['per_class']['HOLD']['recall']:.0%}")

    return meta


def train_etf_model() -> dict:
    """
    Dedicated pipeline for ETF/ETC instruments.

    Why different from stocks?
    - ETFs trend more smoothly (less idiosyncratic noise)
    - Smaller moves are still tradeable (2% threshold vs 3%)
    - Longer horizon (20 days) suits buy-and-hold ETF owners
    - Training data: only broad-market and sector ETFs
    """
    etf_tickers = [
        "SPY", "QQQ", "VOO", "VTI", "IVV",
        "XLK", "XLE", "XLF", "XLV", "XLI", "XLU", "XLP",
        "SMH", "ITA", "GLD", "SLV",
        "EFA", "EEM", "VEA",
        "VWRP.L", "FTAL.L", "SEMI.L", "WENS.L", "DFND.L", "SGLN.L",
    ]
    return train_model(
        training_ticker = "SPY",
        period          = "5y",
        forward_days    = 20,
        threshold       = 0.02,
        extra_tickers   = etf_tickers,
        model_path      = ETF_MODEL_PATH,
        scaler_path     = ETF_SCALER_PATH,
        selector_path   = ETF_SELECTOR_PATH,
        meta_path       = ETF_META_PATH,
    )