# Install necessary libraries
!pip install pandas numpy matplotlib seaborn statsmodels scikit-learn

# Standard library imports
import logging
from collections import deque
import time
import random

# Third-party library imports
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from statsmodels.tsa.stattools import adfuller
from sklearn.metrics import mean_squared_error

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Set seaborn style for better plots
sns.set_theme(style="whitegrid")

def create_adf_state(series_name: str) -> dict:
    """
    Initializes a state dictionary for managing ADF test related data.

    Parameters
    ----------
    series_name : str
        A descriptive name for the time series being analyzed.

    Returns
    -------
    dict
        An initialized state dictionary with keys for data, results, and plots.
    """
    logger.info(f"Initializing ADF state for series: {series_name}")
    state = {
        'series_name': series_name,
        'data': pd.Series(dtype=float),
        'adf_results': {},
        'plot_data': {}
    }
    logger.debug(f"ADF state initialized: {state}")
    return state

def simulate_timeseries_data(state: dict,
                             n_samples: int,
                             series_type: str,
                             ar_coeff: float = 0.7,
                             trend_coeff: float = 0.1,
                             season_period: int = 12,
                             season_strength: float = 5.0,
                             noise_std: float = 1.0,
                             random_seed: int = 42) -> dict:
    """
    Simulates various types of time series data (stationary, random walk, trend, seasonal).

    Parameters
    ----------
    state : dict
        The current state dictionary.
    n_samples : int
        The number of data points to generate.
    series_type : str
        The type of series to generate: 'stationary', 'random_walk', 'trend', 'seasonal'.
    ar_coeff : float, optional
        Auto-regressive coefficient for stationary series, defaults to 0.7.
    trend_coeff : float, optional
        Coefficient for linear trend, defaults to 0.1.
    season_period : int, optional
        Period for seasonality, defaults to 12.
    season_strength : float, optional
        Strength of the seasonal component, defaults to 5.0.
    noise_std : float, optional
        Standard deviation of the random noise, defaults to 1.0.
    random_seed : int, optional
        Seed for random number generation, defaults to 42.

    Returns
    -------
    dict
        The updated state dictionary with the simulated data stored in 'data' key.
    """
    logger.info(f"Simulating {series_type} time series with {n_samples} samples.")
    np.random.seed(random_seed)
    data = np.zeros(n_samples)
    noise = np.random.normal(0, noise_std, n_samples)

    if series_type == 'stationary':
        # AR(1) process: y_t = c + ar_coeff * y_{t-1} + e_t
        data[0] = noise[0]
        for i in range(1, n_samples):
            data[i] = ar_coeff * data[i-1] + noise[i]
        logger.debug("Generated stationary AR(1) series.")
    elif series_type == 'random_walk':
        # Random Walk: y_t = y_{t-1} + e_t
        data[0] = noise[0]
        for i in range(1, n_samples):
            data[i] = data[i-1] + noise[i]
        logger.debug("Generated random walk series.")
    elif series_type == 'trend':
        # Series with a linear trend: y_t = trend_coeff * t + e_t
        data = trend_coeff * np.arange(n_samples) + noise
        logger.debug("Generated series with linear trend.")
    elif series_type == 'seasonal':
        # Series with seasonality: y_t = A * sin(2*pi*t / period) + e_t
        for i in range(n_samples):
            data[i] = season_strength * np.sin(2 * np.pi * i / season_period) + noise[i]
        logger.debug("Generated series with seasonality.")
    else:
        logger.warning(f"Unknown series type: {series_type}. Generating default stationary series.")
        data[0] = noise[0]
        for i in range(1, n_samples):
            data[i] = ar_coeff * data[i-1] + noise[i]

    state['data'] = pd.Series(data, name=state['series_name'])
    logger.info(f"Simulation complete for {state['series_name']}. First 5 values: {state['data'].head().tolist()}")
    return state

