# Install necessary libraries
!pip install pandas numpy matplotlib seaborn arch yfinance --quiet

# Standard library imports
import logging
import collections
import datetime
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
import yfinance as yf
from arch import arch_model
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__)

sns.set_style('whitegrid')
plt.rcParams['figure.figsize'] = (10, 6)
plt.rcParams['figure.dpi'] = 100

def create_garch_state(initial_params: dict = None) -> dict:
    """
    Initializes the state dictionary for GARCH analysis.

    Parameters
    ----------
    initial_params : dict, optional
        A dictionary of initial parameters to override defaults.

    Returns
    -------
    dict
        An initialized state dictionary for GARCH analysis.
    """
    state = {
        'symbol': 'SPY', # Default stock symbol
        'start_date': (datetime.date.today() - datetime.timedelta(days=5*365)).strftime('%Y-%m-%d'), # 5 years ago
        'end_date': datetime.date.today().strftime('%Y-%m-%d'),
        'data': pd.DataFrame(),
        'returns': pd.Series(),
        'model_fit': None,
        'forecast_horizon': 5, # Days to forecast
        'forecasts': pd.DataFrame(),
        'eval_metrics': {},
        'log_messages': collections.deque(maxlen=100) # For logging important events
    }

    if initial_params:
        state.update(initial_params)

    logger.info(f"GARCH state initialized for symbol: {state['symbol']} from {state['start_date']} to {state['end_date']}")
    state['log_messages'].append(f"GARCH state initialized for symbol: {state['symbol']}")
    return state

def fetch_financial_data(state: dict, symbol: str, start_date: str, end_date: str) -> dict:
    """
    Fetches historical financial data for a given symbol and date range.

    Parameters
    ----------
    state : dict
        The current state dictionary.
    symbol : str
        The stock ticker symbol (e.g., 'SPY', 'AAPL').
    start_date : str
        The start date for data fetching in 'YYYY-MM-DD' format.
    end_date : str
        The end date for data fetching in 'YYYY-MM-DD' format.

    Returns
    -------
    dict
        The updated state dictionary with the fetched financial data stored under the 'data' key.
    """
    max_retries = 5
    base_delay = 1.0 # seconds

    for attempt in range(max_retries):
        try:
            logger.info(f"Attempting to fetch data for {symbol} from {start_date} to {end_date} (Attempt {attempt + 1})")
            data = yf.download(symbol, start=start_date, end=end_date)
            if data.empty:
                logger.warning(f"No data fetched for {symbol} in the specified range. Attempt {attempt + 1}.")
                raise ValueError("Empty DataFrame received.")

            state['data'] = data
            state['symbol'] = symbol
            state['start_date'] = start_date
            state['end_date'] = end_date
            logger.info(f"Successfully fetched {len(data)} rows for {symbol}.")
            state['log_messages'].append(f"Fetched {len(data)} rows for {symbol}.")
            return state

        except Exception as e:
            delay = (base_delay * (2 ** attempt)) + (random.random() * 0.5) # Exponential backoff with jitter
            logger.error(f"Failed to fetch data for {symbol}: {e}. Retrying in {delay:.2f} seconds...")
            state['log_messages'].append(f"Failed to fetch data for {symbol}: {e}. Retrying.")
            time.sleep(delay)

    logger.critical(f"Failed to fetch data for {symbol} after {max_retries} attempts.")
    state['log_messages'].append(f"Critical: Failed to fetch data for {symbol} after {max_retries} attempts.")
    return state # Return state even if data fetching failed

def preprocess_data(state: dict) -> dict:
    """
    Calculates daily log returns from the 'Adj Close' price and handles missing values.

    Parameters
    ----------
    state : dict
        The current state dictionary containing the raw 'data' DataFrame.

    Returns
    -------
    dict
        The updated state dictionary with 'returns' Series calculated and any NaN values handled.
    """
    if state['data'].empty:
        logger.warning("Data DataFrame is empty, cannot preprocess. Skipping.")
        state['log_messages'].append("Warning: Data DataFrame is empty, cannot preprocess.")
        return state

    logger.info("Calculating log returns and handling missing values...")
    # Access the 'Close' column using the symbol from the MultiIndex
    returns = 100 * state['data'][('Close', state['symbol'])].pct_change().dropna()
    state['returns'] = returns

    if returns.empty:
        logger.warning("Returns Series is empty after calculation and dropping NaNs.")
        state['log_messages'].append("Warning: Returns Series is empty after preprocessing.")
    else:
        logger.info(f"Successfully calculated {len(returns)} log returns. First 5: {returns.head().tolist()}")
        state['log_messages'].append(f"Calculated {len(returns)} returns.")

    return state

