Blame for src/drift_quant/backtesting/engine.py

"""Event-driven long-only portfolio simulation with next-bar-open execution."""

from __future__ import annotations

import math
from collections.abc import Callable
from dataclasses import asdict, dataclass
from datetime import datetime
from statistics import fmean
from typing import Any
from uuid import uuid4

from drift_quant.backtesting.strategies import Strategy
from drift_quant.nse_data.loader import LoadedNseDataset


@dataclass(frozen=True, slots=True)
class BacktestConfig:
    """Visible simulation assumptions; costs are configurable research assumptions."""

    initial_capital: float = 1_000_000.0
    max_gross_exposure: float = 0.95
    transaction_cost_bps: float = 10.0
    slippage_bps: float = 5.0
    allow_overnight: bool = True
    force_close_at_end: bool = True

    def __post_init__(self) -> None:
        if self.initial_capital <= 0:
            raise ValueError("Initial capital must be positive.")
        if not 0 < self.max_gross_exposure <= 1:
            raise ValueError("Maximum gross exposure must be in (0, 1].")
        if self.transaction_cost_bps < 0 or self.slippage_bps < 0:
            raise ValueError("Cost assumptions must not be negative.")


@dataclass(frozen=True, slots=True)
class Trade:
    symbol: str
    entry_time: datetime
    exit_time: datetime
    quantity: int
    entry_price: float
    exit_price: float
    gross_pnl: float
    costs: float
    net_pnl: float
    exit_reason: str


@dataclass(frozen=True, slots=True)
class EquityPoint:
    timestamp: datetime
    equity: float
    gross_exposure: float


@dataclass(frozen=True, slots=True)
class BacktestResult:
    experiment_id: str
    strategy_name: str
    strategy_version: str
    dataset_id: str
    config: BacktestConfig
    trades: tuple[Trade, ...]
    equity_curve: tuple[EquityPoint, ...]
    metrics: dict[str, float | int | None]

    def to_dict(self, max_equity_points: int | None = 2500) -> dict[str, Any]:
        points = list(self.equity_curve)
        if max_equity_points is not None and len(points) > max_equity_points:
            step = (len(points) - 1) / (max_equity_points - 1)
            indices = {round(i * step) for i in range(max_equity_points)}
            points = [p for i, p in enumerate(points) if i in indices]

        return {
            "experiment_id": self.experiment_id,
            "strategy_name": self.strategy_name,
            "strategy_version": self.strategy_version,
            "dataset_id": self.dataset_id,
            "config": asdict(self.config),
            "trades": [
                {
                    **asdict(trade),
                    "entry_time": trade.entry_time.isoformat(sep=" "),
                    "exit_time": trade.exit_time.isoformat(sep=" "),
                }
                for trade in self.trades
            ],
            "equity_curve": [
                {
                    **asdict(point),
                    "timestamp": point.timestamp.isoformat(sep=" "),
                }
                for point in points
            ],
            "metrics": self.metrics,
        }


@dataclass(slots=True)
class _Position:
    quantity: int
    entry_time: datetime
    entry_price: float
    entry_cost: float


@dataclass(frozen=True, slots=True)
class _PendingOrder:
    symbol: str
    target_long: bool
    decision_time: datetime


def _side_adjusted_price(open_price: float, is_buy: bool, slippage_bps: float) -> float:
    adjustment = slippage_bps / 10_000
    return open_price * (1 + adjustment if is_buy else 1 - adjustment)


def _commission(notional: float, cost_bps: float) -> float:
    return notional * cost_bps / 10_000


def _mark_equity(
    cash: float, positions: dict[str, _Position], latest_prices: dict[str, float]
) -> tuple[float, float]:
    gross_exposure = sum(
        position.quantity * latest_prices[symbol]
        for symbol, position in positions.items()
        if symbol in latest_prices
    )
    return cash + gross_exposure, gross_exposure


def _daily_equity(points: tuple[EquityPoint, ...]) -> list[float]:
    last_by_date: dict[object, float] = {}
    for point in points:
        last_by_date[point.timestamp.date()] = point.equity
    return [last_by_date[key] for key in sorted(last_by_date)]


