"""
Time Series Analysis using Statsmodels

Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.tsa.stattools import adfuller, acf, pacf
from statsmodels.tsa.vector_ar.var_model import VAR
from statsmodels.tsa.holtwinters import ExponentialSmoothing
from statsmodels.stats.diagnostic import acorr_ljungbox
import warnings
warnings.filterwarnings('ignore')


class TimeSeriesModel:
    """
    Time Series Analysis and Forecasting
    
    Author: RSK World
    Website: https://rskworld.in
    Email: help@rskworld.in
    Phone: +91 93305 39277
    """
    
    def __init__(self):
        self.model = None
        self.results = None
        self.data = None
    
    def fit_arima(self, data, order=(1, 1, 1)):
        """
        Fit ARIMA model
        
        Parameters:
        -----------
        data : array-like or Series
            Time series data
        order : tuple
            (p, d, q) order for ARIMA model
        """
        self.data = pd.Series(data) if not isinstance(data, pd.Series) else data
        self.model = ARIMA(self.data, order=order)
        self.results = self.model.fit()
        return self.results
    
    def summary(self):
        """Print model summary"""
        if self.results is not None:
            print(self.results.summary())
        else:
            print("Model not fitted yet. Call fit_arima() first.")
    
    def forecast(self, steps=10, alpha=0.05):
        """
        Generate forecasts
        
        Parameters:
        -----------
        steps : int
            Number of steps ahead to forecast
        alpha : float
            Significance level for confidence intervals
        """
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit_arima() first.")
        
        forecast = self.results.forecast(steps=steps)
        conf_int = self.results.get_forecast(steps=steps).conf_int(alpha=alpha)
        
        return forecast, conf_int
    
    def plot_forecast(self, steps=10, alpha=0.05):
        """Plot original data and forecast"""
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit_arima() first.")
        
        forecast, conf_int = self.forecast(steps=steps, alpha=alpha)
        
        plt.figure(figsize=(12, 6))
        plt.plot(self.data.index, self.data.values, label='Original', linewidth=2)
        
        forecast_index = pd.date_range(start=self.data.index[-1], periods=steps+1, freq='D')[1:]
        plt.plot(forecast_index, forecast, label='Forecast', linewidth=2, color='red')
        plt.fill_between(forecast_index, conf_int.iloc[:, 0], conf_int.iloc[:, 1], 
                         alpha=0.3, color='red', label=f'{int((1-alpha)*100)}% Confidence Interval')
        
        plt.xlabel('Time')
        plt.ylabel('Value')
        plt.title('Time Series Forecast')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()
    
    def decompose(self, model='additive', period=None):
        """
        Decompose time series into trend, seasonal, and residual components
        
        Parameters:
        -----------
        model : str
            'additive' or 'multiplicative'
        period : int
            Period for seasonal decomposition
        """
        if self.data is None:
            raise ValueError("Data not set. Provide data first.")
        
        decomposition = seasonal_decompose(self.data, model=model, period=period)
        
        fig, axes = plt.subplots(4, 1, figsize=(12, 10))
        
        decomposition.observed.plot(ax=axes[0], title='Original', color='blue')
        decomposition.trend.plot(ax=axes[1], title='Trend', color='green')
        decomposition.seasonal.plot(ax=axes[2], title='Seasonal', color='orange')
        decomposition.resid.plot(ax=axes[3], title='Residual', color='red')
        
        for ax in axes:
            ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()
        
        return decomposition
    
    def test_stationarity(self):
        """
        Test for stationarity using Augmented Dickey-Fuller test
        
        Returns:
        --------
        dict : Test results
        """
        if self.data is None:
            raise ValueError("Data not set. Provide data first.")
        
        result = adfuller(self.data.dropna())
        
        print("Augmented Dickey-Fuller Test:")
        print(f"ADF Statistic: {result[0]:.6f}")
        print(f"p-value: {result[1]:.6f}")
        print(f"Critical Values:")
        for key, value in result[4].items():
            print(f"\t{key}: {value:.6f}")
        
        if result[1] <= 0.05:
            print("\nResult: Series is stationary (reject null hypothesis)")
        else:
            print("\nResult: Series is non-stationary (fail to reject null hypothesis)")
        
        return {
            'adf_statistic': result[0],
            'pvalue': result[1],
            'critical_values': result[4],
            'is_stationary': result[1] <= 0.05
        }
    
    def plot_acf_pacf(self, lags=40):
        """Plot ACF and PACF"""
        if self.data is None:
            raise ValueError("Data not set. Provide data first.")
        
        fig, axes = plt.subplots(2, 1, figsize=(12, 8))
        
        acf_values = acf(self.data.dropna(), nlags=lags)
        pacf_values = pacf(self.data.dropna(), nlags=lags)
        
        axes[0].stem(range(len(acf_values)), acf_values)
        axes[0].axhline(y=0, color='black', linestyle='-', linewidth=0.5)
        axes[0].axhline(y=1.96/np.sqrt(len(self.data)), color='r', linestyle='--', alpha=0.5)
        axes[0].axhline(y=-1.96/np.sqrt(len(self.data)), color='r', linestyle='--', alpha=0.5)
        axes[0].set_xlabel('Lag')
        axes[0].set_ylabel('ACF')
        axes[0].set_title('Autocorrelation Function (ACF)')
        axes[0].grid(True, alpha=0.3)
        
        axes[1].stem(range(len(pacf_values)), pacf_values)
        axes[1].axhline(y=0, color='black', linestyle='-', linewidth=0.5)
        axes[1].axhline(y=1.96/np.sqrt(len(self.data)), color='r', linestyle='--', alpha=0.5)
        axes[1].axhline(y=-1.96/np.sqrt(len(self.data)), color='r', linestyle='--', alpha=0.5)
        axes[1].set_xlabel('Lag')
        axes[1].set_ylabel('PACF')
        axes[1].set_title('Partial Autocorrelation Function (PACF)')
        axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()
    
    def fit_exponential_smoothing(self, data, trend=None, seasonal=None, seasonal_periods=None):
        """
        Fit Exponential Smoothing model
        
        Parameters:
        -----------
        data : array-like or Series
            Time series data
        trend : str
            'add', 'mul', or None
        seasonal : str
            'add', 'mul', or None
        seasonal_periods : int
            Number of periods in a season
        """
        self.data = pd.Series(data) if not isinstance(data, pd.Series) else data
        self.model = ExponentialSmoothing(self.data, trend=trend, 
                                         seasonal=seasonal, 
                                         seasonal_periods=seasonal_periods)
        self.results = self.model.fit()
        return self.results
    
    def check_residuals(self):
        """Check model residuals for autocorrelation"""
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit_arima() first.")
        
        residuals = self.results.resid
        lb_test = acorr_ljungbox(residuals, lags=10, return_df=True)
        
        print("Ljung-Box Test for Residual Autocorrelation:")
        print(lb_test)
        
        if lb_test['lb_pvalue'].iloc[-1] < 0.05:
            print("\nWarning: Residuals show autocorrelation (p < 0.05)")
        else:
            print("\nNo significant residual autocorrelation detected")
        
        return lb_test


class VARModel:
    """
    Vector Autoregression (VAR) Model
    
    Author: RSK World
    Website: https://rskworld.in
    Email: help@rskworld.in
    Phone: +91 93305 39277
    """
    
    def __init__(self):
        self.model = None
        self.results = None
        self.data = None
    
    def fit(self, data, maxlags=None):
        """
        Fit VAR model
        
        Parameters:
        -----------
        data : DataFrame
            Multivariate time series data
        maxlags : int
            Maximum number of lags to consider
        """
        self.data = data
        self.model = VAR(data)
        self.results = self.model.fit(maxlags=maxlags)
        return self.results
    
    def summary(self):
        """Print model summary"""
        if self.results is not None:
            print(self.results.summary())
        else:
            print("Model not fitted yet. Call fit() first.")
    
    def forecast(self, steps=10):
        """Generate forecasts"""
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit() first.")
        
        return self.results.forecast(self.data.values[-self.results.k_ar:], steps=steps)


if __name__ == "__main__":
    # Example usage
    print("Time Series Analysis Example")
    print("=" * 50)
    
    # Generate sample time series data
    np.random.seed(42)
    dates = pd.date_range('2020-01-01', periods=100, freq='D')
    trend = np.linspace(0, 10, 100)
    seasonal = 2 * np.sin(2 * np.pi * np.arange(100) / 7)
    noise = np.random.randn(100) * 0.5
    data = pd.Series(trend + seasonal + noise, index=dates)
    
    # Create and fit model
    ts_model = TimeSeriesModel()
    ts_model.data = data
    
    # Test stationarity
    ts_model.test_stationarity()
    
    # Decompose time series
    ts_model.decompose(period=7)
    
    # Plot ACF and PACF
    ts_model.plot_acf_pacf()
    
    # Fit ARIMA model
    ts_model.fit_arima(data, order=(1, 1, 1))
    ts_model.summary()
    
    # Forecast
    ts_model.plot_forecast(steps=20)