#!/usr/bin/env python3
"""
Morning Momentum Day-Trade Backtest
====================================
Strategy:
  - At ENTRY_TIME (e.g. 10:00 AM ET), find top N gainers from market open
  - Filter: S&P 500 (liquid, large-cap)
  - Buy top N at ENTRY_TIME price (equal weight)
  - Sell at EXIT_TIME (e.g. 3:30 PM ET) OR on stop-loss (e.g. -2%)
  - Daily P&L tracking, no overnight holds

Uses Oracle API (Alpaca intraday) for historical 5-min bars.
"""

from __future__ import annotations

import asyncio
import sys
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timedelta, date

import httpx

# ── Configuration ──────────────────────────────────────────────
ORACLE_URL = "http://localhost:18001"

# Strategy parameters
ENTRY_MINUTES_AFTER_OPEN = 30     # minutes after 9:30 AM ET
EXIT_MINUTES_BEFORE_CLOSE = 30    # minutes before 4:00 PM ET
TOP_N = 3                         # number of stocks to buy
STOP_LOSS_PCT = -0.02             # -2% stop loss (None to disable)
MIN_MORNING_GAIN_PCT = 0.01       # minimum 1% gain to qualify as "gainer"
INITIAL_CAPITAL = 10_000.0

# Backtest period
LOOKBACK_TRADING_DAYS = 40        # ~2 months

# Market hours in UTC (ET + 4 in EDT / ET + 5 in EST)
# April = EDT, so 9:30 AM ET = 13:30 UTC, 4:00 PM ET = 20:00 UTC
MARKET_OPEN_UTC_HOUR = 13
MARKET_OPEN_UTC_MIN = 30
MARKET_CLOSE_UTC_HOUR = 20
MARKET_CLOSE_UTC_MIN = 0

INTERVAL = "5min"
INTERVAL_MINUTES = 5


@dataclass
class Trade:
    date: str
    ticker: str
    entry_price: float
    exit_price: float
    entry_time: str
    exit_time: str
    shares: float
    pnl: float
    pnl_pct: float
    exit_reason: str  # "close" or "stop_loss"
    morning_gain_pct: float  # gain at entry from open


@dataclass
class DayResult:
    date: str
    trades: list[Trade] = field(default_factory=list)
    daily_pnl: float = 0.0
    daily_pnl_pct: float = 0.0
    candidates_found: int = 0


async def fetch_json(client: httpx.AsyncClient, url: str, params: dict = None) -> dict:
    resp = await client.get(url, params=params, timeout=30.0)
    resp.raise_for_status()
    return resp.json()


async def get_sp500(client: httpx.AsyncClient) -> list[str]:
    """Get S&P 500 ticker list."""
    data = await fetch_json(client, f"{ORACLE_URL}/api/v1/stocks/index/sp500")
    return [s["symbol"] for s in data.get("constituents", [])]


async def get_trading_days(client: httpx.AsyncClient, ticker: str = "SPY",
                           days: int = 60) -> list[str]:
    """Get recent trading days from SPY daily bars."""
    end = date.today()
    start = end - timedelta(days=days + 30)  # extra buffer for weekends/holidays
    data = await fetch_json(client, f"{ORACLE_URL}/api/v1/price/data/{ticker}", {
        "start_date": start.isoformat(),
        "end_date": end.isoformat(),
    })
    bars = data.get("data", [])
    dates = sorted(b["date"][:10] for b in bars)
    return dates[-days:]  # last N trading days


async def get_intraday(client: httpx.AsyncClient, ticker: str, dt: str,
                       semaphore: asyncio.Semaphore) -> tuple[str, list[dict]]:
    """Fetch 5-min intraday bars for ticker on date, with concurrency limit."""
    async with semaphore:
        try:
            data = await fetch_json(client, f"{ORACLE_URL}/api/v1/alpaca/intraday/{ticker}", {
                "interval": INTERVAL,
                "start_date": dt,
                "end_date": dt,
                "limit": 500,
            })
            return ticker, data.get("candles", [])
        except Exception as e:
            return ticker, []