def _metrics(
    initial_capital: float, points: tuple[EquityPoint, ...], trades: tuple[Trade, ...], costs: float
) -> dict[str, float | int | None]:
    if not points:
        return {
            "cumulative_return": 0.0,
            "annualized_return": None,
            "annualized_volatility": None,
            "sharpe_ratio": None,
            "sortino_ratio": None,
            "maximum_drawdown": 0.0,
            "trade_count": 0,
            "transaction_costs": costs,
        }
    final_equity = points[-1].equity
    cumulative_return = final_equity / initial_capital - 1
    daily_equity = _daily_equity(points)
    daily_returns = [
        daily_equity[index] / daily_equity[index - 1] - 1
        for index in range(1, len(daily_equity))
        if daily_equity[index - 1] > 0
    ]
    annualized_return: float | None = None
    annualized_volatility: float | None = None
    sharpe: float | None = None
    sortino: float | None = None
    if len(daily_equity) >= 2 and daily_equity[0] > 0:
        years = max((points[-1].timestamp - points[0].timestamp).days / 365.25, 1 / 365.25)
        annualized_return = (final_equity / initial_capital) ** (1 / years) - 1
    if len(daily_returns) >= 2:
        mean_return = fmean(daily_returns)
        variance = fmean([(value - mean_return) ** 2 for value in daily_returns])
        annualized_volatility = math.sqrt(variance) * math.sqrt(252)
        if annualized_volatility > 0:
            sharpe = mean_return / math.sqrt(variance) * math.sqrt(252)
        downside = [min(value, 0.0) for value in daily_returns]
        downside_deviation = math.sqrt(fmean([value**2 for value in downside]))
        if downside_deviation > 0:
            sortino = mean_return / downside_deviation * math.sqrt(252)
    peak = initial_capital
    maximum_drawdown = 0.0
    for point in points:
        peak = max(peak, point.equity)
        maximum_drawdown = min(maximum_drawdown, point.equity / peak - 1)
    wins = [trade.net_pnl for trade in trades if trade.net_pnl > 0]
    losses = [trade.net_pnl for trade in trades if trade.net_pnl < 0]
    gross_profit = sum(wins)
    gross_loss = abs(sum(losses))
    holding_minutes = [
        (trade.exit_time - trade.entry_time).total_seconds() / 60 for trade in trades
    ]
    return {
        "cumulative_return": cumulative_return,
        "annualized_return": annualized_return,
        "annualized_volatility": annualized_volatility,
        "sharpe_ratio": sharpe,
        "sortino_ratio": sortino,
        "maximum_drawdown": maximum_drawdown,
        "trade_count": len(trades),
        "win_rate": len(wins) / len(trades) if trades else None,
        "average_win": fmean(wins) if wins else None,
        "average_loss": fmean(losses) if losses else None,
        "profit_factor": gross_profit / gross_loss if gross_loss else None,
        "average_holding_minutes": fmean(holding_minutes) if holding_minutes else None,
        "transaction_costs": costs,
        "final_equity": final_equity,
    }


