"""
Backtesting engine for TradingAgent
Simulates trading with max/min predictions
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from loguru import logger
import json


@dataclass
class Trade:
    """Single trade record"""
    entry_time: datetime
    exit_time: Optional[datetime]
    symbol: str
    side: str  # 'long' or 'short'
    entry_price: float
    exit_price: Optional[float]
    quantity: float
    stop_loss: Optional[float]
    take_profit: Optional[float]
    profit_loss: Optional[float] = None
    profit_loss_pct: Optional[float] = None
    status: str = 'open'  # 'open', 'closed', 'stopped'
    strategy: str = 'maxmin'
    horizon: str = 'scalping'
    
    def close(self, exit_price: float, exit_time: datetime):
        """Close the trade"""
        self.exit_price = exit_price
        self.exit_time = exit_time
        self.status = 'closed'
        
        if self.side == 'long':
            self.profit_loss = (exit_price - self.entry_price) * self.quantity
        else:  # short
            self.profit_loss = (self.entry_price - exit_price) * self.quantity
        
        self.profit_loss_pct = (self.profit_loss / (self.entry_price * self.quantity)) * 100
        
        return self.profit_loss


@dataclass
class BacktestResult:
    """Backtesting results"""
    trades: List[Trade]
    total_trades: int
    winning_trades: int
    losing_trades: int
    win_rate: float
    total_profit: float
    total_profit_pct: float
    max_drawdown: float
    max_drawdown_pct: float
    sharpe_ratio: float
    sortino_ratio: float
    profit_factor: float
    avg_win: float
    avg_loss: float
    best_trade: float
    worst_trade: float
    avg_trade_duration: timedelta
    equity_curve: pd.Series
    metrics: Dict[str, Any] = field(default_factory=dict)


class MaxMinBacktester:
    """Backtesting engine for max/min predictions"""
    
    def __init__(
        self,
        initial_capital: float = 10000,
        position_size: float = 0.1,  # 10% of capital per trade
        max_positions: int = 3,
        commission: float = 0.001,  # 0.1%
        slippage: float = 0.0005  # 0.05%
    ):
        """
        Initialize backtester
        
        Args:
            initial_capital: Starting capital
            position_size: Position size as fraction of capital
            max_positions: Maximum concurrent positions
            commission: Commission rate
            slippage: Slippage rate
        """
        self.initial_capital = initial_capital
        self.position_size = position_size
        self.max_positions = max_positions
        self.commission = commission
        self.slippage = slippage
        
        self.reset()
        
    def reset(self):
        """Reset backtester state"""
        self.capital = self.initial_capital
        self.trades = []
        self.open_trades = []
        self.equity_curve = []
        self.positions = 0
        
    def run(
        self,
        data: pd.DataFrame,
        predictions: pd.DataFrame,
        strategy: str = 'conservative',
        horizon: str = 'scalping'
    ) -> BacktestResult:
        """
        Run backtest with max/min predictions
        
        Args:
            data: OHLCV data
            predictions: DataFrame with prediction columns (pred_high, pred_low, confidence)
            strategy: Trading strategy ('conservative', 'balanced', 'aggressive')
            horizon: Trading horizon
            
        Returns:
            BacktestResult with performance metrics
        """
        self.reset()
        
        # Merge data and predictions
        df = data.join(predictions, how='inner')
        
        # Strategy parameters
        confidence_threshold = {
            'conservative': 0.7,
            'balanced': 0.6,
            'aggressive': 0.5
        }[strategy]
        
        risk_reward_ratio = {
            'conservative': 2.0,
            'balanced': 1.5,
            'aggressive': 1.0
        }[strategy]
        
        # Iterate through data
        for idx, row in df.iterrows():
            current_price = row['close']
            
            # Update open trades
            self._update_open_trades(row, idx)
            
            # Check for entry signals
            if self.positions < self.max_positions:
                signal = self._generate_signal(row, confidence_threshold)
                
                if signal:
                    self._enter_trade(
                        signal=signal,
                        row=row,
                        time=idx,
                        risk_reward_ratio=risk_reward_ratio,
                        horizon=horizon
                    )
            
            # Record equity
            equity = self._calculate_equity(current_price)
            self.equity_curve.append({
                'time': idx,
                'equity': equity,
                'capital': self.capital,
                'positions': self.positions
            })
        
        # Close any remaining trades
        self._close_all_trades(df.iloc[-1]['close'], df.index[-1])
        
        # Calculate metrics
        return self._calculate_metrics()
    
    def _generate_signal(self, row: pd.Series, confidence_threshold: float) -> Optional[str]:
        """
        Generate trading signal based on predictions
        
        Returns:
            'long', 'short', or None
        """
        if 'confidence' not in row or pd.isna(row['confidence']):
            return None
            
        if row['confidence'] < confidence_threshold:
            return None
        
        current_price = row['close']
        pred_high = row.get('pred_high', np.nan)
        pred_low = row.get('pred_low', np.nan)
        
        if pd.isna(pred_high) or pd.isna(pred_low):
            return None
        
        # Calculate potential profits
        long_profit = (pred_high - current_price) / current_price
        short_profit = (current_price - pred_low) / current_price
        
        # Generate signal based on risk/reward
        min_profit_threshold = 0.005  # 0.5% minimum expected profit
        
        if long_profit > min_profit_threshold and long_profit > short_profit:
            # Check if we're closer to predicted low (better entry for long)
            if (current_price - pred_low) / (pred_high - pred_low) < 0.3:
                return 'long'
        elif short_profit > min_profit_threshold:
            # Check if we're closer to predicted high (better entry for short)
            if (pred_high - current_price) / (pred_high - pred_low) < 0.3:
                return 'short'
        
        return None
    
    def _enter_trade(
        self,
        signal: str,
        row: pd.Series,
        time: datetime,
        risk_reward_ratio: float,
        horizon: str
    ):
        """Enter a new trade"""
        entry_price = row['close']
        
        # Apply slippage
        if signal == 'long':
            entry_price *= (1 + self.slippage)
        else:
            entry_price *= (1 - self.slippage)
        
        # Calculate position size
        position_value = self.capital * self.position_size
        quantity = position_value / entry_price
        
        # Apply commission
        commission_cost = position_value * self.commission
        self.capital -= commission_cost
        
        # Set stop loss and take profit
        if signal == 'long':
            stop_loss = row['pred_low'] * 0.98  # 2% below predicted low
            take_profit = row['pred_high'] * 0.98  # 2% below predicted high
        else:
            stop_loss = row['pred_high'] * 1.02  # 2% above predicted high
            take_profit = row['pred_low'] * 1.02  # 2% above predicted low
        
        # Create trade
        trade = Trade(
            entry_time=time,
            exit_time=None,
            symbol='',  # Will be set by caller
            side=signal,
            entry_price=entry_price,
            exit_price=None,
            quantity=quantity,
            stop_loss=stop_loss,
            take_profit=take_profit,
            strategy='maxmin',
            horizon=horizon
        )
        
        self.open_trades.append(trade)
        self.trades.append(trade)
        self.positions += 1
        
        logger.debug(f"📈 Entered {signal} trade at {entry_price:.2f}")
    
    def _update_open_trades(self, row: pd.Series, time: datetime):
        """Update open trades with current prices"""
        current_price = row['close']
        
        for trade in self.open_trades[:]:
            # Check stop loss
            if trade.side == 'long' and current_price <= trade.stop_loss:
                self._close_trade(trade, trade.stop_loss, time, 'stopped')
            elif trade.side == 'short' and current_price >= trade.stop_loss:
                self._close_trade(trade, trade.stop_loss, time, 'stopped')
            
            # Check take profit
            elif trade.side == 'long' and current_price >= trade.take_profit:
                self._close_trade(trade, trade.take_profit, time, 'profit')
            elif trade.side == 'short' and current_price <= trade.take_profit:
                self._close_trade(trade, trade.take_profit, time, 'profit')
    
    def _close_trade(self, trade: Trade, exit_price: float, time: datetime, reason: str):
        """Close a trade"""
        # Apply slippage
        if trade.side == 'long':
            exit_price *= (1 - self.slippage)
        else:
            exit_price *= (1 + self.slippage)
        
        # Close trade
        profit_loss = trade.close(exit_price, time)
        
        # Apply commission
        commission_cost = abs(trade.quantity * exit_price) * self.commission
        profit_loss -= commission_cost
        
        # Update capital
        self.capital += (trade.quantity * exit_price) - commission_cost
        
        # Remove from open trades
        self.open_trades.remove(trade)
        self.positions -= 1
        
        logger.debug(f"📉 Closed {trade.side} trade: {profit_loss:+.2f} ({reason})")
    
    def _close_all_trades(self, price: float, time: datetime):
        """Close all open trades"""
        for trade in self.open_trades[:]:
            self._close_trade(trade, price, time, 'end')
    
    def _calculate_equity(self, current_price: float) -> float:
        """Calculate current equity"""
        equity = self.capital
        
        for trade in self.open_trades:
            if trade.side == 'long':
                unrealized = (current_price - trade.entry_price) * trade.quantity
            else:
                unrealized = (trade.entry_price - current_price) * trade.quantity
            equity += unrealized
        
        return equity
    
    def _calculate_metrics(self) -> BacktestResult:
        """Calculate backtesting metrics"""
        if not self.trades:
            return BacktestResult(
                trades=[], total_trades=0, winning_trades=0, losing_trades=0,
                win_rate=0, total_profit=0, total_profit_pct=0,
                max_drawdown=0, max_drawdown_pct=0, sharpe_ratio=0,
                sortino_ratio=0, profit_factor=0, avg_win=0, avg_loss=0,
                best_trade=0, worst_trade=0,
                avg_trade_duration=timedelta(0),
                equity_curve=pd.Series()
            )
        
        # Filter closed trades
        closed_trades = [t for t in self.trades if t.status == 'closed']
        
        if not closed_trades:
            return BacktestResult(
                trades=self.trades, total_trades=len(self.trades),
                winning_trades=0, losing_trades=0, win_rate=0,
                total_profit=0, total_profit_pct=0,
                max_drawdown=0, max_drawdown_pct=0, sharpe_ratio=0,
                sortino_ratio=0, profit_factor=0, avg_win=0, avg_loss=0,
                best_trade=0, worst_trade=0,
                avg_trade_duration=timedelta(0),
                equity_curve=pd.Series()
            )
        
        # Basic metrics
        profits = [t.profit_loss for t in closed_trades]
        winning_trades = [t for t in closed_trades if t.profit_loss > 0]
        losing_trades = [t for t in closed_trades if t.profit_loss <= 0]
        
        total_profit = sum(profits)
        total_profit_pct = (total_profit / self.initial_capital) * 100
        
        # Win rate
        win_rate = len(winning_trades) / len(closed_trades) if closed_trades else 0
        
        # Average win/loss
        avg_win = np.mean([t.profit_loss for t in winning_trades]) if winning_trades else 0
        avg_loss = np.mean([t.profit_loss for t in losing_trades]) if losing_trades else 0
        
        # Profit factor
        gross_profit = sum(t.profit_loss for t in winning_trades) if winning_trades else 0
        gross_loss = abs(sum(t.profit_loss for t in losing_trades)) if losing_trades else 1
        profit_factor = gross_profit / gross_loss if gross_loss > 0 else 0
        
        # Best/worst trade
        best_trade = max(profits) if profits else 0
        worst_trade = min(profits) if profits else 0
        
        # Trade duration
        durations = [(t.exit_time - t.entry_time) for t in closed_trades if t.exit_time]
        avg_trade_duration = np.mean(durations) if durations else timedelta(0)
        
        # Equity curve
        equity_df = pd.DataFrame(self.equity_curve)
        if not equity_df.empty:
            equity_df.set_index('time', inplace=True)
            equity_series = equity_df['equity']
            
            # Drawdown
            cummax = equity_series.cummax()
            drawdown = (equity_series - cummax) / cummax
            max_drawdown_pct = drawdown.min() * 100
            max_drawdown = (equity_series - cummax).min()
            
            # Sharpe ratio (assuming 0 risk-free rate)
            returns = equity_series.pct_change().dropna()
            if len(returns) > 1:
                sharpe_ratio = np.sqrt(252) * returns.mean() / returns.std()
            else:
                sharpe_ratio = 0
            
            # Sortino ratio
            negative_returns = returns[returns < 0]
            if len(negative_returns) > 0:
                sortino_ratio = np.sqrt(252) * returns.mean() / negative_returns.std()
            else:
                sortino_ratio = sharpe_ratio
        else:
            equity_series = pd.Series()
            max_drawdown = 0
            max_drawdown_pct = 0
            sharpe_ratio = 0
            sortino_ratio = 0
        
        return BacktestResult(
            trades=self.trades,
            total_trades=len(closed_trades),
            winning_trades=len(winning_trades),
            losing_trades=len(losing_trades),
            win_rate=win_rate,
            total_profit=total_profit,
            total_profit_pct=total_profit_pct,
            max_drawdown=max_drawdown,
            max_drawdown_pct=max_drawdown_pct,
            sharpe_ratio=sharpe_ratio,
            sortino_ratio=sortino_ratio,
            profit_factor=profit_factor,
            avg_win=avg_win,
            avg_loss=avg_loss,
            best_trade=best_trade,
            worst_trade=worst_trade,
            avg_trade_duration=avg_trade_duration,
            equity_curve=equity_series,
            metrics={
                'total_commission': len(closed_trades) * 2 * self.commission * self.initial_capital * self.position_size,
                'total_slippage': len(closed_trades) * 2 * self.slippage * self.initial_capital * self.position_size,
                'final_capital': self.capital,
                'roi': ((self.capital - self.initial_capital) / self.initial_capital) * 100
            }
        )
    
    def plot_results(self, result: BacktestResult, save_path: Optional[str] = None):
        """Plot backtesting results"""
        import matplotlib.pyplot as plt
        import seaborn as sns
        
        sns.set_style('darkgrid')
        
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        fig.suptitle('Backtesting Results - Max/Min Strategy', fontsize=16)
        
        # Equity curve
        ax = axes[0, 0]
        result.equity_curve.plot(ax=ax, color='blue', linewidth=2)
        ax.set_title('Equity Curve')
        ax.set_xlabel('Time')
        ax.set_ylabel('Equity ($)')
        ax.grid(True, alpha=0.3)
        
        # Drawdown
        ax = axes[0, 1]
        cummax = result.equity_curve.cummax()
        drawdown = (result.equity_curve - cummax) / cummax * 100
        drawdown.plot(ax=ax, color='red', linewidth=2)
        ax.fill_between(drawdown.index, drawdown.values, 0, alpha=0.3, color='red')
        ax.set_title('Drawdown')
        ax.set_xlabel('Time')
        ax.set_ylabel('Drawdown (%)')
        ax.grid(True, alpha=0.3)
        
        # Trade distribution
        ax = axes[1, 0]
        profits = [t.profit_loss for t in result.trades if t.profit_loss is not None]
        if profits:
            ax.hist(profits, bins=30, color='green', alpha=0.7, edgecolor='black')
            ax.axvline(0, color='red', linestyle='--', linewidth=2)
            ax.set_title('Profit/Loss Distribution')
            ax.set_xlabel('Profit/Loss ($)')
            ax.set_ylabel('Frequency')
        ax.grid(True, alpha=0.3)
        
        # Metrics summary
        ax = axes[1, 1]
        ax.axis('off')
        
        metrics_text = f"""
        Total Trades: {result.total_trades}
        Win Rate: {result.win_rate:.1%}
        Total Profit: ${result.total_profit:,.2f}
        ROI: {result.total_profit_pct:.1f}%
        
        Max Drawdown: {result.max_drawdown_pct:.1f}%
        Sharpe Ratio: {result.sharpe_ratio:.2f}
        Profit Factor: {result.profit_factor:.2f}
        
        Avg Win: ${result.avg_win:,.2f}
        Avg Loss: ${result.avg_loss:,.2f}
        Best Trade: ${result.best_trade:,.2f}
        Worst Trade: ${result.worst_trade:,.2f}
        """
        
        ax.text(0.1, 0.5, metrics_text, fontsize=12, verticalalignment='center',
                fontfamily='monospace')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=100)
            logger.info(f"📊 Saved backtest results to {save_path}")
        
        return fig