"""
Stock Market Time Series Dataset - Analysis Tools

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 Optional, Tuple
import warnings
warnings.filterwarnings('ignore')


class StockAnalyzer:
    """
    Comprehensive analysis tools for stock market time series data.
    """
    
    def __init__(self, data_path: str):
        """
        Initialize the StockAnalyzer.
        
        Args:
            data_path: Path to the stock CSV file
        """
        self.data_path = Path(data_path)
        self.symbol = self.data_path.stem
        self.df = self._load_data()
        self.returns = None
        
        # Set style
        sns.set_style('whitegrid')
        plt.rcParams['figure.figsize'] = (12, 6)
        
    def _load_data(self) -> pd.DataFrame:
        """Load and prepare the data."""
        df = pd.read_csv(self.data_path, parse_dates=['Date'], index_col='Date')
        df.sort_index(inplace=True)
        return df
    
    def summary_statistics(self) -> pd.DataFrame:
        """
        Calculate comprehensive summary statistics.
        
        Returns:
            DataFrame with summary statistics
        """
        stats = {
            'Metric': [],
            'Value': []
        }
        
        # Price statistics
        stats['Metric'].extend(['Average Close', 'Median Close', 'Std Dev Close',
                               'Min Close', 'Max Close', 'Price Range'])
        stats['Value'].extend([
            f"${self.df['Close'].mean():.2f}",
            f"${self.df['Close'].median():.2f}",
            f"${self.df['Close'].std():.2f}",
            f"${self.df['Close'].min():.2f}",
            f"${self.df['Close'].max():.2f}",
            f"${self.df['Close'].max() - self.df['Close'].min():.2f}"
        ])
        
        # Returns
        returns = self.df['Close'].pct_change()
        stats['Metric'].extend(['Average Daily Return', 'Daily Volatility',
                               'Annualized Return', 'Annualized Volatility'])
        stats['Value'].extend([
            f"{returns.mean():.4f} ({returns.mean()*100:.2f}%)",
            f"{returns.std():.4f} ({returns.std()*100:.2f}%)",
            f"{returns.mean() * 252:.4f} ({returns.mean() * 252 * 100:.2f}%)",
            f"{returns.std() * np.sqrt(252):.4f} ({returns.std() * np.sqrt(252) * 100:.2f}%)"
        ])
        
        # Risk metrics
        sharpe_ratio = (returns.mean() / returns.std()) * np.sqrt(252)
        stats['Metric'].extend(['Sharpe Ratio', 'Max Drawdown'])
        stats['Value'].extend([
            f"{sharpe_ratio:.4f}",
            f"{self._calculate_max_drawdown():.2f}%"
        ])
        
        # Volume
        stats['Metric'].extend(['Average Volume', 'Total Volume'])
        stats['Value'].extend([
            f"{self.df['Volume'].mean():,.0f}",
            f"{self.df['Volume'].sum():,.0f}"
        ])
        
        df_stats = pd.DataFrame(stats)
        print(f"\n{'='*60}")
        print(f"Summary Statistics for {self.symbol}")
        print(f"{'='*60}")
        print(df_stats.to_string(index=False))
        print(f"{'='*60}\n")
        
        return df_stats
    
    def _calculate_max_drawdown(self) -> float:
        """Calculate maximum drawdown."""
        cumulative = (1 + self.df['Close'].pct_change()).cumprod()
        running_max = cumulative.expanding().max()
        drawdown = (cumulative - running_max) / running_max
        return drawdown.min() * 100
    
    def plot_price_history(self, save_path: Optional[str] = None):
        """
        Plot price history with volume.
        
        Args:
            save_path: Path to save the plot
        """
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 8), 
                                       gridspec_kw={'height_ratios': [3, 1]})
        
        # Plot closing price
        ax1.plot(self.df.index, self.df['Close'], label='Close', linewidth=2)
        ax1.fill_between(self.df.index, self.df['Low'], self.df['High'], 
                        alpha=0.2, label='High-Low Range')
        
        # Add moving averages
        if 'MA_20' in self.df.columns:
            ax1.plot(self.df.index, self.df['MA_20'], label='MA(20)', 
                    linewidth=1.5, linestyle='--', alpha=0.8)
        if 'MA_50' in self.df.columns:
            ax1.plot(self.df.index, self.df['MA_50'], label='MA(50)', 
                    linewidth=1.5, linestyle='--', alpha=0.8)
        
        ax1.set_title(f'{self.symbol} Stock Price History', fontsize=16, fontweight='bold')
        ax1.set_ylabel('Price ($)', fontsize=12)
        ax1.legend(loc='best')
        ax1.grid(True, alpha=0.3)
        
        # Plot volume
        colors = ['red' if self.df['Close'].iloc[i] < self.df['Open'].iloc[i] 
                 else 'green' for i in range(len(self.df))]
        ax2.bar(self.df.index, self.df['Volume'], color=colors, alpha=0.5)
        ax2.set_ylabel('Volume', fontsize=12)
        ax2.set_xlabel('Date', fontsize=12)
        ax2.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Plot saved to: {save_path}")
        
        plt.show()
    
    def plot_returns_distribution(self, save_path: Optional[str] = None):
        """
        Plot returns distribution.
        
        Args:
            save_path: Path to save the plot
        """
        returns = self.df['Close'].pct_change().dropna()
        
        fig, axes = plt.subplots(1, 2, figsize=(14, 5))
        
        # Histogram
        axes[0].hist(returns, bins=50, alpha=0.7, color='blue', edgecolor='black')
        axes[0].axvline(returns.mean(), color='red', linestyle='--', 
                       linewidth=2, label=f'Mean: {returns.mean():.4f}')
        axes[0].axvline(returns.median(), color='green', linestyle='--', 
                       linewidth=2, label=f'Median: {returns.median():.4f}')
        axes[0].set_title(f'{self.symbol} Returns Distribution', fontsize=14, fontweight='bold')
        axes[0].set_xlabel('Daily Returns', fontsize=12)
        axes[0].set_ylabel('Frequency', fontsize=12)
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)
        
        # Q-Q plot
        from scipy import stats
        stats.probplot(returns, dist="norm", plot=axes[1])
        axes[1].set_title('Q-Q Plot', fontsize=14, fontweight='bold')
        axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Plot saved to: {save_path}")
        
        plt.show()
    
    def plot_technical_indicators(self, save_path: Optional[str] = None):
        """
        Plot technical indicators.
        
        Args:
            save_path: Path to save the plot
        """
        fig, axes = plt.subplots(2, 1, figsize=(14, 10))
        
        # RSI
        if 'RSI' in self.df.columns:
            axes[0].plot(self.df.index, self.df['RSI'], label='RSI', linewidth=2)
            axes[0].axhline(70, color='red', linestyle='--', alpha=0.7, label='Overbought (70)')
            axes[0].axhline(30, color='green', linestyle='--', alpha=0.7, label='Oversold (30)')
            axes[0].fill_between(self.df.index, 30, 70, alpha=0.1)
            axes[0].set_title(f'{self.symbol} - Relative Strength Index (RSI)', 
                            fontsize=14, fontweight='bold')
            axes[0].set_ylabel('RSI', fontsize=12)
            axes[0].legend(loc='best')
            axes[0].grid(True, alpha=0.3)
        
        # MACD
        if 'MACD' in self.df.columns:
            axes[1].plot(self.df.index, self.df['MACD'], label='MACD', linewidth=2)
            axes[1].axhline(0, color='black', linestyle='-', alpha=0.5)
            axes[1].fill_between(self.df.index, 0, self.df['MACD'], 
                               where=(self.df['MACD'] >= 0), alpha=0.3, 
                               color='green', label='Positive')
            axes[1].fill_between(self.df.index, 0, self.df['MACD'], 
                               where=(self.df['MACD'] < 0), alpha=0.3, 
                               color='red', label='Negative')
            axes[1].set_title(f'{self.symbol} - MACD (Moving Average Convergence Divergence)', 
                            fontsize=14, fontweight='bold')
            axes[1].set_ylabel('MACD', fontsize=12)
            axes[1].set_xlabel('Date', fontsize=12)
            axes[1].legend(loc='best')
            axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Plot saved to: {save_path}")
        
        plt.show()
    
    def correlation_analysis(self, other_stock_path: str) -> float:
        """
        Perform correlation analysis with another stock.
        
        Args:
            other_stock_path: Path to the other stock CSV file
            
        Returns:
            Correlation coefficient
        """
        other_df = pd.read_csv(other_stock_path, parse_dates=['Date'], index_col='Date')
        other_symbol = Path(other_stock_path).stem
        
        # Align the data
        combined = pd.DataFrame({
            self.symbol: self.df['Close'],
            other_symbol: other_df['Close']
        }).dropna()
        
        correlation = combined.corr().iloc[0, 1]
        
        print(f"\nCorrelation between {self.symbol} and {other_symbol}: {correlation:.4f}")
        
        # Scatter plot
        plt.figure(figsize=(10, 6))
        plt.scatter(combined[self.symbol], combined[other_symbol], alpha=0.5)
        plt.xlabel(f'{self.symbol} Close Price ($)', fontsize=12)
        plt.ylabel(f'{other_symbol} Close Price ($)', fontsize=12)
        plt.title(f'Price Correlation: {self.symbol} vs {other_symbol}\\nCorrelation: {correlation:.4f}', 
                 fontsize=14, fontweight='bold')
        plt.grid(True, alpha=0.3)
        
        # Add trend line
        z = np.polyfit(combined[self.symbol], combined[other_symbol], 1)
        p = np.poly1d(z)
        plt.plot(combined[self.symbol], p(combined[self.symbol]), 
                "r--", alpha=0.8, linewidth=2, label='Trend Line')
        plt.legend()
        plt.tight_layout()
        plt.show()
        
        return correlation


if __name__ == "__main__":
    print("=" * 70)
    print("Stock Market Time Series Dataset - Analysis Tools")
    print("=" * 70)
    print("\nAuthor: Molla Samser | RSK World")
    print("Website: https://rskworld.in/")
    print("=" * 70)
    
    # Example analysis
    analyzer = StockAnalyzer('data/AAPL.csv')
    
    # Summary statistics
    analyzer.summary_statistics()
    
    # Plot price history
    print("\nGenerating price history plot...")
    analyzer.plot_price_history()
    
    # Plot returns distribution
    print("\nGenerating returns distribution plot...")
    analyzer.plot_returns_distribution()
    
    # Plot technical indicators
    print("\nGenerating technical indicators plot...")
    analyzer.plot_technical_indicators()
    
    print("\n" + "=" * 70)
    print("Visit https://rskworld.in/ for more datasets and tools!")
    print("=" * 70)