def run_backtest(
    dataset: LoadedNseDataset,
    strategy: Strategy,
    config: BacktestConfig,
    progress_callback: Callable[[float, datetime, float, int, Trade | None], None] | None = None,
) -> BacktestResult:
    """Run completed-bar decisions with fills only at the next available bar open.

    The engine allocates the allowed gross exposure equally across the configured
    symbols. It does not use current-bar close to fill an order.
    """

    all_events = sorted(
        (bar for bars in dataset.bars_by_symbol.values() for bar in bars),
        key=lambda bar: (bar.timestamp, bar.symbol),
    )
    if not all_events:
        raise ValueError("Dataset has no valid bars for the selected universe.")
    allocation_per_symbol = (
        config.initial_capital * config.max_gross_exposure / len(dataset.symbols)
    )
    cash = config.initial_capital
    positions: dict[str, _Position] = {}
    pending: dict[str, _PendingOrder] = {}
    latest_prices: dict[str, float] = {}
    trades: list[Trade] = []
    equity_curve: list[EquityPoint] = []
    total_costs = 0.0

    total_events = len(all_events)
    report_interval = max(200, total_events // 60)
    current_ts: datetime | None = None
    last_trade: Trade | None = None

    for idx, bar in enumerate(all_events):
        latest_prices[bar.symbol] = bar.close
        order = pending.pop(bar.symbol, None)
        if order is not None:
            if order.target_long and bar.symbol not in positions:
                fill_price = _side_adjusted_price(bar.open, True, config.slippage_bps)
                quantity = int(allocation_per_symbol // fill_price)
                notional = quantity * fill_price
                cost = _commission(notional, config.transaction_cost_bps)
                if quantity > 0 and cash >= notional + cost:
                    cash -= notional + cost
                    total_costs += cost
                    positions[bar.symbol] = _Position(
                        quantity=quantity,
                        entry_time=bar.timestamp,
                        entry_price=fill_price,
                        entry_cost=cost,
                    )
            elif not order.target_long and bar.symbol in positions:
                position = positions.pop(bar.symbol)
                fill_price = _side_adjusted_price(bar.open, False, config.slippage_bps)
                notional = position.quantity * fill_price
                exit_cost = _commission(notional, config.transaction_cost_bps)
                cash += notional - exit_cost
                total_costs += exit_cost
                gross_pnl = position.quantity * (fill_price - position.entry_price)
                costs = position.entry_cost + exit_cost
                new_trade = Trade(
                    symbol=bar.symbol,
                    entry_time=position.entry_time,
                    exit_time=bar.timestamp,
                    quantity=position.quantity,
                    entry_price=position.entry_price,
                    exit_price=fill_price,
                    gross_pnl=gross_pnl,
                    costs=costs,
                    net_pnl=gross_pnl - costs,
                    exit_reason="strategy_signal",
                )
                trades.append(new_trade)
                last_trade = new_trade

        desired = strategy.target(bar.symbol, bar)
        if desired is not None:
            is_long = bar.symbol in positions
            existing_pending = pending.get(bar.symbol)
            if desired != is_long and (
                existing_pending is None or existing_pending.target_long != desired
            ):
                pending[bar.symbol] = _PendingOrder(
                    symbol=bar.symbol, target_long=desired, decision_time=bar.timestamp
                )

        if current_ts is None or bar.timestamp != current_ts:
            equity, exposure = _mark_equity(cash, positions, latest_prices)
            equity_curve.append(
                EquityPoint(timestamp=bar.timestamp, equity=equity, gross_exposure=exposure)
            )
            current_ts = bar.timestamp

            if progress_callback is not None and (idx % report_interval == 0):
                pct = round((idx + 1) / total_events * 100, 1)
                progress_callback(pct, bar.timestamp, equity, len(trades), last_trade)

    if config.force_close_at_end:
        for symbol, position in list(positions.items()):
            bar = dataset.bars_by_symbol[symbol][-1]
            # Explicit end-of-data liquidation assumption; it is never treated as a next-bar fill.
            fill_price = _side_adjusted_price(bar.close, False, config.slippage_bps)
            notional = position.quantity * fill_price
            exit_cost = _commission(notional, config.transaction_cost_bps)
            cash += notional - exit_cost
            total_costs += exit_cost
            gross_pnl = position.quantity * (fill_price - position.entry_price)
            costs = position.entry_cost + exit_cost
            new_trade = Trade(
                symbol=symbol,
                entry_time=position.entry_time,
                exit_time=bar.timestamp,
                quantity=position.quantity,
                entry_price=position.entry_price,
                exit_price=fill_price,
                gross_pnl=gross_pnl,
                costs=costs,
                net_pnl=gross_pnl - costs,
                exit_reason="end_of_data_close_assumption",
            )
            trades.append(new_trade)
            last_trade = new_trade
            del positions[symbol]
        final_time = max(bar.timestamp for bar in all_events)
        equity_curve.append(EquityPoint(timestamp=final_time, equity=cash, gross_exposure=0.0))

    if progress_callback is not None:
        final_eq = equity_curve[-1].equity if equity_curve else cash
        progress_callback(100.0, all_events[-1].timestamp, final_eq, len(trades), last_trade)

    immutable_points = tuple(equity_curve)
    immutable_trades = tuple(trades)
    return BacktestResult(
        experiment_id=str(uuid4()),
        strategy_name=strategy.name,
        strategy_version=strategy.version,
        dataset_id=dataset.dataset_id,
        config=config,
        trades=immutable_trades,
        equity_curve=immutable_points,
        metrics=_metrics(config.initial_capital, immutable_points, immutable_trades, total_costs),
    )