ahad-quant/prepare_sequences.py

"""
AHAD QUANT Forex V5 — Sequence Dataset Preparation
Converts per-pair tabular features → sliding windows (B, SEQ_LEN, 62) for DL models.

Window i  = X[i : i+SEQ_LEN]          → sequence input for TFT / TransformerGRU
Label i   = y[i+SEQ_LEN-1]            → same label as the last step in the window
Tab row i = X[i+SEQ_LEN-1]            → tabular features at the end of the window
                                          (used to align tabular models in train.py)

Returned arrays are ALL aligned on the same M samples, making it trivial to
train tabular and DL models together with consistent train/val/test splits.

Usage (called from train.py when HAS_DL=True):
    X_tab, X_seq, y = prepare_seq_dataset()
    # X_tab : (M, 62)          — same features as prepare_dataset(), but aligned
    # X_seq : (M, 168, 62)     — sequence tensors
    # y     : (M,)             — labels
"""

import json
import os

import numpy as np

import config
from features import build_features, NUM_FEATURES

# ─── Constants ───────────────────────────────────────────────────────────────

SEQ_LEN  = 168          # 7 days × 24h — window length fed into TFT / TGRU

# Memory cap: 25 pairs × MAX_SEQ_PER_PAIR × 168 × 62 × 4B ~ 1.5 GB at 6000
# Increase if you have >16 GB RAM; decrease if training crashes with OOM.
MAX_SEQ_PER_PAIR: int = int(os.environ.get("MAX_SEQ_PER_PAIR", "6000"))
LOOKAHEAD = 3           # must match LOOKAHEAD in train.py
WARMUP    = 30          # feature warm-up candles (same as train.py)
MIN_CANDLES = SEQ_LEN + LOOKAHEAD + WARMUP + 10   # minimum pair length


# ─── Core builders ───────────────────────────────────────────────────────────

def _load_candles(pair: str) -> list | None:
    path = os.path.join(config.DATA_DIR, f"{pair}_1h.json")
    if not os.path.exists(path):
        return None
    with open(path) as f:
        return json.load(f)


def _make_labels(close: np.ndarray, lookahead: int = LOOKAHEAD) -> np.ndarray:
    labels = np.zeros(len(close))
    for i in range(len(close) - lookahead):
        labels[i] = 1.0 if close[i + lookahead] > close[i] else 0.0
    return labels


