TradeBot/tools/train_lstm_forecaster.py

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import replace
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from crypto_spot_bot.bybit import BybitClient
from crypto_spot_bot.config import Settings, load_settings
from crypto_spot_bot.time_series import _log_returns, _validate_candidates


def main() -> None:
    args = _parse_args()
    settings = load_settings(args.env)
    client = BybitClient(settings)
    symbols = _symbols(args.symbols, settings, client)
    interval = args.interval or settings.base_interval
    output = Path(args.output) if args.output else settings.time_series_lstm_model_path

    artifact: dict[str, Any] = {
        "version": 1,
        "type": "lstm_reservoir_ridge_params",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "interval": interval,
        "limit": args.limit,
        "symbols": {},
    }

    for symbol in symbols:
        result = _train_symbol(
            client=client,
            settings=settings,
            symbol=symbol,
            interval=interval,
            limit=args.limit,
            lookbacks=_ints(args.lookbacks),
            units_values=_ints(args.units),
            ridges=_floats(args.ridges),
        )
        if result is None:
            print(f"{symbol}: skipped, not enough candles or returns")
            continue
        artifact["symbols"][symbol] = result
        print(
            f"{symbol}: lookback={result['lookback']} units={result['units']} "
            f"ridge={result['ridge']} mae={result['validation_mae_percent']:.5f}% "
            f"baseline={result['baseline_mae_percent']:.5f}% skill={result['skill']:.4f}"
        )

    output.parent.mkdir(parents=True, exist_ok=True)
    tmp_output = output.with_name(f"{output.name}.tmp")
    tmp_output.write_text(json.dumps(artifact, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    tmp_output.replace(output)
    print(f"saved {output}")


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train lightweight LSTM forecast params on Bybit spot candles.")
    parser.add_argument("--env", default=None, help="Path to .env file.")
    parser.add_argument("--symbols", default="", help="Comma-separated symbols. Defaults to configured or popular spot pairs.")
    parser.add_argument("--interval", default="", help="Bybit kline interval. Defaults to BASE_INTERVAL.")
    parser.add_argument("--limit", type=int, default=1000, help="Kline limit per symbol.")
    parser.add_argument("--lookbacks", default="16,32", help="Comma-separated LSTM lookback candidates.")
    parser.add_argument("--units", default="4,6", help="Comma-separated LSTM unit candidates.")
    parser.add_argument("--ridges", default="0.001", help="Comma-separated ridge candidates.")
    parser.add_argument("--output", default="", help="Output JSON path. Defaults to TIME_SERIES_LSTM_MODEL_PATH.")
    return parser.parse_args()


def _symbols(raw: str, settings: Settings, client: BybitClient) -> list[str]:
    if raw.strip():
        return [item.strip().upper() for item in raw.split(",") if item.strip()]
    if settings.symbols:
        return list(settings.symbols)
    return client.popular_spot_symbols(settings.top_symbols_count)


def _train_symbol(
    *,
    client: BybitClient,
    settings: Settings,
    symbol: str,
    interval: str,
    limit: int,
    lookbacks: list[int],
    units_values: list[int],
    ridges: list[float],
) -> dict[str, Any] | None:
    candles = client.klines(symbol, interval, limit)
    closes = [float(candle.close) for candle in candles if candle.close > 0]
    returns = _log_returns(closes)
    if len(returns) < 80:
        return None
    validation_window = min(max(8, settings.time_series_validation_window), max(8, len(returns) // 3))
    best: dict[str, Any] | None = None
    for lookback in lookbacks:
        for units in units_values:
            for ridge in ridges:
                candidate_settings = replace(
                    settings,
                    time_series_lstm_enabled=True,
                    time_series_lstm_lookback=lookback,
                    time_series_lstm_units=units,
                    time_series_lstm_ridge=ridge,
                )
                candidates = _validate_candidates(returns, validation_window, candidate_settings, symbol, {})
                baseline = next((item for item in candidates if item["model"] == "naive"), None)
                lstm = next((item for item in candidates if item["model"] == "lstm"), None)
                if baseline is None or lstm is None:
                    continue
                baseline_mae = float(baseline["mae"])
                lstm_mae = float(lstm["mae"])
                skill = (baseline_mae - lstm_mae) / baseline_mae if baseline_mae > 0 else 0.0
                row = {
                    "lookback": lookback,
                    "units": units,
                    "ridge": ridge,
                    "validation_mae_percent": lstm_mae * 100,
                    "baseline_mae_percent": baseline_mae * 100,
                    "skill": skill,
                    "candles": len(candles),
                    "returns": len(returns),
                }
                if best is None or lstm_mae < best["validation_mae_percent"] / 100:
                    best = row
    return best


def _ints(raw: str) -> list[int]:
    return [int(item.strip()) for item in raw.split(",") if item.strip()]


def _floats(raw: str) -> list[float]:
    return [float(item.strip()) for item in raw.split(",") if item.strip()]


if __name__ == "__main__":
    main()