def plot_timeseries(state: dict,
                    title: str,
                    xlabel: str = 'Time',
                    ylabel: str = 'Value',
                    fig_size: tuple = (12, 6)) -> dict:
    """
    Plots a time series from the state dictionary.

    Parameters
    ----------
    state : dict
        The current state dictionary containing the 'data' Series.
    title : str
        The title for the plot.
    xlabel : str, optional
        Label for the x-axis, defaults to 'Time'.
    ylabel : str, optional
        Label for the y-axis, defaults to 'Value'.
    fig_size : tuple, optional
        Dimensions of the plot figure, defaults to (12, 6).

    Returns
    -------
    dict
        The updated state dictionary.
    """
    logger.info(f"Generating plot for {state['series_name']} with title: {title}")
    plt.figure(figsize=fig_size)
    sns.lineplot(x=state['data'].index, y=state['data'].values, color='skyblue')
    plt.title(title, fontsize=16)
    plt.xlabel(xlabel, fontsize=12)
    plt.ylabel(ylabel, fontsize=12)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()
    logger.debug(f"Plot for {state['series_name']} displayed.")
    return state

def perform_adf_test(state: dict,
                     max_d_lags: int = None,
                     regression_type: str = 'c',
                     retries: int = 3,
                     backoff_factor: float = 0.5) -> dict:
    """
    Performs the Augmented Dickey-Fuller (ADF) test on the time series data in the state.

    Parameters
    ----------
    state : dict
        The current state dictionary containing the 'data' Series.
    max_d_lags : int, optional
        Maximum number of lags to use in the regression. If None, it is determined automatically.
        Defaults to None.
    regression_type : str, optional
        Type of regression to run: 'c' (constant), 'ct' (constant and trend),
        'ctt' (constant, trend, and quadratic trend), 'nc' (no constant, no trend).
        Defaults to 'c'.
    retries : int, optional
        Number of retry attempts for the ADF test, defaults to 3.
    backoff_factor : float, optional
        Factor by which to multiply the delay between retries, defaults to 0.5.

    Returns
    -------
    dict
        The updated state dictionary with ADF test results stored in 'adf_results'.
    """
    logger.info(f"Performing ADF test for {state['series_name']} with regression type '{regression_type}'.")

    ts_data = state['data'].dropna().values
    if len(ts_data) < 2:
        logger.warning(f"Insufficient data points ({len(ts_data)}) for ADF test on {state['series_name']}. Skipping.")
        state['adf_results'] = {'error': 'Insufficient data'}
        return state

    for i in range(retries):
        try:
            adf_output = adfuller(ts_data, maxlag=max_d_lags, regression=regression_type)
            adf_statistic = adf_output[0]
            p_value = adf_output[1]
            num_lags = adf_output[2]
            num_observations = adf_output[3]
            critical_values = adf_output[4]

            state['adf_results'] = {
                'adf_statistic': adf_statistic,
                'p_value': p_value,
                'num_lags': num_lags,
                'num_observations': num_observations,
                'critical_values': critical_values
            }
            logger.info(f"ADF test successful for {state['series_name']}. Test Statistic: {adf_statistic:.4f}, p-value: {p_value:.4f}")
            return state
        except Exception as e:
            delay = (backoff_factor * (2 ** i)) + random.uniform(0, 0.1) # Add random jitter
            logger.warning(f"ADF test failed for {state['series_name']} (Attempt {i+1}/{retries}): {e}. Retrying in {delay:.2f} seconds...")
            time.sleep(delay)

    logger.error(f"ADF test failed after {retries} attempts for {state['series_name']}.")
    state['adf_results'] = {'error': f'ADF test failed after {retries} attempts.'}
    return state

