"""
Test Real-time BBRS Strategy with Minute-level Data

This script validates that the incremental BBRS strategy can:
1. Accept minute-level data input (real-time simulation)
2. Internally aggregate to configured timeframes (15min, 1h, etc.)
3. Generate signals only when timeframe bars complete
4. Produce identical results to pre-aggregated data processing
"""

import pandas as pd
import numpy as np
import logging
from datetime import datetime, timedelta
import matplotlib.pyplot as plt

# Import incremental implementation
from cycles.IncStrategies.bbrs_incremental import BBRSIncrementalState

# Import storage utility
from cycles.utils.storage import Storage

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler("test_realtime_bbrs.log"),
        logging.StreamHandler()
    ]
)

def load_minute_data():
    """Load minute-level BTC data for real-time simulation."""
    storage = Storage(logging=logging)
    
    # Load data for testing period
    start_date = "2023-01-01"
    end_date = "2023-01-03"  # 2 days for testing
    
    data = storage.load_data("btcusd_1-min_data.csv", start_date, end_date)
    
    if data.empty:
        logging.error("No data loaded for testing period")
        return None
    
    logging.info(f"Loaded {len(data)} minute-level data points from {data.index[0]} to {data.index[-1]}")
    return data

def test_timeframe_aggregation():
    """Test different timeframe aggregations with minute-level data."""
    
    # Load minute data
    minute_data = load_minute_data()
    if minute_data is None:
        return
    
    # Test different timeframes
    timeframes = [15, 60]  # 15min and 1h
    
    for timeframe_minutes in timeframes:
        logging.info(f"\n{'='*60}")
        logging.info(f"Testing {timeframe_minutes}-minute timeframe")
        logging.info(f"{'='*60}")
        
        # Configuration for this timeframe
        config = {
            "timeframe_minutes": timeframe_minutes,
            "bb_period": 20,
            "rsi_period": 14,
            "bb_width": 0.05,
            "trending": {
                "rsi_threshold": [30, 70],
                "bb_std_dev_multiplier": 2.5,
            },
            "sideways": {
                "rsi_threshold": [40, 60],
                "bb_std_dev_multiplier": 1.8,
            },
            "SqueezeStrategy": True
        }
        
        # Initialize strategy
        strategy = BBRSIncrementalState(config)
        
        # Simulate real-time minute-by-minute processing
        results = []
        minute_count = 0
        bar_count = 0
        
        logging.info(f"Processing {len(minute_data)} minute-level data points...")
        
        for timestamp, row in minute_data.iterrows():
            minute_count += 1
            
            # Prepare minute-level OHLCV data
            minute_ohlcv = {
                'open': row['open'],
                'high': row['high'],
                'low': row['low'],
                'close': row['close'],
                'volume': row['volume']
            }
            
            # Update strategy with minute data
            result = strategy.update_minute_data(timestamp, minute_ohlcv)
            
            if result is not None:
                # A timeframe bar completed
                bar_count += 1
                results.append(result)
                
                # Log significant events
                if result['buy_signal']:
                    logging.info(f"🟢 BUY SIGNAL at {result['timestamp']} (Bar #{bar_count})")
                    logging.info(f"   Price: {result['close']:.2f}, RSI: {result['rsi']:.2f}, Regime: {result['market_regime']}")
                
                if result['sell_signal']:
                    logging.info(f"🔴 SELL SIGNAL at {result['timestamp']} (Bar #{bar_count})")
                    logging.info(f"   Price: {result['close']:.2f}, RSI: {result['rsi']:.2f}, Regime: {result['market_regime']}")
                
                # Log every 10th bar for monitoring
                if bar_count % 10 == 0:
                    logging.info(f"Processed {minute_count} minutes → {bar_count} {timeframe_minutes}min bars")
                    logging.info(f"   Current: Price={result['close']:.2f}, RSI={result['rsi']:.2f}, Regime={result['market_regime']}")
                    
                    # Show current incomplete bar
                    incomplete_bar = strategy.get_current_incomplete_bar()
                    if incomplete_bar:
                        logging.info(f"   Incomplete bar: Volume={incomplete_bar['volume']:.0f}")
        
        # Final statistics
        logging.info(f"\n📊 {timeframe_minutes}-minute Timeframe Results:")
        logging.info(f"   Minutes processed: {minute_count}")
        logging.info(f"   Bars generated: {bar_count}")
        logging.info(f"   Expected bars: ~{minute_count // timeframe_minutes}")
        logging.info(f"   Strategy warmed up: {strategy.is_warmed_up()}")
        
        if results:
            results_df = pd.DataFrame(results)
            buy_signals = results_df['buy_signal'].sum()
            sell_signals = results_df['sell_signal'].sum()
            
            logging.info(f"   Buy signals: {buy_signals}")
            logging.info(f"   Sell signals: {sell_signals}")
            
            # Show regime distribution
            regime_counts = results_df['market_regime'].value_counts()
            logging.info(f"   Market regimes: {dict(regime_counts)}")
            
            # Plot results for this timeframe
            plot_timeframe_results(results_df, timeframe_minutes)