def fit_garch_model(state: dict, p: int = 1, q: int = 1, dist: str = 'normal') -> dict:
    """
    Fits a GARCH(p,q) model to the returns data.

    Parameters
    ----------
    state : dict
        The current state dictionary containing the 'returns' Series.
    p : int, optional
        The order of the ARCH component, defaults to 1.
    q : int, optional
        The order of the GARCH component, defaults to 1.
    dist : str, optional
        The error distribution for the model (e.g., 'normal', 't', 'skewt'), defaults to 'normal'.

    Returns
    -------
    dict
        The updated state dictionary with the fitted GARCH model stored under the 'model_fit' key.
    """
    if state['returns'].empty:
        logger.warning("Returns Series is empty, cannot fit GARCH model. Skipping.")
        state['log_messages'].append("Warning: Returns Series empty, cannot fit GARCH.")
        return state

    logger.info(f"Attempting to fit GARCH({p},{q}) model with {dist} distribution...")
    try:
        # Ensure returns are not all zero or constant, which can cause convergence issues
        if state['returns'].std() < 1e-6:
            logger.warning("Returns are constant or near zero, GARCH model might not converge.")
            state['log_messages'].append("Warning: Returns constant, GARCH might fail.")

        model = arch_model(state['returns'], vol='Garch', p=p, q=q, dist=dist)
        model_fit = model.fit(disp='off') # disp='off' suppresses verbose output
        state['model_fit'] = model_fit
        logger.info(f"GARCH({p},{q}) model fitted successfully. AIC: {model_fit.aic:.2f}")
        state['log_messages'].append(f"GARCH({p},{q}) model fitted. AIC: {model_fit.aic:.2f}")

    except Exception as e:
        logger.error(f"Error fitting GARCH({p},{q}) model: {e}")
        state['log_messages'].append(f"Error fitting GARCH({p},{q}) model: {e}")
        state['model_fit'] = None # Ensure model_fit is None if error occurs

    return state

def forecast_volatility(state: dict) -> dict:
    """
    Generates future volatility forecasts from a fitted GARCH model.

    Parameters
    ----------
    state : dict
        The current state dictionary containing the 'model_fit' and 'forecast_horizon'.

    Returns
    -------
    dict
        The updated state dictionary with 'forecasts' DataFrame containing forecasted conditional variances and annualized volatilities.
    """
    if state['model_fit'] is None:
        logger.warning("No fitted GARCH model found in state, cannot forecast. Skipping.")
        state['log_messages'].append("Warning: No fitted model, cannot forecast.")
        return state

    if not state['returns'].empty:
        last_date = state['returns'].index[-1]
    else:
        last_date = pd.to_datetime(state['end_date']) # Fallback if returns empty

    logger.info(f"Forecasting volatility for {state['forecast_horizon']} steps...")
    try:
        forecast_result = state['model_fit'].forecast(horizon=state['forecast_horizon'], start=last_date)
        mean_forecast = forecast_result.mean.iloc[-1]
        variance_forecast = forecast_result.variance.iloc[-1]
        h_vol_forecast = np.sqrt(variance_forecast)

        # Create a date range for the forecast
        forecast_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), periods=state['forecast_horizon'], freq='B') # Business days
        if len(forecast_dates) > len(h_vol_forecast):
          forecast_dates = forecast_dates[:len(h_vol_forecast)]

        forecast_df = pd.DataFrame({
            'mean_forecast': mean_forecast.values,
            'variance_forecast': variance_forecast.values,
            'h_vol_forecast': h_vol_forecast.values,
            'annualized_volatility': h_vol_forecast.values * np.sqrt(252) # Annualize daily volatility
        }, index=forecast_dates)

        state['forecasts'] = forecast_df
        logger.info(f"Successfully generated {state['forecast_horizon']} volatility forecasts.")
        state['log_messages'].append(f"Generated {state['forecast_horizon']} volatility forecasts.")

    except Exception as e:
        logger.error(f"Error generating volatility forecast: {e}")
        state['log_messages'].append(f"Error during forecasting: {e}")
        state['forecasts'] = pd.DataFrame() # Ensure empty if error

    return state

def evaluate_forecast(state: dict) -> dict:
    """
    Evaluates the performance of the GARCH volatility forecasts.

    Parameters
    ----------
    state : dict
        The current state dictionary containing 'returns', 'model_fit', and 'forecasts'.

    Returns
    -------
    dict
        The updated state dictionary with 'eval_metrics' containing the MSE.
    """
    if state['forecasts'].empty or state['model_fit'] is None or state['returns'].empty:
        logger.warning("Cannot evaluate forecast: missing forecasts, model, or returns.")
        state['log_messages'].append("Warning: Cannot evaluate forecast due to missing data.")
        state['eval_metrics'] = {'mse': np.nan}
        return state

    logger.info("Evaluating GARCH model forecast...")
    try:
        # Get the actual (realized) squared returns for comparison
        # We need to align the dates of forecasts and actuals
        actual_volatility = state['returns'][-len(state['forecasts']):].values**2 # Squared returns as proxy for realized variance
        predicted_variance = state['forecasts']['variance_forecast'].values

        if len(actual_volatility) == len(predicted_variance):
            mse = mean_squared_error(actual_volatility, predicted_variance)
            state['eval_metrics']['mse'] = mse
            logger.info(f"Forecast evaluation successful. MSE: {mse:.4f}")
            state['log_messages'].append(f"Forecast evaluated. MSE: {mse:.4f}")
        else:
            logger.warning("Mismatch in length between actual and predicted volatilities, cannot calculate MSE.")
            state['log_messages'].append("Warning: Length mismatch for forecast evaluation.")
            state['eval_metrics']['mse'] = np.nan

    except Exception as e:
        logger.error(f"Error during forecast evaluation: {e}")
        state['log_messages'].append(f"Error during forecast evaluation: {e}")
        state['eval_metrics']['mse'] = np.nan

    return state

