ml/features.py

"""Sliding-window feature engineering for roulette next-spin prediction.

Features are built per source sequence so windows never cross source boundaries.
Every row of the feature matrix holds the last ``WINDOW`` winning numbers plus
derived counts; labels are the next spin's number, color, parity, dozen, column.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Final

import numpy as np
import pandas as pd

WINDOW: Final[int] = 10

RED_NUMBERS: Final[frozenset[int]] = frozenset(
    {1, 3, 5, 7, 9, 12, 14, 16, 18, 19, 21, 23, 25, 27, 30, 32, 34, 36}
)

NUMBER_CLASSES: Final[int] = 37
COLOR_CLASSES: Final[tuple[str, ...]] = ("red", "black", "green")
PARITY_CLASSES: Final[tuple[str, ...]] = ("odd", "even", "none")
DOZEN_CLASSES: Final[tuple[str, ...]] = ("first", "second", "third", "none")
COLUMN_CLASSES: Final[tuple[str, ...]] = ("first", "second", "third", "none")

TARGETS: Final[tuple[str, ...]] = ("number", "color", "parity", "dozen", "column")


@dataclass(frozen=True)
class WindowedDataset:
    X: np.ndarray
    y: dict[str, np.ndarray]
    feature_names: list[str]


def derive_color(n: int) -> str:
    if n == 0:
        return "green"
    return "red" if n in RED_NUMBERS else "black"


def derive_parity(n: int) -> str:
    if n == 0:
        return "none"
    return "even" if n % 2 == 0 else "odd"


def derive_dozen(n: int) -> str:
    if n == 0:
        return "none"
    if n <= 12:
        return "first"
    if n <= 24:
        return "second"
    return "third"


def derive_column(n: int) -> str:
    if n == 0:
        return "none"
    rem = n % 3
    return "first" if rem == 1 else ("second" if rem == 2 else "third")


def encode_label(value: str, classes: tuple[str, ...]) -> int:
    return classes.index(value)


def build_windows_from_sequence(numbers: np.ndarray, window: int = WINDOW) -> WindowedDataset:
    """Build sliding-window features from a single contiguous number sequence."""
    if len(numbers) <= window:
        empty_y = {name: np.empty(0, dtype=np.int64) for name in TARGETS}
        return WindowedDataset(
            X=np.empty((0, _feature_count(window)), dtype=np.float32),
            y=empty_y,
            feature_names=_feature_names(window),
        )

    n_samples = len(numbers) - window
    X = np.empty((n_samples, _feature_count(window)), dtype=np.float32)
    y_number = np.empty(n_samples, dtype=np.int64)
    y_color = np.empty(n_samples, dtype=np.int64)
    y_parity = np.empty(n_samples, dtype=np.int64)
    y_dozen = np.empty(n_samples, dtype=np.int64)
    y_column = np.empty(n_samples, dtype=np.int64)

    for i in range(n_samples):
        win = numbers[i : i + window]
        target = int(numbers[i + window])
        X[i] = _features_from_window(win)
        y_number[i] = target
        y_color[i] = encode_label(derive_color(target), COLOR_CLASSES)
        y_parity[i] = encode_label(derive_parity(target), PARITY_CLASSES)
        y_dozen[i] = encode_label(derive_dozen(target), DOZEN_CLASSES)
        y_column[i] = encode_label(derive_column(target), COLUMN_CLASSES)

    return WindowedDataset(
        X=X,
        y={
            "number": y_number,
            "color": y_color,
            "parity": y_parity,
            "dozen": y_dozen,
            "column": y_column,
        },
        feature_names=_feature_names(window),
    )


def build_windows_grouped(
    df: pd.DataFrame,
    number_col: str = "number",
    group_col: str = "source",
    window: int = WINDOW,
) -> WindowedDataset:
    """Build windows per source group and concatenate. Never crosses sources."""
    parts: list[WindowedDataset] = []
    for _, group in df.groupby(group_col, sort=False):
        numbers = group[number_col].to_numpy(dtype=np.int64)
        parts.append(build_windows_from_sequence(numbers, window=window))

    non_empty = [p for p in parts if len(p.X) > 0]
    if not non_empty:
        return build_windows_from_sequence(np.empty(0, dtype=np.int64), window=window)

    X = np.vstack([p.X for p in non_empty])
    y = {
        name: np.concatenate([p.y[name] for p in non_empty]) for name in TARGETS
    }
    return WindowedDataset(X=X, y=y, feature_names=non_empty[0].feature_names)


def _features_from_window(win: np.ndarray) -> np.ndarray:
    """Extract features from a window of length WINDOW of integer numbers."""
    window = len(win)
    feats = np.empty(_feature_count(window), dtype=np.float32)

    feats[:window] = win

    red_count = 0
    black_count = 0
    zero_count = 0
    even_count = 0
    odd_count = 0
    low_count = 0
    high_count = 0
    dozen_counts = [0, 0, 0]
    column_counts = [0, 0, 0]
    number_sum = 0

    for n in win:
        n_int = int(n)
        number_sum += n_int
        if n_int == 0:
            zero_count += 1
            continue
        if n_int in RED_NUMBERS:
            red_count += 1
        else:
            black_count += 1
        if n_int % 2 == 0:
            even_count += 1
        else:
            odd_count += 1
        if n_int <= 18:
            low_count += 1
        else:
            high_count += 1
        if n_int <= 12:
            dozen_counts[0] += 1
        elif n_int <= 24:
            dozen_counts[1] += 1
        else:
            dozen_counts[2] += 1
        rem = n_int % 3
        if rem == 1:
            column_counts[0] += 1
        elif rem == 2:
            column_counts[1] += 1
        else:
            column_counts[2] += 1

    offset = window
    feats[offset + 0] = red_count
    feats[offset + 1] = black_count
    feats[offset + 2] = zero_count
    feats[offset + 3] = even_count
    feats[offset + 4] = odd_count
    feats[offset + 5] = low_count
    feats[offset + 6] = high_count
    feats[offset + 7] = dozen_counts[0]
    feats[offset + 8] = dozen_counts[1]
    feats[offset + 9] = dozen_counts[2]
    feats[offset + 10] = column_counts[0]
    feats[offset + 11] = column_counts[1]
    feats[offset + 12] = column_counts[2]
    feats[offset + 13] = number_sum / window
    feats[offset + 14] = int(win[-1])  # last number
    return feats


def _feature_count(window: int) -> int:
    return window + 15


def _feature_names(window: int) -> list[str]:
    lags = [f"lag_{i}" for i in range(window, 0, -1)]
    extras = [
        "red_count",
        "black_count",
        "zero_count",
        "even_count",
        "odd_count",
        "low_count",
        "high_count",
        "dozen1_count",
        "dozen2_count",
        "dozen3_count",
        "col1_count",
        "col2_count",
        "col3_count",
        "mean_number",
        "last_number",
    ]
    return lags + extras