def test_consistency_with_pre_aggregated():
    """Test that minute-level processing produces same results as pre-aggregated data."""
    
    logging.info(f"\n{'='*60}")
    logging.info("Testing consistency: Minute-level vs Pre-aggregated")
    logging.info(f"{'='*60}")
    
    # Load minute data
    minute_data = load_minute_data()
    if minute_data is None:
        return
    
    # Use smaller dataset for detailed comparison
    test_data = minute_data.iloc[:1440].copy()  # 24 hours of minute data
    
    timeframe_minutes = 60  # 1 hour
    
    config = {
        "timeframe_minutes": timeframe_minutes,
        "bb_period": 20,
        "rsi_period": 14,
        "bb_width": 0.05,
        "trending": {
            "rsi_threshold": [30, 70],
            "bb_std_dev_multiplier": 2.5,
        },
        "sideways": {
            "rsi_threshold": [40, 60],
            "bb_std_dev_multiplier": 1.8,
        },
        "SqueezeStrategy": True
    }
    
    # Method 1: Process minute-by-minute (real-time simulation)
    logging.info("Method 1: Processing minute-by-minute...")
    strategy_realtime = BBRSIncrementalState(config)
    realtime_results = []
    
    for timestamp, row in test_data.iterrows():
        minute_ohlcv = {
            'open': row['open'],
            'high': row['high'],
            'low': row['low'],
            'close': row['close'],
            'volume': row['volume']
        }
        
        result = strategy_realtime.update_minute_data(timestamp, minute_ohlcv)
        if result is not None:
            realtime_results.append(result)
    
    # Method 2: Pre-aggregate and process (traditional method)
    logging.info("Method 2: Processing pre-aggregated data...")
    from cycles.utils.data_utils import aggregate_to_hourly
    hourly_data = aggregate_to_hourly(test_data, 1)
    
    strategy_batch = BBRSIncrementalState(config)
    batch_results = []
    
    for timestamp, row in hourly_data.iterrows():
        hourly_ohlcv = {
            'open': row['open'],
            'high': row['high'],
            'low': row['low'],
            'close': row['close'],
            'volume': row['volume']
        }
        
        result = strategy_batch.update(hourly_ohlcv)
        batch_results.append(result)
    
    # Compare results
    logging.info("Comparing results...")
    
    realtime_df = pd.DataFrame(realtime_results)
    batch_df = pd.DataFrame(batch_results)
    
    logging.info(f"Real-time bars: {len(realtime_df)}")
    logging.info(f"Batch bars: {len(batch_df)}")
    
    if len(realtime_df) > 0 and len(batch_df) > 0:
        # Compare after warm-up
        warmup_bars = 25  # Conservative warm-up period
        
        if len(realtime_df) > warmup_bars and len(batch_df) > warmup_bars:
            rt_warmed = realtime_df.iloc[warmup_bars:]
            batch_warmed = batch_df.iloc[warmup_bars:]
            
            # Align by taking minimum length
            min_len = min(len(rt_warmed), len(batch_warmed))
            rt_aligned = rt_warmed.iloc[:min_len]
            batch_aligned = batch_warmed.iloc[:min_len]
            
            logging.info(f"Comparing {min_len} aligned bars after warm-up...")
            
            # Compare key metrics
            comparisons = [
                ('close', 'Close Price'),
                ('rsi', 'RSI'),
                ('upper_band', 'Upper Band'),
                ('lower_band', 'Lower Band'),
                ('middle_band', 'Middle Band'),
                ('buy_signal', 'Buy Signal'),
                ('sell_signal', 'Sell Signal')
            ]
            
            for col, name in comparisons:
                if col in rt_aligned.columns and col in batch_aligned.columns:
                    if col in ['buy_signal', 'sell_signal']:
                        # Boolean comparison
                        match_rate = (rt_aligned[col] == batch_aligned[col]).mean()
                        logging.info(f"{name}: {match_rate:.4f} match rate ({match_rate*100:.2f}%)")
                    else:
                        # Numerical comparison
                        diff = np.abs(rt_aligned[col] - batch_aligned[col])
                        max_diff = diff.max()
                        mean_diff = diff.mean()
                        logging.info(f"{name}: Max diff={max_diff:.6f}, Mean diff={mean_diff:.6f}")
            
            # Plot comparison
            plot_consistency_comparison(rt_aligned, batch_aligned)

