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.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(feature_matrix: np.ndarray, out_path: Path, epochs: int = 20) -> None:
    """Train the Stage 3 TCN on the feature matrix.

    Uses self-supervised pretraining: predict the next timestep's features
    from the current window. This learns the "normal" pattern.
    """
    logger.info("=" * 60)
    logger.info("TRAINING STAGE 3 — TEMPORAL CONVOLUTIONAL NETWORK")
    logger.info("=" * 60)

    # Use the first 17 features (F1-F4) for Stage 3
    stage3_features = feature_matrix[:, :17]
    logger.info("Stage 3 input shape: %s", stage3_features.shape)

    # Build sequences: sliding window of 100 timesteps
    seq_len = 100
    n_sequences = len(stage3_features) - seq_len
    if n_sequences < 100:
        logger.warning("Not enough data for TCN training (need >%d ticks, got %d)",
                       seq_len, len(stage3_features))
        return

    logger.info("Building %d sequences of length %d...", n_sequences, seq_len)
    sequences = []
    for i in range(0, n_sequences, max(1, n_sequences // 5000)):  # limit to 5000 sequences
        seq = stage3_features[i:i + seq_len]
        sequences.append(seq)

    sequences = np.array(sequences, dtype=np.float32)
    logger.info("Sequences shape: %s", sequences.shape)

    # Split 80/20 train/val
    split = int(len(sequences) * 0.8)
    train_data = sequences[:split]
    val_data = sequences[split:]
    logger.info("Train: %d sequences, Val: %d sequences", len(train_data), len(val_data))

    # Train the TCN
    config = TCNConfig(sequence_length=seq_len)
    model = Stage3TCN(config)
    model.train(train_data, val_data, 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(feature_matrix, out_dir / "stage3_tcn.pt", epochs=args.epochs)

    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())