def parse_utc_ts(ts_str: str) -> datetime:
    """Parse timestamp string to datetime (UTC)."""
    # Handle both 2026-04-09T08:00:00Z and 2026-04-09T08:00:00+00:00
    ts_str = ts_str.replace("Z", "+00:00")
    if "+" not in ts_str and "-" not in ts_str[10:]:
        ts_str += "+00:00"
    dt = datetime.fromisoformat(ts_str)
    return dt.replace(tzinfo=None)  # work in UTC naively


def filter_market_hours(candles: list[dict]) -> list[dict]:
    """Keep only regular trading hours candles (9:30-16:00 ET = 13:30-20:00 UTC)."""
    filtered = []
    for c in candles:
        ts = parse_utc_ts(c["timestamp"])
        market_open = ts.replace(hour=MARKET_OPEN_UTC_HOUR, minute=MARKET_OPEN_UTC_MIN, second=0)
        market_close = ts.replace(hour=MARKET_CLOSE_UTC_HOUR, minute=MARKET_CLOSE_UTC_MIN, second=0)
        if market_open <= ts < market_close:
            filtered.append(c)
    return filtered


def get_candle_at_time(candles: list[dict], target_hour_utc: int, target_min_utc: int,
                       tolerance_minutes: int = 10) -> dict | None:
    """Find candle closest to target time within tolerance."""
    best = None
    best_diff = float("inf")
    for c in candles:
        ts = parse_utc_ts(c["timestamp"])
        target = ts.replace(hour=target_hour_utc, minute=target_min_utc)
        diff = abs((ts - target).total_seconds())
        if diff < best_diff and diff <= tolerance_minutes * 60:
            best = c
            best_diff = diff
    return best


def simulate_day(candles_by_ticker: dict[str, list[dict]], day: str,
                 capital_per_trade: float) -> DayResult:
    """Simulate one day of trading."""
    result = DayResult(date=day)

    # Entry time in UTC: 9:30 + ENTRY_MINUTES_AFTER_OPEN = e.g. 10:00 AM ET = 14:00 UTC
    entry_h = MARKET_OPEN_UTC_HOUR
    entry_m = MARKET_OPEN_UTC_MIN + ENTRY_MINUTES_AFTER_OPEN
    entry_h += entry_m // 60
    entry_m = entry_m % 60

    # Exit time in UTC: 16:00 - EXIT_MINUTES_BEFORE_CLOSE = e.g. 15:30 ET = 19:30 UTC
    exit_total_min = (MARKET_CLOSE_UTC_HOUR * 60 + MARKET_CLOSE_UTC_MIN) - EXIT_MINUTES_BEFORE_CLOSE
    exit_h = exit_total_min // 60
    exit_m = exit_total_min % 60

    # Step 1: Calculate morning gain for each ticker
    morning_gains = {}
    for ticker, candles in candles_by_ticker.items():
        mkt_candles = filter_market_hours(candles)
        if len(mkt_candles) < 5:
            continue

        # Open price = first candle's open
        open_price = mkt_candles[0]["open"]
        if open_price <= 0:
            continue

        # Entry candle
        entry_candle = get_candle_at_time(mkt_candles, entry_h, entry_m)
        if not entry_candle:
            continue

        entry_price = entry_candle["close"]
        gain_pct = (entry_price - open_price) / open_price

        if gain_pct >= MIN_MORNING_GAIN_PCT:
            morning_gains[ticker] = {
                "gain_pct": gain_pct,
                "entry_price": entry_price,
                "entry_time": entry_candle["timestamp"],
                "candles": mkt_candles,
            }

    result.candidates_found = len(morning_gains)

    if not morning_gains:
        return result

    # Step 2: Pick top N
    top_tickers = sorted(morning_gains.keys(), key=lambda t: morning_gains[t]["gain_pct"],
                         reverse=True)[:TOP_N]

    # Step 3: Simulate each trade
    for ticker in top_tickers:
        info = morning_gains[ticker]
        entry_price = info["entry_price"]
        shares = capital_per_trade / entry_price
        candles = info["candles"]

        # Find exit: iterate candles after entry to check stop loss
        exit_price = entry_price
        exit_time = info["entry_time"]
        exit_reason = "close"

        entry_ts = parse_utc_ts(info["entry_time"])

        for c in candles:
            ts = parse_utc_ts(c["timestamp"])
            if ts <= entry_ts:
                continue

            # Check stop loss on low
            if STOP_LOSS_PCT is not None:
                low_pnl = (c["low"] - entry_price) / entry_price
                if low_pnl <= STOP_LOSS_PCT:
                    # Stop loss triggered - assume exit at stop price
                    exit_price = entry_price * (1 + STOP_LOSS_PCT)
                    exit_time = c["timestamp"]
                    exit_reason = "stop_loss"
                    break

            # Check if at/past exit time
            exit_target = ts.replace(hour=exit_h, minute=exit_m)
            if ts >= exit_target:
                exit_price = c["close"]
                exit_time = c["timestamp"]
                exit_reason = "close"
                break
            else:
                # Update running exit (last seen candle)
                exit_price = c["close"]
                exit_time = c["timestamp"]

        pnl = (exit_price - entry_price) * shares
        pnl_pct = (exit_price - entry_price) / entry_price

        trade = Trade(
            date=day,
            ticker=ticker,
            entry_price=round(entry_price, 2),
            exit_price=round(exit_price, 2),
            entry_time=info["entry_time"],
            exit_time=exit_time,
            shares=round(shares, 4),
            pnl=round(pnl, 2),
            pnl_pct=round(pnl_pct, 4),
            exit_reason=exit_reason,
            morning_gain_pct=round(info["gain_pct"], 4),
        )
        result.trades.append(trade)
        result.daily_pnl += trade.pnl

    if result.trades:
        total_invested = capital_per_trade * len(result.trades)
        result.daily_pnl_pct = result.daily_pnl / total_invested

    return result