def plot_timeframe_results(results_df, timeframe_minutes):
    """Plot results for a specific timeframe."""
    
    if len(results_df) < 10:
        logging.warning(f"Not enough data to plot for {timeframe_minutes}min timeframe")
        return
    
    fig, axes = plt.subplots(3, 1, figsize=(15, 10))
    
    # Plot 1: Price and Bollinger Bands
    axes[0].plot(results_df.index, results_df['close'], 'k-', label='Close Price', alpha=0.8)
    axes[0].plot(results_df.index, results_df['upper_band'], 'b-', label='Upper Band', alpha=0.7)
    axes[0].plot(results_df.index, results_df['middle_band'], 'g-', label='Middle Band', alpha=0.7)
    axes[0].plot(results_df.index, results_df['lower_band'], 'r-', label='Lower Band', alpha=0.7)
    
    # Mark signals
    buy_signals = results_df[results_df['buy_signal']]
    sell_signals = results_df[results_df['sell_signal']]
    
    if len(buy_signals) > 0:
        axes[0].scatter(buy_signals.index, buy_signals['close'], 
                       color='green', marker='^', s=100, label='Buy Signal', zorder=5)
    
    if len(sell_signals) > 0:
        axes[0].scatter(sell_signals.index, sell_signals['close'], 
                       color='red', marker='v', s=100, label='Sell Signal', zorder=5)
    
    axes[0].set_title(f'{timeframe_minutes}-minute Timeframe: Price and Bollinger Bands')
    axes[0].legend()
    axes[0].grid(True)
    
    # Plot 2: RSI
    axes[1].plot(results_df.index, results_df['rsi'], 'purple', label='RSI', alpha=0.8)
    axes[1].axhline(y=70, color='red', linestyle='--', alpha=0.5, label='Overbought')
    axes[1].axhline(y=30, color='green', linestyle='--', alpha=0.5, label='Oversold')
    axes[1].set_title('RSI')
    axes[1].legend()
    axes[1].grid(True)
    axes[1].set_ylim(0, 100)
    
    # Plot 3: Market Regime
    regime_numeric = [1 if regime == 'sideways' else 0 for regime in results_df['market_regime']]
    axes[2].plot(results_df.index, regime_numeric, 'orange', label='Market Regime', alpha=0.8)
    axes[2].set_title('Market Regime (1=Sideways, 0=Trending)')
    axes[2].legend()
    axes[2].grid(True)
    axes[2].set_ylim(-0.1, 1.1)
    
    plt.tight_layout()
    save_path = f"realtime_bbrs_{timeframe_minutes}min.png"
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    logging.info(f"Plot saved to {save_path}")
    plt.show()

def plot_consistency_comparison(realtime_df, batch_df):
    """Plot comparison between real-time and batch processing."""
    
    fig, axes = plt.subplots(2, 1, figsize=(15, 8))
    
    # Plot 1: Price and signals comparison
    axes[0].plot(realtime_df.index, realtime_df['close'], 'k-', label='Price', alpha=0.8)
    
    # Real-time signals
    rt_buy = realtime_df[realtime_df['buy_signal']]
    rt_sell = realtime_df[realtime_df['sell_signal']]
    
    if len(rt_buy) > 0:
        axes[0].scatter(rt_buy.index, rt_buy['close'], 
                       color='green', marker='^', s=80, label='Real-time Buy', alpha=0.8)
    
    if len(rt_sell) > 0:
        axes[0].scatter(rt_sell.index, rt_sell['close'], 
                       color='red', marker='v', s=80, label='Real-time Sell', alpha=0.8)
    
    # Batch signals
    batch_buy = batch_df[batch_df['buy_signal']]
    batch_sell = batch_df[batch_df['sell_signal']]
    
    if len(batch_buy) > 0:
        axes[0].scatter(batch_buy.index, batch_buy['close'], 
                       color='lightgreen', marker='s', s=60, label='Batch Buy', alpha=0.6)
    
    if len(batch_sell) > 0:
        axes[0].scatter(batch_sell.index, batch_sell['close'], 
                       color='lightcoral', marker='s', s=60, label='Batch Sell', alpha=0.6)
    
    axes[0].set_title('Signal Comparison: Real-time vs Batch Processing')
    axes[0].legend()
    axes[0].grid(True)
    
    # Plot 2: RSI comparison
    axes[1].plot(realtime_df.index, realtime_df['rsi'], 'b-', label='Real-time RSI', alpha=0.8)
    axes[1].plot(batch_df.index, batch_df['rsi'], 'r--', label='Batch RSI', alpha=0.8)
    axes[1].set_title('RSI Comparison')
    axes[1].legend()
    axes[1].grid(True)
    
    plt.tight_layout()
    save_path = "realtime_vs_batch_comparison.png"
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    logging.info(f"Comparison plot saved to {save_path}")
    plt.show()

def main():
    """Main test function."""
    logging.info("Starting real-time BBRS strategy validation test")
    
    try:
        # Test 1: Different timeframe aggregations
        test_timeframe_aggregation()
        
        # Test 2: Consistency with pre-aggregated data
        test_consistency_with_pre_aggregated()
        
        logging.info("Real-time BBRS strategy test completed successfully!")
    except Exception as e:
        logging.error(f"Test failed with error: {e}")
        raise

if __name__ == "__main__":
    main()