"""
Stock Market Time Series Dataset - Forecasting 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 StockForecaster:
    """
    Time series forecasting tools for stock market data.
    """
    
    def __init__(self, data_path: str):
        """
        Initialize the StockForecaster.
        
        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()
        
        # 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 simple_moving_average_forecast(self, days: int = 30, window: int = 20) -> pd.DataFrame:
        """
        Simple Moving Average forecast.
        
        Args:
            days: Number of days to forecast
            window: Moving average window
            
        Returns:
            DataFrame with forecasted values
        """
        last_date = self.df.index[-1]
        last_value = self.df['Close'].rolling(window=window).mean().iloc[-1]
        
        # Generate forecast dates
        forecast_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), 
                                       periods=days, freq='D')
        
        # Simple forecast: assume the MA continues
        forecast_values = [last_value] * days
        
        forecast_df = pd.DataFrame({
            'Date': forecast_dates,
            'Forecast': forecast_values,
            'Lower_Bound': [v * 0.95 for v in forecast_values],
            'Upper_Bound': [v * 1.05 for v in forecast_values]
        })
        forecast_df.set_index('Date', inplace=True)
        
        return forecast_df
    
    def exponential_smoothing_forecast(self, days: int = 30, alpha: float = 0.3) -> pd.DataFrame:
        """
        Exponential Smoothing forecast.
        
        Args:
            days: Number of days to forecast
            alpha: Smoothing parameter (0-1)
            
        Returns:
            DataFrame with forecasted values
        """
        # Calculate exponentially weighted moving average
        ewm = self.df['Close'].ewm(alpha=alpha, adjust=False).mean()
        last_value = ewm.iloc[-1]
        
        # Calculate trend
        recent_trend = (ewm.iloc[-1] - ewm.iloc[-20]) / 20
        
        last_date = self.df.index[-1]
        forecast_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), 
                                       periods=days, freq='D')
        
        # Generate forecast with trend
        forecast_values = [last_value + recent_trend * i for i in range(1, days + 1)]
        
        # Calculate confidence intervals
        std_dev = self.df['Close'].pct_change().std() * self.df['Close'].iloc[-1]
        
        forecast_df = pd.DataFrame({
            'Date': forecast_dates,
            'Forecast': forecast_values,
            'Lower_Bound': [v - 1.96 * std_dev * np.sqrt(i) 
                           for i, v in enumerate(forecast_values, 1)],
            'Upper_Bound': [v + 1.96 * std_dev * np.sqrt(i) 
                           for i, v in enumerate(forecast_values, 1)]
        })
        forecast_df.set_index('Date', inplace=True)
        
        return forecast_df
    
    def arima_forecast(self, days: int = 30) -> pd.DataFrame:
        """
        ARIMA forecast (simplified implementation).
        
        Args:
            days: Number of days to forecast
            
        Returns:
            DataFrame with forecasted values
        """
        print("\nNote: This is a simplified ARIMA implementation.")
        print("For production use, install statsmodels: pip install statsmodels")
        print("and use: from statsmodels.tsa.arima.model import ARIMA\n")
        
        # Simplified approach using moving average and trend
        prices = self.df['Close'].values
        
        # Calculate trend using linear regression
        x = np.arange(len(prices))
        z = np.polyfit(x[-60:], prices[-60:], 1)  # Use last 60 days
        trend_func = np.poly1d(z)
        
        # Generate forecast
        last_date = self.df.index[-1]
        forecast_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), 
                                       periods=days, freq='D')
        
        forecast_x = np.arange(len(prices), len(prices) + days)
        forecast_values = trend_func(forecast_x)
        
        # Calculate confidence intervals
        residuals = prices[-60:] - trend_func(x[-60:])
        std_error = np.std(residuals)
        
        forecast_df = pd.DataFrame({
            'Date': forecast_dates,
            'Forecast': forecast_values,
            'Lower_Bound': forecast_values - 1.96 * std_error,
            'Upper_Bound': forecast_values + 1.96 * std_error
        })
        forecast_df.set_index('Date', inplace=True)
        
        return forecast_df
    
    def plot_forecast(self, forecast_df: pd.DataFrame, method_name: str = "Forecast",
                     save_path: Optional[str] = None):
        """
        Plot historical data with forecast.
        
        Args:
            forecast_df: DataFrame with forecast data
            method_name: Name of forecasting method
            save_path: Path to save the plot
        """
        plt.figure(figsize=(14, 7))
        
        # Plot historical data
        plt.plot(self.df.index, self.df['Close'], label='Historical', 
                linewidth=2, color='blue')
        
        # Plot forecast
        plt.plot(forecast_df.index, forecast_df['Forecast'], 
                label=f'{method_name}', linewidth=2, color='red', linestyle='--')
        
        # Plot confidence interval
        plt.fill_between(forecast_df.index, 
                        forecast_df['Lower_Bound'], 
                        forecast_df['Upper_Bound'],
                        alpha=0.2, color='red', label='95% Confidence Interval')
        
        plt.title(f'{self.symbol} Stock Price Forecast - {method_name}', 
                 fontsize=16, fontweight='bold')
        plt.xlabel('Date', fontsize=12)
        plt.ylabel('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"Plot saved to: {save_path}")
        
        plt.show()
        
        # Print forecast summary
        print(f"\n{'='*60}")
        print(f"{method_name} Summary for {self.symbol}")
        print(f"{'='*60}")
        print(f"Forecast Period: {len(forecast_df)} days")
        print(f"Start Date: {forecast_df.index[0].strftime('%Y-%m-%d')}")
        print(f"End Date: {forecast_df.index[-1].strftime('%Y-%m-%d')}")
        print(f"Starting Price: ${forecast_df['Forecast'].iloc[0]:.2f}")
        print(f"Ending Price: ${forecast_df['Forecast'].iloc[-1]:.2f}")
        print(f"Expected Change: ${forecast_df['Forecast'].iloc[-1] - forecast_df['Forecast'].iloc[0]:.2f}")
        print(f"Expected Change %: {((forecast_df['Forecast'].iloc[-1] / forecast_df['Forecast'].iloc[0]) - 1) * 100:.2f}%")
        print(f"{'='*60}\n")
    
    def compare_forecasts(self, days: int = 30, save_path: Optional[str] = None):
        """
        Compare multiple forecasting methods.
        
        Args:
            days: Number of days to forecast
            save_path: Path to save the plot
        """
        # Generate forecasts
        sma_forecast = self.simple_moving_average_forecast(days)
        es_forecast = self.exponential_smoothing_forecast(days)
        arima_forecast = self.arima_forecast(days)
        
        plt.figure(figsize=(14, 7))
        
        # Plot historical data (last 60 days for clarity)
        historical_subset = self.df.tail(60)
        plt.plot(historical_subset.index, historical_subset['Close'], 
                label='Historical', linewidth=2, color='blue')
        
        # Plot forecasts
        plt.plot(sma_forecast.index, sma_forecast['Forecast'], 
                label='SMA Forecast', linewidth=2, linestyle='--')
        plt.plot(es_forecast.index, es_forecast['Forecast'], 
                label='Exp. Smoothing', linewidth=2, linestyle='--')
        plt.plot(arima_forecast.index, arima_forecast['Forecast'], 
                label='ARIMA Forecast', linewidth=2, linestyle='--')
        
        plt.title(f'{self.symbol} - Forecast Comparison', 
                 fontsize=16, fontweight='bold')
        plt.xlabel('Date', fontsize=12)
        plt.ylabel('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()
        
        # Print comparison
        print(f"\n{'='*60}")
        print(f"Forecast Comparison for {self.symbol}")
        print(f"{'='*60}")
        print(f"{'Method':<20} {'Final Price':<15} {'Change %':<15}")
        print(f"{'-'*60}")
        
        current_price = self.df['Close'].iloc[-1]
        
        for name, forecast in [('SMA', sma_forecast), 
                              ('Exp. Smoothing', es_forecast),
                              ('ARIMA', arima_forecast)]:
            final_price = forecast['Forecast'].iloc[-1]
            change_pct = ((final_price / current_price) - 1) * 100
            print(f"{name:<20} ${final_price:<14.2f} {change_pct:>13.2f}%")
        
        print(f"{'='*60}\n")


if __name__ == "__main__":
    print("=" * 70)
    print("Stock Market Time Series Dataset - Forecasting Tools")
    print("=" * 70)
    print("\nAuthor: Molla Samser | RSK World")
    print("Website: https://rskworld.in/")
    print("=" * 70)
    
    # Example forecasting
    forecaster = StockForecaster('data/AAPL.csv')
    
    # Simple Moving Average forecast
    print("\n1. Simple Moving Average Forecast...")
    sma_forecast = forecaster.simple_moving_average_forecast(days=30)
    forecaster.plot_forecast(sma_forecast, "Simple Moving Average")
    
    # Exponential Smoothing forecast
    print("\n2. Exponential Smoothing Forecast...")
    es_forecast = forecaster.exponential_smoothing_forecast(days=30)
    forecaster.plot_forecast(es_forecast, "Exponential Smoothing")
    
    # ARIMA forecast
    print("\n3. ARIMA Forecast...")
    arima_forecast = forecaster.arima_forecast(days=30)
    forecaster.plot_forecast(arima_forecast, "ARIMA")
    
    # Compare all methods
    print("\n4. Comparing Forecasting Methods...")
    forecaster.compare_forecasts(days=30)
    
    print("\n" + "=" * 70)
    print("Visit https://rskworld.in/ for more datasets and tools!")
    print("=" * 70)