# Initialize the GARCH state
initial_state = create_garch_state({
    'symbol': 'SPY',
    'start_date': (datetime.date.today() - datetime.timedelta(days=5*365)).strftime('%Y-%m-%d'),
    'end_date': datetime.date.today().strftime('%Y-%m-%d'),
    'forecast_horizon': 20 # Forecast 20 business days into the future
})

# Fetch financial data
state = fetch_financial_data(initial_state, initial_state['symbol'], initial_state['start_date'], initial_state['end_date'])

print("\n--- Raw Data Head ---")
display(state['data'].head())

print("\n--- Raw Data Tail ---")
display(state['data'].tail())

# Preprocess the data to get returns
state = preprocess_data(state)

print("\n--- Returns Series Head ---")
display(state['returns'].head())

print("\n--- Returns Series Info ---")
state['returns'].info()

# Plot the returns series
fig, ax = plt.subplots(figsize=(12, 6))
state['returns'].plot(ax=ax, title=f'{state["symbol"]} Daily Log Returns')
ax.set_xlabel('Date')
ax.set_ylabel('Log Returns (%)')
plt.tight_layout()
plt.show()

# Fit the GARCH model (GARCH(1,1) with normal distribution)
state = fit_garch_model(state, p=1, q=1, dist='normal')

if state['model_fit']:
    print("\n--- GARCH Model Summary ---")
    print(state['model_fit'].summary())

    # Plot conditional volatility
    fig, ax = plt.subplots(figsize=(12, 6))
    state['model_fit'].conditional_volatility.plot(ax=ax, title=f'{state["symbol"]} Conditional Volatility (GARCH(1,1))')
    ax.set_xlabel('Date')
    ax.set_ylabel('Conditional Volatility (%)')
    plt.tight_layout()
    plt.show()

    # Plot standardized residuals
    fig, axes = plt.subplots(1, 2, figsize=(15, 6))

    state['model_fit'].resid.plot(ax=axes[0], title='Standardized Residuals')
    axes[0].set_xlabel('Date')
    axes[0].set_ylabel('Residual')

    sns.histplot(state['model_fit'].std_resid, kde=True, ax=axes[1])
    axes[1].set_title('Histogram of Standardized Residuals')
    axes[1].set_xlabel('Standardized Residual')
    axes[1].set_ylabel('Density')

    plt.tight_layout()
    plt.show()
else:
    print("GARCH model could not be fitted.")

# Forecast volatility
state = forecast_volatility(state)

if not state['forecasts'].empty:
    print(f"\n--- Volatility Forecasts for next {state['forecast_horizon']} days ---")
    display(state['forecasts'])

    # Plot historical conditional volatility and future forecasts
    fig, ax = plt.subplots(figsize=(12, 7))

    # Historical conditional volatility
    if state['model_fit']:
        state['model_fit'].conditional_volatility.plot(ax=ax, label='Historical Conditional Volatility', color='blue')

    # Forecasted conditional volatility
    state['forecasts']['h_vol_forecast'].plot(ax=ax, label='Forecasted Volatility', color='red', linestyle='--')

    ax.set_title(f'{state["symbol"]} Historical and Forecasted Conditional Volatility')
    ax.set_xlabel('Date')
    ax.set_ylabel('Volatility (%)')
    ax.legend()
    plt.tight_layout()
    plt.show()

    # Plot annualized volatility forecast separately
    fig, ax = plt.subplots(figsize=(12, 6))
    state['forecasts']['annualized_volatility'].plot(ax=ax, color='green', marker='o', linestyle='-')
    ax.set_title(f'{state["symbol"]} Annualized Volatility Forecast (Next {state["forecast_horizon"]} Days)')
    ax.set_xlabel('Date')
    ax.set_ylabel('Annualized Volatility (%)')
    ax.grid(True, linestyle='--', alpha=0.6)
    plt.tight_layout()
    plt.show()

else:
    print("No volatility forecasts generated.")

# Evaluate the forecast
state = evaluate_forecast(state)

if 'mse' in state['eval_metrics'] and not np.isnan(state['eval_metrics']['mse']):
    print(f"\n--- Forecast Evaluation Metrics ---")
    display(pd.DataFrame([state['eval_metrics']]))
else:
    print("Forecast evaluation could not be performed or resulted in NaN.")

# Display recent log messages
print("\n--- Recent Log Messages ---")
for msg in state['log_messages']:
    print(msg)