"""
Stock Market Time Series Dataset - Backtesting Framework

Author: Molla Samser
Organization: RSK World
Designer & Tester: Rima Khatun
Website: https://rskworld.in/
Email: help@rskworld.in
Phone: +91 93305 39277
Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
from typing import Dict, List, Optional, Callable
import warnings
warnings.filterwarnings('ignore')


class TradingStrategy:
    """
    Base class for trading strategies.
    """
    
    def __init__(self, name: str = "Strategy"):
        """
        Initialize trading strategy.
        
        Args:
            name: Strategy name
        """
        self.name = name
    
    def generate_signals(self, df: pd.DataFrame) -> pd.Series:
        """
        Generate trading signals.
        
        Args:
            df: DataFrame with stock data
            
        Returns:
            Series with signals (1: buy, -1: sell, 0: hold)
        """
        raise NotImplementedError("Subclasses must implement generate_signals()")


class MovingAverageCrossover(TradingStrategy):
    """
    Moving Average Crossover Strategy.
    """
    
    def __init__(self, short_window: int = 20, long_window: int = 50):
        """
        Initialize MA Crossover strategy.
        
        Args:
            short_window: Short MA period
            long_window: Long MA period
        """
        super().__init__(f"MA_Crossover_{short_window}_{long_window}")
        self.short_window = short_window
        self.long_window = long_window
    
    def generate_signals(self, df: pd.DataFrame) -> pd.Series:
        """Generate signals based on MA crossover."""
        signals = pd.Series(0, index=df.index)
        
        # Calculate MAs
        short_ma = df['Close'].rolling(window=self.short_window).mean()
        long_ma = df['Close'].rolling(window=self.long_window).mean()
        
        # Generate signals
        signals[short_ma > long_ma] = 1  # Buy
        signals[short_ma < long_ma] = -1  # Sell
        
        return signals


class RSIStrategy(TradingStrategy):
    """
    RSI-based Trading Strategy.
    """
    
    def __init__(self, oversold: int = 30, overbought: int = 70):
        """
        Initialize RSI strategy.
        
        Args:
            oversold: RSI oversold threshold
            overbought: RSI overbought threshold
        """
        super().__init__(f"RSI_{oversold}_{overbought}")
        self.oversold = oversold
        self.overbought = overbought
    
    def generate_signals(self, df: pd.DataFrame) -> pd.Series:
        """Generate signals based on RSI."""
        signals = pd.Series(0, index=df.index)
        
        if 'RSI' in df.columns:
            rsi = df['RSI']
        else:
            # Calculate RSI if not present
            delta = df['Close'].diff()
            gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
            loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
            rs = gain / loss
            rsi = 100 - (100 / (1 + rs))
        
        # Generate signals
        signals[rsi < self.oversold] = 1  # Buy (oversold)
        signals[rsi > self.overbought] = -1  # Sell (overbought)
        
        return signals


class MACDStrategy(TradingStrategy):
    """
    MACD-based Trading Strategy.
    """
    
    def __init__(self):
        """Initialize MACD strategy."""
        super().__init__("MACD_Crossover")
    
    def generate_signals(self, df: pd.DataFrame) -> pd.Series:
        """Generate signals based on MACD."""
        signals = pd.Series(0, index=df.index)
        
        if 'MACD' in df.columns:
            macd = df['MACD']
            # Calculate signal line (9-day EMA of MACD)
            signal_line = macd.ewm(span=9, adjust=False).mean()
            
            # Generate signals
            signals[macd > signal_line] = 1  # Buy
            signals[macd < signal_line] = -1  # Sell
        
        return signals


class Backtester:
    """
    Comprehensive backtesting framework.
    """
    
    def __init__(self, data_path: str, initial_capital: float = 100000):
        """
        Initialize backtester.
        
        Args:
            data_path: Path to stock CSV file
            initial_capital: Starting capital in USD
        """
        self.data_path = Path(data_path)
        self.symbol = self.data_path.stem
        self.df = self._load_data()
        self.initial_capital = initial_capital
        self.results = {}
        
        sns.set_style('whitegrid')
        plt.rcParams['figure.figsize'] = (14, 8)
    
    def _load_data(self) -> pd.DataFrame:
        """Load and prepare data."""
        df = pd.read_csv(self.data_path, parse_dates=['Date'], index_col='Date')
        df.sort_index(inplace=True)
        return df
    
    def run_backtest(self, strategy: TradingStrategy, 
                    commission: float = 0.001,
                    slippage: float = 0.0005) -> Dict:
        """
        Run backtest for a strategy.
        
        Args:
            strategy: Trading strategy instance
            commission: Commission per trade (as fraction)
            slippage: Slippage per trade (as fraction)
            
        Returns:
            Dictionary with backtest results
        """
        print(f"\n{'='*70}")
        print(f"Backtesting {strategy.name} on {self.symbol}")
        print(f"{'='*70}\n")
        
        # Generate signals
        signals = strategy.generate_signals(self.df)
        
        # Initialize variables
        capital = self.initial_capital
        shares = 0
        portfolio_value = []
        trades = []
        
        # Run backtest
        for i, (date, signal) in enumerate(signals.items()):
            price = self.df.loc[date, 'Close']
            
            # Calculate portfolio value
            current_value = capital + (shares * price)
            portfolio_value.append(current_value)
            
            # Execute trades
            if signal == 1 and shares == 0:  # Buy signal
                # Buy as many shares as possible
                cost_per_share = price * (1 + slippage + commission)
                shares = int(capital / cost_per_share)
                cost = shares * cost_per_share
                capital -= cost
                
                trades.append({
                    'date': date,
                    'action': 'BUY',
                    'price': price,
                    'shares': shares,
                    'value': cost
                })
                
            elif signal == -1 and shares > 0:  # Sell signal
                # Sell all shares
                revenue_per_share = price * (1 - slippage - commission)
                revenue = shares * revenue_per_share
                capital += revenue
                
                trades.append({
                    'date': date,
                    'action': 'SELL',
                    'price': price,
                    'shares': shares,
                    'value': revenue
                })
                
                shares = 0
        
        # Close any open positions
        if shares > 0:
            final_price = self.df['Close'].iloc[-1]
            revenue = shares * final_price * (1 - slippage - commission)
            capital += revenue
            portfolio_value[-1] = capital
        
        # Calculate metrics
        portfolio_series = pd.Series(portfolio_value, index=self.df.index)
        returns = portfolio_series.pct_change()
        
        metrics = self._calculate_metrics(portfolio_series, returns, trades)
        
        # Print results
        self._print_results(metrics, trades)
        
        # Store results
        results = {
            'strategy': strategy.name,
            'signals': signals,
            'portfolio_value': portfolio_series,
            'returns': returns,
            'trades': trades,
            'metrics': metrics
        }
        
        self.results[strategy.name] = results
        
        return results
    
    def _calculate_metrics(self, portfolio_series: pd.Series, 
                          returns: pd.Series, trades: List[Dict]) -> Dict:
        """Calculate performance metrics."""
        final_value = portfolio_series.iloc[-1]
        total_return = (final_value - self.initial_capital) / self.initial_capital
        
        # Annualized metrics
        days = len(portfolio_series)
        years = days / 252
        cagr = (final_value / self.initial_capital) ** (1 / years) - 1
        
        # Risk metrics
        volatility = returns.std() * np.sqrt(252)
        sharpe_ratio = (returns.mean() * 252) / (returns.std() * np.sqrt(252)) if returns.std() > 0 else 0
        
        # Drawdown
        cumulative = (1 + returns).cumprod()
        running_max = cumulative.expanding().max()
        drawdown = (cumulative - running_max) / running_max
        max_drawdown = drawdown.min()
        
        # Win rate
        winning_trades = [t for i, t in enumerate(trades[1::2]) 
                         if i < len(trades[::2]) and 
                         t['value'] > trades[::2][i+1]['value']]
        total_closed_trades = len(trades) // 2
        win_rate = len(winning_trades) / total_closed_trades if total_closed_trades > 0 else 0
        
        return {
            'initial_capital': self.initial_capital,
            'final_value': final_value,
            'total_return': total_return,
            'cagr': cagr,
            'volatility': volatility,
            'sharpe_ratio': sharpe_ratio,
            'max_drawdown': max_drawdown,
            'total_trades': len(trades),
            'win_rate': win_rate,
            'total_days': days
        }
    
    def _print_results(self, metrics: Dict, trades: List[Dict]):
        """Print backtest results."""
        print(f"{'Metric':<25} {'Value':<25}")
        print(f"{'-'*50}")
        print(f"{'Initial Capital:':<25} ${metrics['initial_capital']:,.2f}")
        print(f"{'Final Value:':<25} ${metrics['final_value']:,.2f}")
        print(f"{'Total Return:':<25} {metrics['total_return']:.2%}")
        print(f"{'CAGR:':<25} {metrics['cagr']:.2%}")
        print(f"{'Volatility (Annual):':<25} {metrics['volatility']:.2%}")
        print(f"{'Sharpe Ratio:':<25} {metrics['sharpe_ratio']:.4f}")
        print(f"{'Max Drawdown:':<25} {metrics['max_drawdown']:.2%}")
        print(f"{'Total Trades:':<25} {metrics['total_trades']}")
        print(f"{'Win Rate:':<25} {metrics['win_rate']:.2%}")
        print(f"{'Trading Days:':<25} {metrics['total_days']}")
        print(f"{'='*50}\n")
    
    def plot_results(self, strategy_name: str, save_path: Optional[str] = None):
        """
        Plot backtest results.
        
        Args:
            strategy_name: Name of strategy to plot
            save_path: Path to save the plot
        """
        if strategy_name not in self.results:
            print(f"No results found for strategy: {strategy_name}")
            return
        
        results = self.results[strategy_name]
        
        fig, axes = plt.subplots(3, 1, figsize=(14, 12))
        
        # Portfolio value
        axes[0].plot(results['portfolio_value'], linewidth=2, label='Strategy')
        buy_and_hold = self.initial_capital * (self.df['Close'] / self.df['Close'].iloc[0])
        axes[0].plot(buy_and_hold, linewidth=2, alpha=0.7, label='Buy & Hold', linestyle='--')
        axes[0].set_title(f'{strategy_name} - Portfolio Value', fontsize=14, fontweight='bold')
        axes[0].set_ylabel('Portfolio Value ($)', fontsize=12)
        axes[0].legend(loc='best')
        axes[0].grid(True, alpha=0.3)
        
        # Mark trades
        trades_df = pd.DataFrame(results['trades'])
        if len(trades_df) > 0:
            buys = trades_df[trades_df['action'] == 'BUY']
            sells = trades_df[trades_df['action'] == 'SELL']
            
            for _, trade in buys.iterrows():
                axes[0].axvline(x=trade['date'], color='green', alpha=0.3, linestyle=':')
            for _, trade in sells.iterrows():
                axes[0].axvline(x=trade['date'], color='red', alpha=0.3, linestyle=':')
        
        # Returns distribution
        axes[1].hist(results['returns'].dropna() * 100, bins=50, alpha=0.7, color='blue', edgecolor='black')
        axes[1].axvline(results['returns'].mean() * 100, color='red', linestyle='--', 
                       linewidth=2, label=f'Mean: {results["returns"].mean()*100:.2f}%')
        axes[1].set_title('Daily Returns Distribution', fontsize=12, fontweight='bold')
        axes[1].set_xlabel('Daily Returns (%)', fontsize=10)
        axes[1].set_ylabel('Frequency', fontsize=10)
        axes[1].legend()
        axes[1].grid(True, alpha=0.3)
        
        # Drawdown
        cumulative = (1 + results['returns']).cumprod()
        running_max = cumulative.expanding().max()
        drawdown = (cumulative - running_max) / running_max
        axes[2].fill_between(drawdown.index, drawdown * 100, 0, color='red', alpha=0.3)
        axes[2].set_title('Drawdown', fontsize=12, fontweight='bold')
        axes[2].set_xlabel('Date', fontsize=10)
        axes[2].set_ylabel('Drawdown (%)', fontsize=10)
        axes[2].grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Results plot saved to: {save_path}")
        
        plt.show()
    
    def compare_strategies(self, strategies: List[TradingStrategy],
                          save_path: Optional[str] = None):
        """
        Compare multiple strategies.
        
        Args:
            strategies: List of strategy instances
            save_path: Path to save the plot
        """
        print(f"\n{'='*70}")
        print(f"Comparing {len(strategies)} Strategies on {self.symbol}")
        print(f"{'='*70}\n")
        
        # Run all strategies
        for strategy in strategies:
            self.run_backtest(strategy)
        
        # Compare results
        comparison_df = pd.DataFrame({
            name: {
                'Final Value': results['metrics']['final_value'],
                'Total Return': results['metrics']['total_return'],
                'CAGR': results['metrics']['cagr'],
                'Sharpe Ratio': results['metrics']['sharpe_ratio'],
                'Max Drawdown': results['metrics']['max_drawdown'],
                'Win Rate': results['metrics']['win_rate']
            }
            for name, results in self.results.items()
        }).T
        
        print("\nStrategy Comparison:")
        print("="*70)
        print(comparison_df.to_string())
        print("="*70)
        
        # Plot comparison
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        
        # Portfolio values
        for name, results in self.results.items():
            axes[0, 0].plot(results['portfolio_value'], label=name, linewidth=2)
        axes[0, 0].set_title('Portfolio Value Comparison', fontsize=12, fontweight='bold')
        axes[0, 0].set_ylabel('Portfolio Value ($)')
        axes[0, 0].legend(loc='best', fontsize=8)
        axes[0, 0].grid(True, alpha=0.3)
        
        # Returns
        comparison_df['Total Return'].plot(kind='bar', ax=axes[0, 1], color='steelblue')
        axes[0, 1].set_title('Total Return Comparison', fontsize=12, fontweight='bold')
        axes[0, 1].set_ylabel('Total Return')
        axes[0, 1].set_xticklabels(axes[0, 1].get_xticklabels(), rotation=45, ha='right')
        axes[0, 1].grid(True, alpha=0.3, axis='y')
        
        # Sharpe Ratio
        comparison_df['Sharpe Ratio'].plot(kind='bar', ax=axes[1, 0], color='green')
        axes[1, 0].set_title('Sharpe Ratio Comparison', fontsize=12, fontweight='bold')
        axes[1, 0].set_ylabel('Sharpe Ratio')
        axes[1, 0].set_xticklabels(axes[1, 0].get_xticklabels(), rotation=45, ha='right')
        axes[1, 0].grid(True, alpha=0.3, axis='y')
        
        # Max Drawdown
        comparison_df['Max Drawdown'].plot(kind='bar', ax=axes[1, 1], color='red')
        axes[1, 1].set_title('Max Drawdown Comparison', fontsize=12, fontweight='bold')
        axes[1, 1].set_ylabel('Max Drawdown')
        axes[1, 1].set_xticklabels(axes[1, 1].get_xticklabels(), rotation=45, ha='right')
        axes[1, 1].grid(True, alpha=0.3, axis='y')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Comparison plot saved to: {save_path}")
        
        plt.show()
        
        return comparison_df


if __name__ == "__main__":
    print("=" * 70)
    print("Stock Market Time Series Dataset - Backtesting Framework")
    print("=" * 70)
    print("\nAuthor: Molla Samser | RSK World")
    print("Website: https://rskworld.in/")
    print("=" * 70)
    
    # Initialize backtester
    backtester = Backtester('data/AAPL.csv', initial_capital=100000)
    
    # Define strategies
    strategies = [
        MovingAverageCrossover(short_window=20, long_window=50),
        RSIStrategy(oversold=30, overbought=70),
        MACDStrategy()
    ]
    
    # Compare strategies
    comparison = backtester.compare_strategies(strategies)
    
    # Plot individual strategy
    print("\nPlotting MA Crossover strategy details...")
    backtester.plot_results('MA_Crossover_20_50')
    
    print("\n" + "=" * 70)
    print("Visit https://rskworld.in/ for more datasets and tools!")
    print("=" * 70)