"""
Stock Market Time Series Dataset - Visualization 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 List, Optional
import warnings
warnings.filterwarnings('ignore')


class StockVisualizer:
    """
    Advanced visualization tools for stock market data.
    """
    
    def __init__(self):
        """Initialize the StockVisualizer."""
        sns.set_style('whitegrid')
        self.colors = sns.color_palette("husl", 10)
        
    def plot_candlestick(self, data_path: str, last_n_days: int = 60, 
                        save_path: Optional[str] = None):
        """
        Create a candlestick chart.
        
        Args:
            data_path: Path to stock CSV file
            last_n_days: Number of recent days to plot
            save_path: Path to save the plot
        """
        df = pd.read_csv(data_path, parse_dates=['Date'], index_col='Date')
        symbol = Path(data_path).stem
        df = df.tail(last_n_days)
        
        fig, ax = plt.subplots(figsize=(14, 7))
        
        # Plot candlesticks
        for i, (date, row) in enumerate(df.iterrows()):
            color = 'green' if row['Close'] >= row['Open'] else 'red'
            
            # Draw the high-low line
            ax.plot([i, i], [row['Low'], row['High']], color=color, linewidth=1)
            
            # Draw the open-close rectangle
            height = abs(row['Close'] - row['Open'])
            bottom = min(row['Open'], row['Close'])
            ax.add_patch(plt.Rectangle((i - 0.3, bottom), 0.6, height, 
                                       facecolor=color, edgecolor=color, alpha=0.8))
        
        ax.set_xlim(-0.5, len(df) - 0.5)
        ax.set_title(f'{symbol} Candlestick Chart (Last {last_n_days} Days)', 
                    fontsize=16, fontweight='bold')
        ax.set_xlabel('Trading Days', fontsize=12)
        ax.set_ylabel('Price ($)', fontsize=12)
        ax.grid(True, alpha=0.3)
        
        # Set x-axis labels (show every 5th date)
        tick_positions = range(0, len(df), max(1, len(df) // 10))
        tick_labels = [df.index[i].strftime('%Y-%m-%d') for i in tick_positions]
        ax.set_xticks(tick_positions)
        ax.set_xticklabels(tick_labels, rotation=45)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Candlestick chart saved to: {save_path}")
        
        plt.show()
    
    def plot_multiple_stocks(self, stock_paths: List[str], 
                           normalize: bool = True,
                           save_path: Optional[str] = None):
        """
        Plot multiple stocks for comparison.
        
        Args:
            stock_paths: List of paths to stock CSV files
            normalize: Whether to normalize prices to 100
            save_path: Path to save the plot
        """
        plt.figure(figsize=(14, 7))
        
        for i, path in enumerate(stock_paths):
            df = pd.read_csv(path, parse_dates=['Date'], index_col='Date')
            symbol = Path(path).stem
            
            prices = df['Close']
            
            if normalize:
                prices = (prices / prices.iloc[0]) * 100
            
            plt.plot(df.index, prices, label=symbol, linewidth=2, 
                    color=self.colors[i % len(self.colors)])
        
        title = 'Stock Price Comparison'
        if normalize:
            title += ' (Normalized to 100)'
        
        plt.title(title, fontsize=16, fontweight='bold')
        plt.xlabel('Date', fontsize=12)
        plt.ylabel('Price ($)' if not normalize else 'Normalized Price', fontsize=12)
        plt.legend(loc='best', fontsize=10)
        plt.grid(True, alpha=0.3)
        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()
    
    def plot_correlation_matrix(self, stock_paths: List[str], 
                               save_path: Optional[str] = None):
        """
        Plot correlation matrix of multiple stocks.
        
        Args:
            stock_paths: List of paths to stock CSV files
            save_path: Path to save the plot
        """
        # Load and merge data
        dfs = {}
        for path in stock_paths:
            df = pd.read_csv(path, parse_dates=['Date'], index_col='Date')
            symbol = Path(path).stem
            dfs[symbol] = df['Close']
        
        combined = pd.DataFrame(dfs)
        correlation = combined.corr()
        
        plt.figure(figsize=(10, 8))
        sns.heatmap(correlation, annot=True, fmt='.3f', cmap='coolwarm', 
                   center=0, square=True, linewidths=1, cbar_kws={"shrink": 0.8})
        plt.title('Stock Price Correlation Matrix', fontsize=16, fontweight='bold')
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Correlation matrix saved to: {save_path}")
        
        plt.show()
    
    def plot_volume_analysis(self, data_path: str, save_path: Optional[str] = None):
        """
        Create volume analysis visualization.
        
        Args:
            data_path: Path to stock CSV file
            save_path: Path to save the plot
        """
        df = pd.read_csv(data_path, parse_dates=['Date'], index_col='Date')
        symbol = Path(data_path).stem
        
        fig, axes = plt.subplots(3, 1, figsize=(14, 10), 
                                gridspec_kw={'height_ratios': [2, 1, 1]})
        
        # Price with volume overlay
        ax1 = axes[0]
        ax2 = ax1.twinx()
        
        ax1.plot(df.index, df['Close'], color='blue', linewidth=2, label='Close Price')
        colors = ['green' if df['Close'].iloc[i] >= df['Open'].iloc[i] 
                 else 'red' for i in range(len(df))]
        ax2.bar(df.index, df['Volume'], color=colors, alpha=0.3, label='Volume')
        
        ax1.set_ylabel('Price ($)', fontsize=12, color='blue')
        ax2.set_ylabel('Volume', fontsize=12, color='gray')
        ax1.set_title(f'{symbol} Price and Volume Analysis', 
                     fontsize=16, fontweight='bold')
        ax1.grid(True, alpha=0.3)
        ax1.legend(loc='upper left')
        ax2.legend(loc='upper right')
        
        # Volume by day of week
        df['DayOfWeek'] = df.index.dayofweek
        day_names = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday']
        volume_by_day = df.groupby('DayOfWeek')['Volume'].mean()
        
        axes[1].bar(range(5), volume_by_day.values, color=self.colors[:5])
        axes[1].set_xticks(range(5))
        axes[1].set_xticklabels(day_names, rotation=45)
        axes[1].set_title('Average Volume by Day of Week', fontsize=12, fontweight='bold')
        axes[1].set_ylabel('Average Volume', fontsize=10)
        axes[1].grid(True, alpha=0.3, axis='y')
        
        # Volume distribution
        axes[2].hist(df['Volume'], bins=50, color='steelblue', alpha=0.7, edgecolor='black')
        axes[2].axvline(df['Volume'].mean(), color='red', linestyle='--', 
                       linewidth=2, label=f'Mean: {df["Volume"].mean():,.0f}')
        axes[2].set_title('Volume Distribution', fontsize=12, fontweight='bold')
        axes[2].set_xlabel('Volume', fontsize=10)
        axes[2].set_ylabel('Frequency', fontsize=10)
        axes[2].legend()
        axes[2].grid(True, alpha=0.3, axis='y')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Volume analysis saved to: {save_path}")
        
        plt.show()
    
    def plot_returns_heatmap(self, data_path: str, save_path: Optional[str] = None):
        """
        Create a heatmap of returns.
        
        Args:
            data_path: Path to stock CSV file
            save_path: Path to save the plot
        """
        df = pd.read_csv(data_path, parse_dates=['Date'], index_col='Date')
        symbol = Path(data_path).stem
        
        # Calculate daily returns
        df['Returns'] = df['Close'].pct_change() * 100
        
        # Create pivot table for heatmap
        df['Year'] = df.index.year
        df['Month'] = df.index.month
        df['Day'] = df.index.day
        
        # Monthly returns heatmap
        pivot = df.pivot_table(values='Returns', index='Month', 
                              columns='Year', aggfunc='sum')
        
        month_names = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 
                      'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
        
        plt.figure(figsize=(12, 8))
        sns.heatmap(pivot, annot=True, fmt='.1f', cmap='RdYlGn', center=0, 
                   cbar_kws={'label': 'Monthly Return (%)'}, linewidths=0.5)
        plt.title(f'{symbol} Monthly Returns Heatmap', fontsize=16, fontweight='bold')
        plt.xlabel('Year', fontsize=12)
        plt.ylabel('Month', fontsize=12)
        plt.yticks(range(12), month_names, rotation=0)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Returns heatmap saved to: {save_path}")
        
        plt.show()
    
    def create_dashboard(self, data_path: str, save_path: Optional[str] = None):
        """
        Create a comprehensive dashboard.
        
        Args:
            data_path: Path to stock CSV file
            save_path: Path to save the plot
        """
        df = pd.read_csv(data_path, parse_dates=['Date'], index_col='Date')
        symbol = Path(data_path).stem
        
        fig = plt.figure(figsize=(16, 12))
        gs = fig.add_gridspec(3, 2, hspace=0.3, wspace=0.3)
        
        # 1. Price History
        ax1 = fig.add_subplot(gs[0, :])
        ax1.plot(df.index, df['Close'], linewidth=2, color='blue')
        if 'MA_20' in df.columns:
            ax1.plot(df.index, df['MA_20'], linewidth=1.5, linestyle='--', 
                    label='MA(20)', alpha=0.7)
        if 'MA_50' in df.columns:
            ax1.plot(df.index, df['MA_50'], linewidth=1.5, linestyle='--', 
                    label='MA(50)', alpha=0.7)
        ax1.set_title(f'{symbol} Price History', fontsize=14, fontweight='bold')
        ax1.set_ylabel('Price ($)')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # 2. Volume
        ax2 = fig.add_subplot(gs[1, 0])
        colors = ['green' if df['Close'].iloc[i] >= df['Open'].iloc[i] 
                 else 'red' for i in range(len(df))]
        ax2.bar(df.index, df['Volume'], color=colors, alpha=0.5, width=1)
        ax2.set_title('Volume', fontsize=12, fontweight='bold')
        ax2.set_ylabel('Volume')
        ax2.grid(True, alpha=0.3)
        
        # 3. Returns Distribution
        ax3 = fig.add_subplot(gs[1, 1])
        returns = df['Close'].pct_change().dropna() * 100
        ax3.hist(returns, bins=50, alpha=0.7, color='steelblue', edgecolor='black')
        ax3.axvline(returns.mean(), color='red', linestyle='--', linewidth=2)
        ax3.set_title('Returns Distribution', fontsize=12, fontweight='bold')
        ax3.set_xlabel('Daily Returns (%)')
        ax3.set_ylabel('Frequency')
        ax3.grid(True, alpha=0.3)
        
        # 4. RSI
        ax4 = fig.add_subplot(gs[2, 0])
        if 'RSI' in df.columns:
            ax4.plot(df.index, df['RSI'], linewidth=2)
            ax4.axhline(70, color='red', linestyle='--', alpha=0.7)
            ax4.axhline(30, color='green', linestyle='--', alpha=0.7)
            ax4.fill_between(df.index, 30, 70, alpha=0.1)
        ax4.set_title('RSI', fontsize=12, fontweight='bold')
        ax4.set_ylabel('RSI')
        ax4.grid(True, alpha=0.3)
        
        # 5. MACD
        ax5 = fig.add_subplot(gs[2, 1])
        if 'MACD' in df.columns:
            ax5.plot(df.index, df['MACD'], linewidth=2)
            ax5.axhline(0, color='black', linestyle='-', alpha=0.5)
            ax5.fill_between(df.index, 0, df['MACD'], 
                           where=(df['MACD'] >= 0), alpha=0.3, color='green')
            ax5.fill_between(df.index, 0, df['MACD'], 
                           where=(df['MACD'] < 0), alpha=0.3, color='red')
        ax5.set_title('MACD', fontsize=12, fontweight='bold')
        ax5.set_ylabel('MACD')
        ax5.grid(True, alpha=0.3)
        
        fig.suptitle(f'{symbol} Stock Analysis Dashboard', 
                    fontsize=18, fontweight='bold', y=0.995)
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Dashboard saved to: {save_path}")
        
        plt.show()


if __name__ == "__main__":
    print("=" * 70)
    print("Stock Market Time Series Dataset - Visualization Tools")
    print("=" * 70)
    print("\nAuthor: Molla Samser | RSK World")
    print("Website: https://rskworld.in/")
    print("=" * 70)
    
    visualizer = StockVisualizer()
    
    # Candlestick chart
    print("\n1. Creating candlestick chart...")
    visualizer.plot_candlestick('data/AAPL.csv', last_n_days=60)
    
    # Multiple stocks comparison
    print("\n2. Comparing multiple stocks...")
    visualizer.plot_multiple_stocks(['data/AAPL.csv', 'data/GOOGL.csv', 
                                    'data/MSFT.csv'], normalize=True)
    
    # Correlation matrix
    print("\n3. Creating correlation matrix...")
    visualizer.plot_correlation_matrix(['data/AAPL.csv', 'data/GOOGL.csv', 
                                       'data/MSFT.csv', 'data/AMZN.csv', 'data/TSLA.csv'])
    
    # Volume analysis
    print("\n4. Volume analysis...")
    visualizer.plot_volume_analysis('data/AAPL.csv')
    
    # Dashboard
    print("\n5. Creating comprehensive dashboard...")
    visualizer.create_dashboard('data/AAPL.csv')
    
    print("\n" + "=" * 70)
    print("Visit https://rskworld.in/ for more datasets and tools!")
    print("=" * 70)