import sys

# Install required libraries
!{sys.executable} -m pip install pandas numpy scikit-learn matplotlib seaborn shap loguru

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestRegressor
from sklearn.inspection import permutation_importance
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import shap
import logging
from collections import deque
import random
import time

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

def create_initial_state() -> dict:
    """
    Initializes an empty state dictionary.

    Returns
    -------
    dict
        An empty dictionary to manage the state of the notebook.

    Examples
    --------
    >>> state = create_initial_state()
    >>> print(state)
    {}
    """
    logger.info("Initializing new state dictionary.")
    return {}

def simulate_time_series_data(state: dict, n_samples: int = 1000, n_features: int = 5, seed: int = 42) -> dict:
    """
    Simulates time-series data with evolving feature contributions.

    Parameters
    ----------
    state : dict
        Current state dictionary.
    n_samples : int, optional
        Number of samples to generate, defaults to 1000.
    n_features : int, optional
        Number of features to generate, defaults to 5.
    seed : int, optional
        Random seed for reproducibility, defaults to 42.

    Returns
    -------
    dict
        Updated state with a 'data' key containing the generated DataFrame.

    Examples
    --------
    >>> state = create_initial_state()
    >>> state = simulate_time_series_data(state, n_samples=100, n_features=3, seed=1)
    >>> print(state['data'].head())
    """
    logger.info(f"Simulating time-series data with {n_samples} samples and {n_features} features.")
    np.random.seed(seed)
    random.seed(seed)

    # Generate time index
    time_index = pd.date_range(start='2020-01-01', periods=n_samples, freq='D')

    # Generate features
    features = pd.DataFrame(np.random.randn(n_samples, n_features), columns=[f'feature_{i}' for i in range(n_features)])

    # Introduce time-varying effects for some features
    # Feature_0 might have a strong, consistent effect
    # Feature_1 might have an effect that increases over time
    # Feature_2 might have an effect that decreases over time
    # Other features have random noise contribution

    weights = np.array([1.5] + [0.5] * (n_features - 1))

    # Introduce a time-dependent weight for feature_1 and feature_2
    time_effect_feature_1 = np.linspace(0.1, 1.0, n_samples) # increasing importance
    time_effect_feature_2 = np.linspace(1.0, 0.1, n_samples) # decreasing importance

    # Calculate target based on features and time-varying weights
    # A simple linear relationship for demonstration
    target = (1.5 * features['feature_0'] + # Strong constant contributor
              0.8 * features['feature_1'] * time_effect_feature_1 + # Increasing importance
              0.6 * features['feature_2'] * time_effect_feature_2 + # Decreasing importance
              np.sum(features.iloc[:, 3:] * weights[3:], axis=1) + # Other features (if any) with constant low weights
              np.random.randn(n_samples) * 0.5) # Noise

    df = pd.DataFrame({
        'timestamp': time_index,
        **{col: features[col] for col in features.columns},
        'target': target
    }).set_index('timestamp')

    state['data'] = df
    logger.info("Time-series data generated and stored in state['data'].")
    return state

def train_model(state: dict, X: pd.DataFrame, y: pd.Series, model_params: dict = None) -> dict:
    """
    Trains a RandomForestRegressor model.

    Parameters
    ----------
    state : dict
        Current state dictionary.
    X : pd.DataFrame
        Feature matrix for training.
    y : pd.Series
        Target variable for training.
    model_params : dict, optional
        Parameters for the RandomForestRegressor, defaults to {'random_state': 42}.

    Returns
    -------
    dict
        Updated state with the trained model under the 'model' key.

    Examples
    --------
    >>> state = create_initial_state()
    >>> state = simulate_time_series_data(state)
    >>> X = state['data'].drop('target', axis=1)
    >>> y = state['data']['target']
    >>> state = train_model(state, X, y)
    >>> print(state['model'])
    RandomForestRegressor(random_state=42)
    """
    if model_params is None:
        model_params = {'random_state': 42}

    logger.info(f"Training RandomForestRegressor with parameters: {model_params}")
    model = RandomForestRegressor(**model_params)
    model.fit(X, y)
    state['model'] = model
    logger.info("Model trained and stored in state['model'].")
    return state