def summarize_adf_results(state: dict, significance_level: float = 0.05) -> dict:
    """
    Summarizes the ADF test results and determines stationarity.

    Parameters
    ----------
    state : dict
        The current state dictionary containing 'adf_results'.
    significance_level : float, optional
        The significance level to use for determining stationarity, defaults to 0.05.

    Returns
    -------
    dict
        The updated state dictionary.
    """
    logger.info(f"Summarizing ADF results for {state['series_name']}.")
    results = state.get('adf_results', {})

    if 'error' in results:
        logger.error(f"Cannot summarize ADF results due to error: {results['error']}")
        print(f"Error: {results['error']}")
        return state

    adf_statistic = results.get('adf_statistic')
    p_value = results.get('p_value')
    critical_values = results.get('critical_values', {})

    is_stationary_p_value = "Reject H0 (Stationary)" if p_value < significance_level else "Fail to Reject H0 (Non-Stationary)"

    # Check stationarity based on critical values (more robust for ADF)
    is_stationary_critical = ("Reject H0 (Stationary)" if (
        adf_statistic is not None and
        critical_values.get(f'{int(significance_level*100)}%', -np.inf) is not None and
        adf_statistic < critical_values.get(f'{int(significance_level*100)}%', -np.inf)
    ) else "Fail to Reject H0 (Non-Stationary)")

    summary_df = pd.DataFrame({
        'Metric': [
            'ADF Statistic',
            'p-value',
            'Number of Lags Used',
            'Number of Observations',
            '1% Critical Value',
            '5% Critical Value',
            '10% Critical Value',
            f'Stationarity (p < {significance_level})',
            f'Stationarity (ADF Stat < {int(significance_level*100)}% Crit Value)'
        ],
        'Value': [
            f'{adf_statistic:.4f}' if adf_statistic is not None else 'N/A',
            f'{p_value:.4f}' if p_value is not None else 'N/A',
            results.get('num_lags', 'N/A'),
            results.get('num_observations', 'N/A'),
            critical_values.get('1%', 'N/A'),
            critical_values.get('5%', 'N/A'),
            critical_values.get('10%', 'N/A'),
            is_stationary_p_value,
            is_stationary_critical
        ]
    })

    print(f"\n--- ADF Test Results Summary for {state['series_name']} ---")
    display(summary_df)
    logger.info(f"ADF results summary displayed for {state['series_name']}.")
    state['adf_summary_df'] = summary_df # Store for potential later use
    return state

# Initialize state for a stationary series
stationary_state = create_adf_state('Stationary AR(1) Series')

# Simulate stationary data
stationary_state = simulate_timeseries_data(stationary_state, n_samples=200, series_type='stationary', ar_coeff=0.8)

# Plot the stationary series
stationary_state = plot_timeseries(stationary_state, title='Simulated Stationary AR(1) Series')

# Perform ADF test on the stationary series
stationary_state = perform_adf_test(stationary_state)

# Summarize ADF test results
stationary_state = summarize_adf_results(stationary_state)

# Initialize state for a random walk series
random_walk_state = create_adf_state('Random Walk Series')

# Simulate random walk data
random_walk_state = simulate_timeseries_data(random_walk_state, n_samples=200, series_type='random_walk')

# Plot the random walk series
random_walk_state = plot_timeseries(random_walk_state, title='Simulated Random Walk Series')

# Perform ADF test on the random walk series
random_walk_state = perform_adf_test(random_walk_state)

# Summarize ADF test results
random_walk_state = summarize_adf_results(random_walk_state)

# Initialize state for a series with a trend
trend_state = create_adf_state('Series with Linear Trend')

# Simulate data with a linear trend
trend_state = simulate_timeseries_data(trend_state, n_samples=200, series_type='trend', trend_coeff=0.2)

# Plot the series with trend
trend_state = plot_timeseries(trend_state, title='Simulated Series with Linear Trend')

# Perform ADF test on the series with trend
trend_state = perform_adf_test(trend_state, regression_type='ct') # Use 'ct' for trend and constant

# Summarize ADF test results
trend_state = summarize_adf_results(trend_state)

# Initialize state for a seasonal series
seasonal_state = create_adf_state('Series with Seasonality')

# Simulate seasonal data
seasonal_state = simulate_timeseries_data(seasonal_state, n_samples=200, series_type='seasonal', season_period=24, season_strength=10.0)

# Plot the seasonal series
seasonal_state = plot_timeseries(seasonal_state, title='Simulated Series with Seasonality')

# Perform ADF test on the seasonal series (using 'c' as default, assuming trend is handled by constant)
seasonal_state = perform_adf_test(seasonal_state)

# Summarize ADF test results
seasonal_state = summarize_adf_results(seasonal_state)