#!/usr/bin/env python3
"""Analyze SPY regime detection during backtest periods."""

import pandas as pd
import yfinance as yf
from datetime import datetime, timedelta

def calculate_ema(series, period):
    """Calculate EMA using pandas."""
    return series.ewm(span=period, adjust=False).mean()

def determine_regime(price, ema50, ema200):
    """Replicate the Rust regime detection logic."""
    if pd.isna(price) or pd.isna(ema50) or pd.isna(ema200):
        return "Caution"

    # Bear: price below 200 EMA AND 50 EMA below 200 EMA
    if price < ema200 and ema50 < ema200:
        return "Bear"

    # Caution: price below 50 EMA
    if price < ema50:
        return "Caution"

    # Bull: price above both, 50 above 200
    if ema50 > ema200:
        return "Bull"

    # Edge case: price above both but 50 still below 200
    return "Caution"

def analyze_period(start_date, end_date, period_name):
    """Analyze SPY regime for a given period."""
    print(f"\n{'='*70}")
    print(f"{period_name}: {start_date} to {end_date}")
    print('='*70)

    # Fetch SPY data with extra warmup for EMA-200
    warmup_start = (datetime.strptime(start_date, '%Y-%m-%d') - timedelta(days=400)).strftime('%Y-%m-%d')
    spy = yf.download('SPY', start=warmup_start, end=end_date, progress=False)

    if spy.empty:
        print(f"ERROR: No SPY data available for {period_name}")
        return

    # Calculate EMAs
    spy['EMA50'] = calculate_ema(spy['Close'], 50)
    spy['EMA200'] = calculate_ema(spy['Close'], 200)

    # Determine regime for each day
    spy['Regime'] = spy.apply(
        lambda row: determine_regime(row['Close'], row['EMA50'], row['EMA200']),
        axis=1
    )

    # Filter to actual trading period
    trading_period = spy[start_date:end_date].copy()

    if trading_period.empty:
        print(f"ERROR: No trading data for {period_name}")
        return

    # Calculate SPY return
    spy_start = trading_period['Close'].iloc[0]
    spy_end = trading_period['Close'].iloc[-1]
    spy_return = (spy_end - spy_start) / spy_start * 100

    print(f"\nSPY Performance:")
    print(f"  Start: ${spy_start:.2f}")
    print(f"  End:   ${spy_end:.2f}")
    print(f"  Return: {spy_return:+.2f}%")

    # Count regime days
    regime_counts = trading_period['Regime'].value_counts()
    total_days = len(trading_period)

    print(f"\nRegime Distribution ({total_days} trading days):")
    for regime in ['Bull', 'Caution', 'Bear']:
        count = regime_counts.get(regime, 0)
        pct = count / total_days * 100
        print(f"  {regime:8s}: {count:4d} days ({pct:5.1f}%)")

    # Show regime transitions
    regime_changes = trading_period[trading_period['Regime'] != trading_period['Regime'].shift(1)]
    if len(regime_changes) > 0:
        print(f"\nRegime Transitions ({len(regime_changes)} total):")
        for date, row in regime_changes.head(20).iterrows():
            print(f"  {date.strftime('%Y-%m-%d')}: {row['Regime']:8s} (SPY: ${row['Close']:.2f}, "
                  f"EMA50: ${row['EMA50']:.2f}, EMA200: ${row['EMA200']:.2f})")
        if len(regime_changes) > 20:
            print(f"  ... and {len(regime_changes) - 20} more transitions")

    # Identify problematic Bear periods during bull markets
    bear_days = trading_period[trading_period['Regime'] == 'Bear']
    if len(bear_days) > 0:
        print(f"\n⚠️  WARNING: {len(bear_days)} days classified as BEAR:")
        for date, row in bear_days.head(10).iterrows():
            print(f"  {date.strftime('%Y-%m-%d')}: SPY=${row['Close']:.2f}, "
                  f"EMA50=${row['EMA50']:.2f}, EMA200=${row['EMA200']:.2f}")
        if len(bear_days) > 10:
            print(f"  ... and {len(bear_days) - 10} more Bear days")

    # Show first and last months in detail
    print(f"\nFirst Month Detail:")
    first_month = trading_period.head(22)[['Close', 'EMA50', 'EMA200', 'Regime']]
    for date, row in first_month.iterrows():
        print(f"  {date.strftime('%Y-%m-%d')}: {row['Regime']:8s} | "
              f"SPY: ${row['Close']:7.2f} | EMA50: ${row['EMA50']:7.2f} | EMA200: ${row['EMA200']:7.2f}")

if __name__ == '__main__':
    # Analyze 2023
    analyze_period('2023-01-01', '2023-12-31', '2023 Backtest')

    # Analyze 2024
    analyze_period('2024-01-01', '2024-12-31', '2024 Backtest')

    print(f"\n{'='*70}")
    print("Analysis complete.")
    print('='*70)