import logging
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from typing import Dict, Any

logger = logging.getLogger(__name__)

class BacktestEngine:
    def __init__(self, initial_capital: float = 10000.0):
        self.initial_capital = initial_capital
        self.capital = initial_capital
        self.positions = 0.0 # simplified sizing (e.g. holding 100 shares of YES)
        self.trades = []
        self.equity_curve = []

    def load_data(self, data: pd.DataFrame):
        """
        Expects a DataFrame with ['timestamp', 'bid', 'ask', 'mid', 'volume']
        and external features like ['sentiment'] if testing the momentum strategy.
        """
        self.df = data.sort_values("timestamp").reset_index(drop=True)
        # Assuming slippage model: market orders execute at worst of bid/ask + slippage
        self.slippage_bps = 5 

    def run_macrostem_simulation(self, strategy_func):
        """
        Iterate through tick/candle data, evaluating the strategy function.
        strategy_func takes a dict row and returns a trade action dict or None.
        """
        logger.info(f"Starting backtest with {len(self.df)} ticks. Capital: ${self.capital}")
        
        for idx, row in self.df.iterrows():
            # 1. Update Equity based on current holdings marked to market (mid price)
            mtm_value = self.capital + (self.positions * row['mid'])
            self.equity_curve.append({'timestamp': row['timestamp'], 'equity': mtm_value})
            
            # 2. Evaluate Strategy Signal
            signal = strategy_func(row)
            if not signal:
                 continue
                 
            action = signal.get("action")
            
            # Execute Trade Simulation
            if action == "BUY_YES":
                # Buy 100 shares at ask price + slippage
                execute_price = row['ask'] + (row['ask'] * self.slippage_bps / 10000)
                cost = execute_price * 100
                
                if self.capital >= cost:
                    self.capital -= cost
                    self.positions += 100
                    self.trades.append({
                        "timestamp": row['timestamp'], "action": action, 
                        "price": execute_price, "size": 100, "pnl": 0
                    })
                    
            elif action == "SELL_YES" and self.positions > 0:
                # Close position at bid price - slippage
                execute_price = row['bid'] - (row['bid'] * self.slippage_bps / 10000)
                revenue = execute_price * self.positions
                pnl = revenue - sum([t['price']*t['size'] for t in self.trades if t['action']=='BUY_YES']) # Naive PnL tracking for demo
                
                self.capital += revenue
                self.trades.append({
                    "timestamp": row['timestamp'], "action": action, 
                    "price": execute_price, "size": self.positions, "pnl": pnl
                })
                self.positions = 0.0
                
        # Final MTM
        final_equity = self.capital + (self.positions * self.df.iloc[-1]['mid'])
        logger.info(f"Backtest Complete. Final Equity: ${final_equity:.2f}")

    def calculate_metrics(self) -> Dict[str, Any]:
        """Calculate Sharpe, Sortino, Max Drawdown, etc."""
        if not self.equity_curve:
            return {}
            
        equity_df = pd.DataFrame(self.equity_curve)
        equity_df.set_index('timestamp', inplace=True)
        
        # Calculate returns
        returns = equity_df['equity'].pct_change().dropna()
        
        # Annualization factor (assuming hourly data for example, 252*24 approx 6048)
        # Using a generic 252 * 252 for Crypto/Prediction markets pseudo-continuous
        annualization = 365 * 24 
        
        mean_ret = returns.mean() * annualization
        std_ret = returns.std() * np.sqrt(annualization)
        
        sharpe = mean_ret / std_ret if std_ret != 0 else 0
        
        # Max Drawdown
        cum_ret = (1 + returns).cumprod()
        rolling_max = cum_ret.cummax()
        drawdowns = (cum_ret - rolling_max) / rolling_max
        max_dd = drawdowns.min()
        
        # Win Rate
        winning_trades = len([t for t in self.trades if t.get('pnl', 0) > 0])
        total_closed = len([t for t in self.trades if t['action'] == 'SELL_YES'])
        win_rate = winning_trades / total_closed if total_closed > 0 else 0
        
        return {
            "Total Return (%)": ((equity_df['equity'].iloc[-1] / self.initial_capital) - 1) * 100,
            "Sharpe Ratio": sharpe,
            "Max Drawdown (%)": max_dd * 100,
            "Win Rate (%)": win_rate * 100,
            "Total Trades": len(self.trades)
        }

    def generate_report(self):
        """Optional: Output Plotly charts or console summary."""
        metrics = self.calculate_metrics()
        logger.info("=== Backtest Performance Report ===")
        for k, v in metrics.items():
            logger.info(f"{k}: {v:.2f}")