quant-web/engine.py

"""Backtest and optimization engine."""

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pandas as pd
import vectorbt as vbt

from strategies.executor import (
    CUSTOM_TEMPLATE,
    StrategyError,
    load_custom_strategy,
    optimize_custom,
    run_builtin_signals,
    run_custom_signals,
)
from strategies.registry import BuiltinStrategy, get_builtin


@dataclass(frozen=True)
class BacktestResult:
    ticker: str
    strategy_key: str
    params: dict[str, Any]
    sharpe_ratio: float
    sortino_ratio: float
    max_drawdown: float
    total_return: float
    win_rate: float
    total_trades: int
    equity_curve: pd.Series
    price: pd.Series
    entries: pd.Series
    exits: pd.Series


@dataclass(frozen=True)
class OptimizationResult:
    ticker: str
    strategy_key: str
    metric: str
    best_params: dict[str, Any]
    best_score: float
    results: pd.DataFrame
    combinations_tested: int


def parquet_dir() -> Path:
    return Path(os.environ.get("PARQUET_DIR", "/data/parquet"))


def load_ohlcv(ticker: str) -> pd.DataFrame:
    path = parquet_dir() / f"{ticker.upper()}.parquet"
    if not path.exists():
        raise FileNotFoundError(f"No Parquet file for {ticker.upper()} at {path}")

    df = pd.read_parquet(path)
    if "Date" in df.columns:
        df = df.set_index("Date")
    df.index = pd.to_datetime(df.index)
    return df.sort_index()


from metrics import run_from_signals, safe_float as _safe_float
def _portfolio_from_signals(
    close: pd.Series,
    entries: pd.Series,
    exits: pd.Series,
    init_cash: float,
    fees: float,
) -> vbt.Portfolio:
    return vbt.Portfolio.from_signals(
        close,
        entries=entries,
        exits=exits,
        init_cash=init_cash,
        fees=fees,
        freq="1D",
    )


def run_backtest(
    ticker: str,
    strategy_key: str,
    params: dict[str, Any],
    source_code: str | None = None,
    init_cash: float = 10_000.0,
    fees: float = 0.001,
) -> BacktestResult:
    ohlcv = load_ohlcv(ticker)
    close = ohlcv["Close"].astype(float)

    runtime_params = {k: v for k, v in params.items() if k not in ("init_cash", "fees")}
    init_cash = float(params.get("init_cash", init_cash))
    fees = float(params.get("fees", fees))

    if strategy_key == "custom":
        if not source_code:
            raise StrategyError("Custom strategy requires source_code")
        entries, exits, _, merged = run_custom_signals(source_code, ohlcv, runtime_params)
        params = merged
    else:
        builtin = get_builtin(strategy_key)
        entries, exits = run_builtin_signals(builtin, ohlcv, runtime_params)
        params = {**builtin.default_params, **runtime_params}

    portfolio = _portfolio_from_signals(close, entries, exits, init_cash, fees)
    stats = portfolio.stats()
    equity = portfolio.value()
    if isinstance(equity, pd.DataFrame):
        equity = equity.iloc[:, 0]

    return BacktestResult(
        ticker=ticker.upper(),
        strategy_key=strategy_key,
        params=params,
        sharpe_ratio=_safe_float(stats.get("Sharpe Ratio")),
        sortino_ratio=_safe_float(stats.get("Sortino Ratio")),
        max_drawdown=_safe_float(stats.get("Max Drawdown [%]")) / 100.0,
        total_return=_safe_float(stats.get("Total Return [%]")) / 100.0,
        win_rate=_safe_float(stats.get("Win Rate [%]")) / 100.0,
        total_trades=int(stats.get("Total Trades", 0) or 0),
        equity_curve=equity,
        price=close,
        entries=entries,
        exits=exits,
    )


def run_optimization(
    ticker: str,
    strategy_key: str,
    metric: str = "sharpe_ratio",
    init_cash: float = 10_000.0,
    fees: float = 0.001,
    source_code: str | None = None,
    param_grid: dict | None = None,
) -> OptimizationResult:
    ohlcv = load_ohlcv(ticker)
    close = ohlcv["Close"].astype(float)

    if strategy_key == "custom":
        if not source_code:
            raise StrategyError("Custom strategy requires source_code")
        results = optimize_custom(
            source_code,
            close,
            ohlcv,
            init_cash=init_cash,
            fees=fees,
            metric=metric,
            param_grid=param_grid,
        )
    else:
        builtin = get_builtin(strategy_key)
        results = builtin.optimize(close, init_cash, fees, metric, grid_override=param_grid)

    if results.empty:
        raise StrategyError("Optimization produced no valid parameter combinations.")

    best = results.iloc[0]
    param_cols = [c for c in results.columns if c not in {
        "score", "sharpe_ratio", "sortino_ratio", "max_drawdown", "total_return", "win_rate", "total_trades",
    }]
    best_params = {col: best[col] for col in param_cols}

    return OptimizationResult(
        ticker=ticker.upper(),
        strategy_key=strategy_key,
        metric=metric,
        best_params=best_params,
        best_score=float(best["score"]),
        results=results,
        combinations_tested=len(results),
    )


def get_strategy_source(strategy_key: str, source_code: str | None = None) -> str:
    if strategy_key == "custom":
        return source_code or CUSTOM_TEMPLATE
    return get_builtin(strategy_key).source_code