"""
Advanced Time Series Models (SARIMA, GARCH, etc.)

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.statespace.sarimax import SARIMAX
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.stattools import adfuller, kpss
from statsmodels.stats.diagnostic import acorr_ljungbox
import warnings
warnings.filterwarnings('ignore')


class SARIMAModel:
    """
    Seasonal ARIMA 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, order=(1, 1, 1), seasonal_order=(1, 1, 1, 12)):
        """
        Fit SARIMA model
        
        Parameters:
        -----------
        data : array-like or Series
            Time series data
        order : tuple
            (p, d, q) order
        seasonal_order : tuple
            (P, D, Q, s) seasonal order
        """
        self.data = pd.Series(data) if not isinstance(data, pd.Series) else data
        self.model = SARIMAX(self.data, order=order, seasonal_order=seasonal_order)
        self.results = self.model.fit(disp=False)
        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, alpha=0.05):
        """Generate forecasts"""
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit() 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 forecast"""
        if self.results is None:
            raise ValueError("Model not fitted yet. Call fit() 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)}% CI')
        
        plt.xlabel('Time')
        plt.ylabel('Value')
        plt.title('SARIMA Forecast')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()


class AutoARIMA:
    """
    Automatic ARIMA Model Selection
    
    Author: RSK World
    Website: https://rskworld.in
    Email: help@rskworld.in
    Phone: +91 93305 39277
    """
    
    def __init__(self):
        self.best_model = None
        self.best_order = None
        self.best_seasonal = None
        self.best_aic = float('inf')
    
    def auto_select(self, data, max_p=3, max_d=2, max_q=3, 
                    seasonal=False, max_P=2, max_D=1, max_Q=2, s=12):
        """
        Automatically select best ARIMA/SARIMA order
        
        Parameters:
        -----------
        data : array-like or Series
            Time series data
        max_p, max_d, max_q : int
            Maximum orders for ARIMA
        seasonal : bool
            Whether to include seasonal component
        max_P, max_D, max_Q : int
            Maximum seasonal orders
        s : int
            Seasonal period
        """
        data = pd.Series(data) if not isinstance(data, pd.Series) else data
        best_results = None
        
        print("Searching for best model...")
        print("=" * 70)
        
        for p in range(max_p + 1):
            for d in range(max_d + 1):
                for q in range(max_q + 1):
                    try:
                        if seasonal:
                            for P in range(max_P + 1):
                                for D in range(max_D + 1):
                                    for Q in range(max_Q + 1):
                                        try:
                                            model = SARIMAX(data, 
                                                          order=(p, d, q),
                                                          seasonal_order=(P, D, Q, s))
                                            results = model.fit(disp=False)
                                            
                                            if results.aic < self.best_aic:
                                                self.best_aic = results.aic
                                                self.best_order = (p, d, q)
                                                self.best_seasonal = (P, D, Q, s)
                                                self.best_model = results
                                                best_results = results
                                        except:
                                            continue
                        else:
                            try:
                                model = ARIMA(data, order=(p, d, q))
                                results = model.fit(disp=False)
                                
                                if results.aic < self.best_aic:
                                    self.best_aic = results.aic
                                    self.best_order = (p, d, q)
                                    self.best_model = results
                                    best_results = results
                            except:
                                continue
                    except:
                        continue
        
        if best_results is not None:
            print(f"Best model order: {self.best_order}")
            if seasonal:
                print(f"Best seasonal order: {self.best_seasonal}")
            print(f"Best AIC: {self.best_aic:.4f}")
            print("\nModel Summary:")
            print(best_results.summary())
        else:
            print("No suitable model found.")
        
        return best_results
    
    def forecast(self, steps=10):
        """Generate forecast from best model"""
        if self.best_model is None:
            raise ValueError("No model selected. Call auto_select() first.")
        
        return self.best_model.forecast(steps=steps)


class StationarityTests:
    """
    Comprehensive Stationarity Testing
    
    Author: RSK World
    Website: https://rskworld.in
    Email: help@rskworld.in
    Phone: +91 93305 39277
    """
    
    @staticmethod
    def adf_test(data, maxlag=None):
        """Augmented Dickey-Fuller test"""
        result = adfuller(data.dropna(), maxlag=maxlag)
        
        print("Augmented Dickey-Fuller Test:")
        print(f"ADF Statistic: {result[0]:.6f}")
        print(f"p-value: {result[1]:.6f}")
        print("Critical Values:")
        for key, value in result[4].items():
            print(f"  {key}: {value:.6f}")
        
        if result[1] <= 0.05:
            print("Result: Series is stationary (reject null hypothesis)")
        else:
            print("Result: Series is non-stationary (fail to reject null hypothesis)")
        
        return result
    
    @staticmethod
    def kpss_test(data, regression='c'):
        """KPSS test for stationarity"""
        result = kpss(data.dropna(), regression=regression)
        
        print("KPSS Test:")
        print(f"KPSS Statistic: {result[0]:.6f}")
        print(f"p-value: {result[1]:.6f}")
        print("Critical Values:")
        for key, value in result[3].items():
            print(f"  {key}: {value:.6f}")
        
        if result[1] < 0.05:
            print("Result: Series is non-stationary (reject null hypothesis)")
        else:
            print("Result: Series is stationary (fail to reject null hypothesis)")
        
        return result
    
    @staticmethod
    def comprehensive_test(data):
        """Run both ADF and KPSS tests"""
        print("=" * 70)
        print("Comprehensive Stationarity Tests")
        print("=" * 70)
        
        adf_result = StationarityTests.adf_test(data)
        print("\n" + "-" * 70)
        kpss_result = StationarityTests.kpss_test(data)
        
        return {
            'adf': adf_result,
            'kpss': kpss_result
        }


if __name__ == "__main__":
    # Example usage
    print("Advanced Time Series Models Example")
    print("=" * 70)
    
    # Generate sample seasonal data
    np.random.seed(42)
    dates = pd.date_range('2020-01-01', periods=200, freq='D')
    trend = np.linspace(0, 10, 200)
    seasonal = 5 * np.sin(2 * np.pi * np.arange(200) / 30)  # Monthly seasonality
    noise = np.random.randn(200) * 0.5
    data = pd.Series(trend + seasonal + noise, index=dates)
    
    # SARIMA model
    print("\nSARIMA Model:")
    print("=" * 70)
    sarima = SARIMAModel()
    sarima.fit(data, order=(1, 1, 1), seasonal_order=(1, 1, 1, 30))
    sarima.summary()
    
    # Auto ARIMA
    print("\n" + "=" * 70)
    print("Auto ARIMA Selection:")
    print("=" * 70)
    auto_arima = AutoARIMA()
    auto_arima.auto_select(data, max_p=2, max_d=1, max_q=2, seasonal=False)
    
    # Stationarity tests
    print("\n" + "=" * 70)
    print("Stationarity Tests:")
    print("=" * 70)
    StationarityTests.comprehensive_test(data)