def calculate_permutation_importance(state: dict, model, X: pd.DataFrame, y: pd.Series, n_repeats: int = 10) -> dict:
    """
    Calculates permutation importance for a given model and dataset.

    Parameters
    ----------
    state : dict
        Current state dictionary (model expected in state['model'] if not passed directly).
    model : estimator
        Trained machine learning model.
    X : pd.DataFrame
        Feature matrix.
    y : pd.Series
        True target values.
    n_repeats : int, optional
        Number of times to permute a feature, defaults to 10.

    Returns
    -------
    dict
        Updated state dictionary. This function directly returns the importance scores, not modifying state for generality.

    Examples
    --------
    >>> state = create_initial_state()
    >>> state = simulate_time_series_data(state, n_samples=100)
    >>> X = state['data'].drop('target', axis=1)
    >>> y = state['data']['target']
    >>> state = train_model(state, X, y)
    >>> model = state['model']
    >>> importance_results = calculate_permutation_importance(state, model, X, y)
    >>> print(importance_results['importances_mean'])
    """
    logger.info(f"Calculating permutation importance with {n_repeats} repeats.")
    result = permutation_importance(model, X, y, n_repeats=n_repeats, random_state=state.get('seed', 42), scoring='neg_mean_squared_error')

    # Store results for easier access
    state['permutation_importance_mean'] = pd.Series(result.importances_mean, index=X.columns)
    state['permutation_importance_std'] = pd.Series(result.importances_std, index=X.columns)

    logger.info("Permutation importance calculated.")
    return state

def track_rolling_feature_contributions(state: dict, data: pd.DataFrame, window_size: int = 200, step_size: int = 10, n_repeats: int = 5) -> dict:
    """
    Tracks feature contributions over time using a rolling window approach.

    Parameters
    ----------
    state : dict
        Current state dictionary.
    data : pd.DataFrame
        Time-series data with features and target.
    window_size : int, optional
        Size of the rolling window, defaults to 200.
    step_size : int, optional
        Step size to advance the window, defaults to 10.
    n_repeats : int, optional
        Number of repeats for permutation importance, defaults to 5.

    Returns
    -------
    dict
        Updated state with 'rolling_contributions' (DataFrame of feature importances over time)
        and 'rolling_metrics' (DataFrame of model evaluation metrics over time).

    Examples
    --------
    >>> state = create_initial_state()
    >>> state = simulate_time_series_data(state, n_samples=500)
    >>> state = track_rolling_feature_contributions(state, state['data'], window_size=100, step_size=20)
    >>> print(state['rolling_contributions'].head())
    >>> print(state['rolling_metrics'].head())
    """
    logger.info(f"Starting rolling window analysis with window_size={window_size} and step_size={step_size}.")

    feature_columns = [col for col in data.columns if col.startswith('feature_')]
    target_column = 'target'

    all_importances = []
    all_metrics = []
    window_timestamps = []

    # Use deque for an efficient rolling window
    # For this example, we'll manually slice the dataframe, but deque is good for streaming data.

    for i in range(0, len(data) - window_size + 1, step_size):
        window_data = data.iloc[i : i + window_size]

        if len(window_data) < window_size:
            continue # Skip incomplete windows at the end

        X_window = window_data[feature_columns]
        y_window = window_data[target_column]
        window_end_time = window_data.index[-1]
        window_timestamps.append(window_end_time)

        # Train model for the current window
        current_state = create_initial_state() # Create a temporary state for the model
        current_state = train_model(current_state, X_window, y_window)
        model = current_state['model']

        # Calculate permutation importance
        importance_state = create_initial_state()
        importance_state = calculate_permutation_importance(importance_state, model, X_window, y_window, n_repeats=n_repeats)
        importance_series = importance_state['permutation_importance_mean']
        all_importances.append(importance_series)

        # Evaluate model performance
        predictions = model.predict(X_window)
        rmse = np.sqrt(mean_squared_error(y_window, predictions))
        all_metrics.append({'rmse': rmse})

        # Add a small random jitter for demonstration of 'timing/backoff' mechanism concept
        time.sleep(random.uniform(0.01, 0.05))
        logger.debug(f"Processed window ending at {window_end_time}. RMSE: {rmse:.3f}")

    if not all_importances:
        logger.warning("No windows were processed. Check window_size and step_size.")
        state['rolling_contributions'] = pd.DataFrame()
        state['rolling_metrics'] = pd.DataFrame()
        return state

    rolling_contributions_df = pd.DataFrame(all_importances, index=window_timestamps)
    rolling_metrics_df = pd.DataFrame(all_metrics, index=window_timestamps)

    state['rolling_contributions'] = rolling_contributions_df
    state['rolling_metrics'] = rolling_metrics_df
    logger.info("Rolling feature contributions and metrics calculated and stored.")
    return state

