#!/usr/bin/env python3
"""
RSI 50-line momentum template.

Bring your own OHLC CSV. Expected columns:
datetime, open, high, low, close

Example:
python3 run_backtest.py --csv your_ohlc.csv --cost-bps 2 --max-hold 20
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd


def rsi(close: pd.Series, length: int) -> pd.Series:
    change = close.diff()
    gain = change.clip(lower=0)
    loss = -change.clip(upper=0)
    avg_gain = gain.ewm(alpha=1 / length, adjust=False, min_periods=length).mean()
    avg_loss = loss.ewm(alpha=1 / length, adjust=False, min_periods=length).mean()
    rs = avg_gain / avg_loss.replace(0, np.nan)
    return 100 - (100 / (1 + rs))


def load_ohlc(path: Path) -> pd.DataFrame:
    df = pd.read_csv(path)
    required = {"open", "high", "low", "close"}
    missing = required - set(df.columns.str.lower())
    if missing:
        raise SystemExit(f"CSV is missing required columns: {', '.join(sorted(missing))}")

    df.columns = [column.lower() for column in df.columns]
    if "datetime" in df.columns:
        df["datetime"] = pd.to_datetime(df["datetime"])
        df = df.sort_values("datetime")
    else:
        df["datetime"] = pd.RangeIndex(len(df))

    for column in ["open", "high", "low", "close"]:
        df[column] = pd.to_numeric(df[column], errors="coerce")

    return df.dropna(subset=["open", "high", "low", "close"]).reset_index(drop=True)


def run_backtest(df: pd.DataFrame, args: argparse.Namespace) -> tuple[pd.DataFrame, dict[str, float]]:
    work = df.copy()
    work["rsi"] = rsi(work["close"], args.rsi_length)
    work["rsi_ma"] = work["rsi"].rolling(args.rsi_ma_length).mean()
    work["regime"] = np.select(
        [
            (work["rsi"] > 50 + args.neutral_zone) & (work["rsi_ma"] > 50 + args.neutral_zone),
            (work["rsi"] < 50 - args.neutral_zone) & (work["rsi_ma"] < 50 - args.neutral_zone),
        ],
        [1, -1],
        default=0,
    )

    trades = []
    position = 0
    entry_index = None
    entry_price = None

    for i in range(1, len(work)):
        signal = int(work.loc[i - 1, "regime"])
        open_price = float(work.loc[i, "open"])

        should_exit = False
        if position:
            held = i - entry_index
            if signal == -position or signal == 0 or held >= args.max_hold:
                should_exit = True

        if should_exit:
            gross = position * (open_price / entry_price - 1)
            net = gross - args.cost_bps / 10000
            trades.append(
                {
                    "entry_time": work.loc[entry_index, "datetime"],
                    "exit_time": work.loc[i, "datetime"],
                    "side": "long" if position == 1 else "short",
                    "entry": entry_price,
                    "exit": open_price,
                    "bars_held": i - entry_index,
                    "gross_return": gross,
                    "net_return": net,
                }
            )
            position = 0
            entry_index = None
            entry_price = None

        if position == 0 and signal:
            position = signal
            entry_index = i
            entry_price = open_price

    trades_df = pd.DataFrame(trades)
    if trades_df.empty:
        return trades_df, {
            "trades": 0,
            "win_rate": 0.0,
            "avg_trade": 0.0,
            "median_trade": 0.0,
            "total_net_return": 0.0,
            "profit_factor": 0.0,
        }

    wins = trades_df.loc[trades_df["net_return"] > 0, "net_return"]
    losses = trades_df.loc[trades_df["net_return"] < 0, "net_return"]
    profit_factor = wins.sum() / abs(losses.sum()) if abs(losses.sum()) > 0 else np.inf
    summary = {
        "trades": int(len(trades_df)),
        "win_rate": float((trades_df["net_return"] > 0).mean()),
        "avg_trade": float(trades_df["net_return"].mean()),
        "median_trade": float(trades_df["net_return"].median()),
        "total_net_return": float((1 + trades_df["net_return"]).prod() - 1),
        "profit_factor": float(profit_factor),
    }
    return trades_df, summary


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv", required=True, type=Path, help="Path to your OHLC CSV")
    parser.add_argument("--out-dir", type=Path, default=Path("rsi_50_results"))
    parser.add_argument("--rsi-length", type=int, default=14)
    parser.add_argument("--rsi-ma-length", type=int, default=9)
    parser.add_argument("--neutral-zone", type=float, default=0.0, help="Ignore RSI values within this distance of 50")
    parser.add_argument("--cost-bps", type=float, default=2.0, help="Round-trip cost in basis points")
    parser.add_argument("--max-hold", type=int, default=20)
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    df = load_ohlc(args.csv)
    trades, summary = run_backtest(df, args)
    trades.to_csv(args.out_dir / "trades.csv", index=False)
    pd.DataFrame([summary]).to_csv(args.out_dir / "summary.csv", index=False)
    print(pd.DataFrame([summary]).to_string(index=False))


if __name__ == "__main__":
    main()