async def run_backtest():
    """Main backtest loop."""
    print("=" * 70)
    print("  Morning Momentum Day-Trade Backtest")
    print("=" * 70)
    print(f"\n📋 Parameters:")
    print(f"   Entry:     {ENTRY_MINUTES_AFTER_OPEN} min after open (10:00 AM ET)")
    print(f"   Exit:      {EXIT_MINUTES_BEFORE_CLOSE} min before close (3:30 PM ET)")
    print(f"   Top N:     {TOP_N} stocks per day")
    print(f"   Stop Loss: {STOP_LOSS_PCT*100 if STOP_LOSS_PCT else 'None'}%")
    print(f"   Min Gain:  {MIN_MORNING_GAIN_PCT*100}% to qualify")
    print(f"   Capital:   ${INITIAL_CAPITAL:,.0f}")
    print(f"   Period:    last {LOOKBACK_TRADING_DAYS} trading days")
    print()

    async with httpx.AsyncClient() as client:
        # Step 1: Get universe
        print("⏳ Fetching S&P 500 constituents...")
        sp500 = await get_sp500(client)
        print(f"   Got {len(sp500)} tickers")

        # Step 2: Get trading days
        print("⏳ Fetching trading days...")
        trading_days = await get_trading_days(client, days=LOOKBACK_TRADING_DAYS + 10)
        trading_days = trading_days[-LOOKBACK_TRADING_DAYS:]
        print(f"   Period: {trading_days[0]} to {trading_days[-1]} ({len(trading_days)} days)")

        # Step 3: Pre-screen using daily bars to find "big mover" days
        # Instead of fetching intraday for ALL 500 stocks every day,
        # first get daily bars to identify which stocks had big moves
        print("\n⏳ Phase 1: Fetching daily bars for pre-screening...")
        daily_by_ticker: dict[str, list[dict]] = {}
        batch_size = 50
        semaphore = asyncio.Semaphore(20)

        async def fetch_daily(ticker):
            async with semaphore:
                try:
                    data = await fetch_json(client, f"{ORACLE_URL}/api/v1/price/data/{ticker}", {
                        "start_date": trading_days[0],
                        "end_date": trading_days[-1],
                    })
                    return ticker, data.get("data", [])
                except Exception:
                    return ticker, []

        for i in range(0, len(sp500), batch_size):
            batch = sp500[i:i + batch_size]
            tasks = [fetch_daily(t) for t in batch]
            results = await asyncio.gather(*tasks)
            for ticker, bars in results:
                if bars:
                    daily_by_ticker[ticker] = bars
            pct = min(100, (i + batch_size) / len(sp500) * 100)
            sys.stdout.write(f"\r   Progress: {pct:.0f}% ({len(daily_by_ticker)} tickers with data)")
            sys.stdout.flush()
        print()

        # Step 4: For each day, identify potential big movers using daily data
        # Heuristic: stocks where (close - open)/open > 1% OR (high - open)/open > 2%
        print("\n⏳ Phase 2: Identifying daily big movers...")
        day_candidates: dict[str, list[str]] = defaultdict(list)  # date -> tickers

        for ticker, bars in daily_by_ticker.items():
            bar_by_date = {b["date"][:10]: b for b in bars}
            for day in trading_days:
                if day not in bar_by_date:
                    continue
                b = bar_by_date[day]
                open_p = b["open"]
                if open_p <= 0:
                    continue
                # Use (high - open) / open as proxy for morning momentum potential
                intraday_range = (b["high"] - open_p) / open_p
                if intraday_range >= 0.015:  # at least 1.5% above open at some point
                    day_candidates[day].append(ticker)

        total_candidates = sum(len(v) for v in day_candidates.values())
        print(f"   Found {total_candidates} ticker-day candidates across {len(day_candidates)} days")
        # Cap to top 20 per day to limit API calls
        for day in day_candidates:
            tickers = day_candidates[day]
            # Sort by daily range (high-open)/open descending
            bar_map = {}
            for t in tickers:
                for b in daily_by_ticker.get(t, []):
                    if b["date"][:10] == day:
                        bar_map[t] = (b["high"] - b["open"]) / b["open"] if b["open"] > 0 else 0
            day_candidates[day] = sorted(tickers, key=lambda t: bar_map.get(t, 0), reverse=True)[:20]

        capped_total = sum(len(v) for v in day_candidates.values())
        print(f"   Capped to top 20 per day: {capped_total} total intraday fetches needed")

        # Step 5: Fetch intraday data for candidates
        print("\n⏳ Phase 3: Fetching intraday data for candidates...")
        intraday_sem = asyncio.Semaphore(10)  # lower concurrency for intraday
        day_intraday: dict[str, dict[str, list[dict]]] = defaultdict(dict)

        fetch_count = 0
        for day_idx, day in enumerate(trading_days):
            if day not in day_candidates or not day_candidates[day]:
                continue

            tasks = [get_intraday(client, t, day, intraday_sem) for t in day_candidates[day]]
            results = await asyncio.gather(*tasks)
            for ticker, candles in results:
                if candles:
                    day_intraday[day][ticker] = candles
            fetch_count += len(tasks)
            sys.stdout.write(f"\r   Progress: day {day_idx+1}/{len(trading_days)} | {fetch_count} API calls")
            sys.stdout.flush()
        print()

        # Step 6: Simulate trades
        print("\n⏳ Phase 4: Simulating trades...")
        capital_per_trade = INITIAL_CAPITAL / TOP_N
        all_results: list[DayResult] = []
        all_trades: list[Trade] = []

        for day in trading_days:
            candles = day_intraday.get(day, {})
            if not candles:
                continue
            day_result = simulate_day(candles, day, capital_per_trade)
            all_results.append(day_result)
            all_trades.extend(day_result.trades)

    # ── Results ──────────────────────────────────────────────────
    print("\n" + "=" * 70)
    print("  RESULTS")
    print("=" * 70)

    trading_days_with_trades = [r for r in all_results if r.trades]
    total_trades = len(all_trades)

    if total_trades == 0:
        print("\n❌ No trades executed. Try relaxing MIN_MORNING_GAIN_PCT.")
        return

    # Basic stats
    wins = [t for t in all_trades if t.pnl > 0]
    losses = [t for t in all_trades if t.pnl <= 0]
    stop_losses = [t for t in all_trades if t.exit_reason == "stop_loss"]

    total_pnl = sum(t.pnl for t in all_trades)
    avg_pnl_per_trade = total_pnl / total_trades
    win_rate = len(wins) / total_trades
    avg_win = sum(t.pnl_pct for t in wins) / len(wins) if wins else 0
    avg_loss = sum(t.pnl_pct for t in losses) / len(losses) if losses else 0

    # Daily stats
    daily_returns = [r.daily_pnl_pct for r in all_results if r.trades]
    avg_daily_return = sum(daily_returns) / len(daily_returns) if daily_returns else 0

    # Cumulative equity
    equity = INITIAL_CAPITAL
    equity_curve = [equity]
    for r in all_results:
        equity += r.daily_pnl
        equity_curve.append(equity)

    max_equity = INITIAL_CAPITAL
    max_dd = 0
    for eq in equity_curve:
        max_equity = max(max_equity, eq)
        dd = (eq - max_equity) / max_equity
        max_dd = min(max_dd, dd)

    # Annualized return (rough)
    total_return_pct = (equity - INITIAL_CAPITAL) / INITIAL_CAPITAL
    n_days = len(trading_days)
    ann_factor = 252 / n_days if n_days > 0 else 1
    ann_return = total_return_pct * ann_factor

    print(f"\n📊 Summary ({trading_days[0]} → {trading_days[-1]})")
    print(f"   Trading days:       {n_days}")
    print(f"   Days with trades:   {len(trading_days_with_trades)}")
    print(f"   Total trades:       {total_trades}")
    print(f"   Stop-loss exits:    {len(stop_losses)}")
    print(f"\n💰 P&L:")
    print(f"   Total P&L:          ${total_pnl:+,.2f}  ({total_return_pct:+.2%})")
    print(f"   Avg P&L/trade:      ${avg_pnl_per_trade:+,.2f}")
    print(f"   Annualized return:  {ann_return:+.1%}")
    print(f"\n📈 Win/Loss:")
    print(f"   Win rate:           {win_rate:.1%}  ({len(wins)}W / {len(losses)}L)")
    print(f"   Avg winner:         {avg_win:+.2%}")
    print(f"   Avg loser:          {avg_loss:+.2%}")
    profit_factor = abs(sum(t.pnl for t in wins) / sum(t.pnl for t in losses)) if losses and sum(t.pnl for t in losses) != 0 else float("inf")
    print(f"   Profit factor:      {profit_factor:.2f}")
    print(f"\n📉 Risk:")
    print(f"   Max drawdown:       {max_dd:.2%}")
    print(f"   Avg daily return:   {avg_daily_return:+.4%}")
    print(f"   Final equity:       ${equity:,.2f}")

    # Show top 10 best and worst trades
    sorted_trades = sorted(all_trades, key=lambda t: t.pnl_pct, reverse=True)
    print(f"\n🏆 Top 5 Winners:")
    for t in sorted_trades[:5]:
        print(f"   {t.date} {t.ticker:6} entry={t.entry_price:8.2f} exit={t.exit_price:8.2f}"
              f"  {t.pnl_pct:+.2%}  (${t.pnl:+.2f})  morn_gain={t.morning_gain_pct:+.2%}")

    print(f"\n💀 Top 5 Losers:")
    for t in sorted_trades[-5:]:
        print(f"   {t.date} {t.ticker:6} entry={t.entry_price:8.2f} exit={t.exit_price:8.2f}"
              f"  {t.pnl_pct:+.2%}  (${t.pnl:+.2f})  morn_gain={t.morning_gain_pct:+.2%}  [{t.exit_reason}]")

    # Daily breakdown
    print(f"\n📅 Daily P&L (days with trades):")
    print(f"   {'Date':12} {'#Trades':>8} {'P&L':>10} {'Return':>8} {'Tickers'}")
    print(f"   {'-'*12} {'-'*8} {'-'*10} {'-'*8} {'-'*30}")
    for r in all_results:
        if not r.trades:
            continue
        tickers = ", ".join(f"{t.ticker}({t.pnl_pct:+.1%})" for t in r.trades)
        print(f"   {r.date:12} {len(r.trades):>8} ${r.daily_pnl:>+9.2f} {r.daily_pnl_pct:>+7.2%} {tickers}")


    # Also test variations
    print(f"\n{'='*70}")
    print("  SENSITIVITY: Entry at 1 hour after open")
    print(f"{'='*70}")


if __name__ == "__main__":
    asyncio.run(run_backtest())
