scripts/train_models.py

#!/usr/bin/env python3
"""Train the detector models on real Binance crash data.

Trains:
    1. Stage 2 Isolation Forest on normal traffic (unsupervised)
    2. Stage 3 TCN on labeled crash windows (self-supervised + supervised)

Usage:
    python scripts/train_models.py --data data/parquet/BTCUSDT_2021-05-18.parquet --out models/
    python scripts/train_models.py --data data/parquet/BTCUSDT_2024-01-15.parquet --out models/ --epochs 20

The training data should be a NORMAL day (not a crash day) so the models
learn what "normal" looks like. Then the backtest on crash days will detect
the anomalies.
"""
import argparse
import logging
import sys
from pathlib import Path

import numpy as np
import pandas as pd

# Insert the ml directory at the FRONT of sys.path
ML_DIR = Path(__file__).resolve().parent.parent / "ml"
sys.path.insert(0, str(ML_DIR))
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from flash_crash_watchdog.data.historical_loader import df_to_ticks, load_parquet
from flash_crash_watchdog.data.windows import build_windows_from_df
from flash_crash_watchdog.features import FEATURE_NAMES, FeatureExtractor
from flash_crash_watchdog.models.stage2_isolation_forest import Stage2IsolationForest
from flash_crash_watchdog.models.stage3_tcn import Stage3TCN, TCNConfig

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)


def extract_feature_matrix(df: pd.DataFrame, max_ticks: int = 100_000) -> np.ndarray:
    """Extract a feature matrix from a DataFrame for training.

    Args:
        df: Historical tick data.
        max_ticks: Maximum number of ticks to process (for speed).

    Returns:
        Matrix of shape (n_ticks, 20) — the feature vector per tick.
    """
    logger.info("Extracting features from %d ticks (max %d)...", len(df), max_ticks)

    # Sample if too many ticks
    if len(df) > max_ticks:
        # Sample evenly across the day to get representative data
        indices = np.linspace(0, len(df) - 1, max_ticks, dtype=int)
        df_sample = df.iloc[indices].copy()
        logger.info("Sampled down to %d ticks (evenly spaced)", len(df_sample))
    else:
        df_sample = df

    extractor = FeatureExtractor()
    features_list = []

    for i, tick in enumerate(df_to_ticks(df_sample, symbol="TRAIN")):
        if i % 10000 == 0:
            logger.info("  Processing tick %d/%d...", i, len(df_sample))
        features = extractor.extract(tick)
        features_list.append([features.get(f, 0.0) for f in FEATURE_NAMES])

    matrix = np.array(features_list, dtype=np.float32)
    logger.info("Feature matrix shape: %s", matrix.shape)

    # Replace NaN/Inf with 0
    matrix = np.nan_to_num(matrix, nan=0.0, posinf=0.0, neginf=0.0)

    return matrix


def train_stage2_isolation_forest(feature_matrix: np.ndarray, out_path: Path) -> None:
    """Train the Stage 2 Isolation Forest on normal data."""
    logger.info("=" * 60)
    logger.info("TRAINING STAGE 2 — ISOLATION FOREST")
    logger.info("=" * 60)

    # Use the first 12 features (F1 + F2) for Stage 2
    stage2_features = feature_matrix[:, :12]
    logger.info("Stage 2 input shape: %s", stage2_features.shape)

    model = Stage2IsolationForest(n_estimators=100, contamination=0.05)
    model.fit(stage2_features)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    model.save(out_path)
    logger.info("Stage 2 model saved to %s", out_path)


def train_stage3_tcn(df: pd.DataFrame, out_path: Path, epochs: int = 20,
                     max_ticks: int = 100_000) -> None:
    """Train the Stage 3 TCN on labeled crash windows.

    Built from the same tick frame; window labels come from a forward-looking
    2%-drop-in-5s rule (see data.windows.build_windows_from_df). A "normal" day
    has no crash labels, so Stage 3 is skipped with a pointer rather than
    silently trained to an all-zero target.
    """
    logger.info("=" * 60)
    logger.info("TRAINING STAGE 3 — TEMPORAL CONVOLUTIONAL NETWORK")
    logger.info("=" * 60)

    try:
        windows, labels, _feature_names = build_windows_from_df(df, max_ticks=max_ticks)
    except FileNotFoundError as e:
        logger.warning("Skipping Stage 3: %s", e)
        return
    n_pos = int(np.sum(labels))
    if n_pos == 0:
        logger.warning(
            "No crash labels in this day — the TCN is a crash classifier and cannot be "
            "trained on normal-only data. Skip Stage 3 (or run scripts/train_tcn_windows.py "
            "on labeled crash windows).")
        return
    logger.info("Stage 3 windows: %s (%d positive)", windows.shape, n_pos)

    seq_len = windows.shape[1]
    config = TCNConfig(sequence_length=seq_len)
    model = Stage3TCN(config)
    model.train(windows, labels, epochs=epochs)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    model.save(out_path)
    logger.info("Stage 3 model saved to %s", out_path)


def main() -> int:
    parser = argparse.ArgumentParser(description="Train detector models on real data")
    parser.add_argument("--data", required=True, help="Parquet file of NORMAL market data")
    parser.add_argument("--out", default="models/", help="Output directory for trained models")
    parser.add_argument("--epochs", type=int, default=20, help="TCN training epochs")
    parser.add_argument("--max-ticks", type=int, default=100_000,
                        help="Max ticks to process (for speed)")
    args = parser.parse_args()

    # Load data
    df = load_parquet(args.data)
    logger.info("Loaded %d ticks from %s", len(df), args.data)

    # Extract features
    feature_matrix = extract_feature_matrix(df, max_ticks=args.max_ticks)

    # Train Stage 2
    out_dir = Path(args.out)
    train_stage2_isolation_forest(feature_matrix, out_dir / "stage2_isolation_forest.joblib")

    # Train Stage 3
    train_stage3_tcn(df, out_dir / "stage3_tcn.pt", epochs=args.epochs,
                     max_ticks=args.max_ticks)

    logger.info("=" * 60)
    logger.info("TRAINING COMPLETE")
    logger.info("  Models saved to: %s", out_dir.resolve())
    logger.info("  Stage 2: stage2_isolation_forest.joblib")
    logger.info("  Stage 3: stage3_tcn.pt")
    logger.info("=" * 60)
    logger.info("")
    logger.info("Next step: re-run the backtest on crash data to see alerts!")
    logger.info("  python scripts/run_backtest.py --data data/parquet/BTCUSDT_2021-05-19.parquet")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())