def build_sequences(
    X: np.ndarray,
    y: np.ndarray,
    seq_len: int = SEQ_LEN,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convert a single pair's aligned (X, y) arrays into sliding windows.

    Parameters
    ----------
    X       : (N, 62) feature matrix — already NaN-free and warmed up
    y       : (N,)    label vector
    seq_len : window length

    Returns
    -------
    X_seq  : (M, seq_len, 62)  — sequences
    X_tab  : (M, 62)           — tabular row at end of each window (aligned)
    y_seq  : (M,)              — labels aligned with end of each window

    where M = N - seq_len - LOOKAHEAD
    (last LOOKAHEAD samples excluded to avoid future-leakage in labels)
    """
    N = len(X)
    # Last valid window ends at index N-LOOKAHEAD-1
    n_windows = N - seq_len - LOOKAHEAD
    if n_windows <= 0:
        raise ValueError(
            f"Not enough samples ({N}) for seq_len={seq_len} + lookahead={LOOKAHEAD}. "
            f"Need at least {seq_len + LOOKAHEAD + 1}."
        )

    X_seq = np.empty((n_windows, seq_len, X.shape[1]), dtype=np.float32)
    for i in range(n_windows):
        X_seq[i] = X[i : i + seq_len]

    # Tabular: last row of each window
    X_tab = X[seq_len - 1 : seq_len - 1 + n_windows].astype(np.float32)

    # Label: at the end of the window (same index as X_tab row)
    y_seq = y[seq_len - 1 : seq_len - 1 + n_windows].astype(np.float32)

    return X_seq, X_tab, y_seq


# ─── Full dataset builder ────────────────────────────────────────────────────

def prepare_seq_dataset(
    seq_len: int = SEQ_LEN,
    verbose: bool = True,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Build the full aligned (X_tab, X_seq, y) dataset across all pairs.

    Mirrors prepare_dataset() in train.py but:
      1. Uses a sliding window of `seq_len` candles
      2. Returns BOTH tabular (aligned) AND sequence data
      3. All three output arrays share the same index space → consistent splits

    Returns
    -------
    X_tab : (M_total, 62)          — tabular features at end of each window
    X_seq : (M_total, seq_len, 62) — sequence windows
    y     : (M_total,)             — binary labels
    """
    all_X_tab: list[np.ndarray] = []
    all_X_seq: list[np.ndarray] = []
    all_y:     list[np.ndarray] = []

    # Reference pair for correlation feature (same logic as train.py)
    _ref = "EURUSD"
    ref_data  = _load_candles(_ref)
    ref_close = np.array([c["c"] for c in ref_data]) if ref_data else None

    for pair in config.COINS:
        data = _load_candles(pair)
        if data is None or len(data) < MIN_CANDLES:
            if verbose:
                print(f"  [{pair:8s}] skipped — insufficient data ({len(data) if data else 0} candles)")
            continue

        close = np.array([c["c"] for c in data])
        pair_ref = (
            ref_close[-len(close) :]
            if ref_close is not None and len(ref_close) >= len(close)
            else None
        )

        X_all = build_features(data, btc_closes=pair_ref)
        y_all = _make_labels(close)

        # Apply same warm-up crop as train.py
        X_all = X_all[WARMUP : -LOOKAHEAD]
        y_all = y_all[WARMUP : -LOOKAHEAD]

        # Drop NaN rows (should be none after warm-up, but safety check)
        valid = ~np.isnan(X_all).any(axis=1)
        X_all, y_all = X_all[valid], y_all[valid]

        if len(X_all) < MIN_CANDLES:
            if verbose:
                print(f"  [{pair:8s}] skipped after cleaning ({len(X_all)} samples left)")
            continue

        try:
            X_seq_pair, X_tab_pair, y_seq_pair = build_sequences(X_all, y_all, seq_len)
        except ValueError as e:
            if verbose:
                print(f"  [{pair:8s}] skipped — {e}")
            continue

        # ── Memory cap: stratified sample to MAX_SEQ_PER_PAIR ──────────
        n_pair = len(y_seq_pair)
        if n_pair > MAX_SEQ_PER_PAIR:
            rng = np.random.default_rng(seed=42)
            # Preserve class balance (stratified)
            idx_pos = np.where(y_seq_pair == 1)[0]
            idx_neg = np.where(y_seq_pair == 0)[0]
            half    = MAX_SEQ_PER_PAIR // 2
            sel_pos = rng.choice(idx_pos, size=min(half, len(idx_pos)), replace=False)
            sel_neg = rng.choice(idx_neg, size=min(half, len(idx_neg)), replace=False)
            sel     = np.sort(np.concatenate([sel_pos, sel_neg]))
            X_seq_pair = X_seq_pair[sel]
            X_tab_pair = X_tab_pair[sel]
            y_seq_pair = y_seq_pair[sel]
            if verbose:
                print(f"    → sampled {len(y_seq_pair):,} / {n_pair:,} sequences (RAM cap)")

        all_X_tab.append(X_tab_pair)
        all_X_seq.append(X_seq_pair)
        all_y.append(y_seq_pair)

        if verbose:
            print(f"  [{pair:8s}] {len(y_seq_pair):>7,} sequences")

    if not all_X_tab:
        raise RuntimeError(
            "No pairs produced sequences. Check DATA_DIR and MIN_CANDLES."
        )

    X_tab = np.concatenate(all_X_tab, axis=0)
    X_seq = np.concatenate(all_X_seq, axis=0)
    y     = np.concatenate(all_y,     axis=0)

    if verbose:
        print(f"\n  Total : {len(y):,} sequences | shape X_seq={X_seq.shape} | "
              f"{y.mean():.2%} long labels")

    return X_tab, X_seq, y


# ─── Sequence normalisation helpers ─────────────────────────────────────────

def fit_seq_scaler(X_tab: np.ndarray):
    """
    Fit a StandardScaler on tabular training data.
    The same scaler is applied to sequences by reshaping (B, T, F) → (B*T, F).
    Returns a fitted sklearn StandardScaler.
    """
    from sklearn.preprocessing import StandardScaler
    scaler = StandardScaler()
    scaler.fit(X_tab)
    return scaler


def transform_sequences(scaler, X_seq: np.ndarray) -> np.ndarray:
    """
    Apply a fitted StandardScaler to a 3-D sequence array.

    Parameters
    ----------
    scaler : fitted sklearn StandardScaler
    X_seq  : (B, T, F) float32

    Returns
    -------
    X_norm : (B, T, F) float32
    """
    B, T, F = X_seq.shape
    flat    = X_seq.reshape(-1, F)
    normed  = scaler.transform(flat).astype(np.float32)
    return normed.reshape(B, T, F)