# 1. Initialize State
state = create_initial_state()
state['seed'] = 42 # Set seed for reproducibility
logger.info("Initial state created.")

# 2. Simulate Time-Series Data
n_samples_sim = 1000
n_features_sim = 5
state = simulate_time_series_data(state, n_samples=n_samples_sim, n_features=n_features_sim, seed=state['seed'])
df = state['data']

print("\nSimulated Data Head:")
display(df.head())
print("\nSimulated Data Info:")
df.info()

X_full = df.drop('target', axis=1)
y_full = df['target']

# Train model on full data
state = train_model(state, X_full, y_full)
full_data_model = state['model']

# Calculate permutation importance on full data
state = calculate_permutation_importance(state, full_data_model, X_full, y_full)
initial_importance = state['permutation_importance_mean']

# Plot initial feature importance
plt.figure(figsize=(10, 6))
sns.barplot(x=initial_importance.index, y=initial_importance.values, palette='viridis')
plt.title('Initial Feature Importance (Full Dataset)')
plt.xlabel('Features')
plt.ylabel('Permutation Importance Mean')
plt.xticks(rotation=45, ha='right')
plt.tight_layout()
plt.show()

print("\nInitial Feature Importance:")
display(initial_importance.sort_values(ascending=False).to_frame(name='Importance Mean'))

window_s = 200 # Size of the training window
step_s = 20  # How many steps to advance the window
n_repeats_pi = 3 # Fewer repeats for faster demo

state = track_rolling_feature_contributions(state, df, window_size=window_s, step_size=step_s, n_repeats=n_repeats_pi)

rolling_contributions = state['rolling_contributions']
rolling_metrics = state['rolling_metrics']

print("\nRolling Feature Contributions Head:")
display(rolling_contributions.head())
print("\nRolling Metrics Head:")
display(rolling_metrics.head())

plt.figure(figsize=(15, 8))

# Plot each feature's importance over time
for feature in rolling_contributions.columns:
    plt.plot(rolling_contributions.index, rolling_contributions[feature], label=feature)

plt.title('Evolving Feature Contributions Over Time (Permutation Importance)')
plt.xlabel('Window End Timestamp')
plt.ylabel('Permutation Importance Mean')
plt.grid(True, linestyle='--', alpha=0.7)
plt.legend(title='Features', bbox_to_anchor=(1.05, 1), loc='upper left')
plt.xticks(rotation=45, ha='right')
plt.tight_layout()
plt.show()

plt.figure(figsize=(16, 8))
sns.heatmap(rolling_contributions.T, cmap='viridis', annot=False, fmt=".2f", cbar_kws={'label': 'Permutation Importance Mean'})
plt.title('Heatmap of Rolling Feature Contributions')
plt.xlabel('Window End Timestamp')
plt.ylabel('Features')
plt.xticks(rotation=45, ha='right')
plt.tight_layout()
plt.show()

plt.figure(figsize=(15, 6))
sns.lineplot(x=rolling_metrics.index, y=rolling_metrics['rmse'])
plt.title('Model RMSE Over Time (Rolling Window)')
plt.xlabel('Window End Timestamp')
plt.ylabel('RMSE')
plt.grid(True, linestyle='--', alpha=0.7)
plt.xticks(rotation=45, ha='right')
plt.tight_layout()
plt.show()