backtest.py

import json
from pathlib import Path
from typing import Dict

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np
import pandas as pd


def _get_outputs_dir() -> Path:
    out = Path("outputs")
    out.mkdir(exist_ok=True)
    return out


def _compute_performance(daily_returns: pd.Series) -> Dict[str, float]:
    daily_returns = daily_returns.fillna(0)
    equity = (1 + daily_returns).cumprod()

    total_return = float(equity.iloc[-1] - 1)
    ann_factor = 252 / max(len(equity), 1)
    cagr = float(equity.iloc[-1] ** ann_factor - 1) if len(equity) > 0 else 0.0
    vol = float(daily_returns.std() * np.sqrt(252))
    sharpe = float((daily_returns.mean() / (daily_returns.std() + 1e-9)) * np.sqrt(252))
    max_dd = float(((equity / equity.cummax()) - 1).min())

    return {
        "total_return": total_return,
        "cagr": cagr,
        "annualized_volatility": vol,
        "sharpe_ratio": sharpe,
        "max_drawdown": max_dd,
    }


def run_backtest(signals_path: str, initial_capital: float = 1.0) -> Dict:
    signals_path = Path(signals_path)
    if not signals_path.exists():
        raise FileNotFoundError(f"Signals not found at {signals_path}")

    df = pd.read_csv(signals_path, parse_dates=["date"])
    if "future_return" not in df.columns:
        raise ValueError("Signals file must include 'future_return' for PnL calculation.")
    if "signal" not in df.columns:
        raise ValueError("Signals file must include 'signal' column.")

    if "strategy_return" not in df.columns:
        df["strategy_return"] = df["signal"] * df["future_return"]

    daily_strategy = df.groupby("date")["strategy_return"].mean().sort_index().fillna(0)
    daily_baseline = df.groupby("date")["future_return"].mean().sort_index().fillna(0)

    strategy_equity = (1 + daily_strategy).cumprod() * initial_capital
    buy_hold_equity = (1 + daily_baseline).cumprod() * initial_capital

    perf_strategy = _compute_performance(daily_strategy)
    perf_baseline = _compute_performance(daily_baseline)

    equity_df = pd.DataFrame(
        {
            "date": daily_strategy.index,
            "strategy_equity": strategy_equity.values,
            "buy_hold_equity": buy_hold_equity.values,
        }
    )

    outputs_dir = _get_outputs_dir()
    equity_path = outputs_dir / f"equity_curve_{signals_path.stem}.csv"
    equity_df.to_csv(equity_path, index=False)

    plt.figure(figsize=(7, 4))
    plt.plot(equity_df["date"], equity_df["strategy_equity"], label="Strategy")
    plt.plot(equity_df["date"], equity_df["buy_hold_equity"], label="Buy & Hold", linestyle="--")
    plt.title("Equity Curve")
    plt.xlabel("Date")
    plt.ylabel("Equity")
    plt.legend()
    plt.tight_layout()
    plot_path = outputs_dir / f"equity_curve_{signals_path.stem}.png"
    plt.savefig(plot_path)
    plt.close()

    metrics = {
        "strategy": perf_strategy,
        "buy_and_hold": perf_baseline,
        "equity_curve_path": str(equity_path),
        "equity_plot_path": str(plot_path),
    }
    metrics_path = outputs_dir / f"backtest_metrics_{signals_path.stem}.json"
    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    return metrics


if __name__ == "__main__":
    example_signals = next(Path("outputs").glob("signals_*.csv"), None)
    if example_signals:
        stats = run_backtest(str(example_signals))
        print(json.dumps(stats, indent=2))
    else:
        print("No signals found in outputs/. Generate predictions and signals first.")