# Install necessary libraries
!pip install -qq numpy pandas scipy scikit-learn matplotlib seaborn evidently

# Standard Library Imports
import logging
import random
import time
from collections import deque

# Third-party Imports
import numpy as np
import pandas as pd
import scipy.stats as stats
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score
import matplotlib.pyplot as plt
import seaborn as sns

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

def create_initial_state() -> dict:
    """
    Initializes the application state dictionary with default parameters.

    Returns
    -------
    dict
        An initialized state dictionary containing simulation parameters,
        drift thresholds, and placeholders for data and model.
    """
    logger.info("Initializing application state.")
    state = {
        'n_samples': 1000, # Number of samples for initial data
        'n_features': 5, # Number of features
        'n_categorical': 2, # Number of categorical features
        'random_state': 42,
        'drift_magnitude': 0.1, # How much to shift for data drift simulation
        'model_performance_threshold': 0.85, # F1-score threshold for model drift
        'ks_alpha': 0.05, # Significance level for KS-test
        'chi2_alpha': 0.05, # Significance level for Chi-squared test
        'reference_data': None,
        'production_data_window': deque(maxlen=500), # Rolling window for production data
        'trained_model': None,
        'feature_names': [f'feature_{i}' for i in range(5)],
        'categorical_feature_names': [f'feature_{i}' for i in range(3, 5)],
        'numerical_feature_names': [f'feature_{i}' for i in range(3)],
        'target_name': 'target'
    }
    logger.debug(f"Initial state created: {state}")
    return state

def simulate_data(state: dict, n_samples: int, drift_type: str = 'none', drift_magnitude: float = 0.0) -> pd.DataFrame:
    """
    Simulates synthetic tabular data with optional drift.

    Parameters
    ----------
    state : dict
        The current application state dictionary.
    n_samples : int
        The number of samples to generate.
    drift_type : str, optional
        Type of drift to simulate. Options: 'none', 'data_drift_numerical',
        'data_drift_categorical', 'model_drift'. Defaults to 'none'.
    drift_magnitude : float, optional
        The magnitude of drift to apply, e.g., standard deviation shift.
        Defaults to 0.0.

    Returns
    -------
    pd.DataFrame
        A pandas DataFrame containing simulated features and a target variable.
    """
    np.random.seed(state['random_state'] + int(time.time() * 1000) % 1000) # Jitter for reproducibility
    logger.info(f"Simulating data with {n_samples} samples and drift type: {drift_type}")

    data = {}
    # Simulate numerical features
    for i in range(state['n_features'] - state['n_categorical']):
        mean_val = 0.0
        std_val = 1.0
        if drift_type == 'data_drift_numerical':
            mean_val += drift_magnitude # Shift mean for drift
        data[f'feature_{i}'] = np.random.normal(loc=mean_val, scale=std_val, size=n_samples)

    # Simulate categorical features
    for i in range(state['n_features'] - state['n_categorical'], state['n_features']):
        categories = [0, 1, 2]
        probabilities = [0.4, 0.3, 0.3]
        if drift_type == 'data_drift_categorical':
            # Change probabilities for drift
            probabilities = [0.2, 0.5, 0.3] # Example shift

        data[f'feature_{i}'] = np.random.choice(categories, size=n_samples, p=probabilities)

    df = pd.DataFrame(data)

    # Simulate target variable (e.g., binary classification)
    # Target depends on features, creating a linear relationship with some noise
    target_weights = np.random.rand(state['n_features'])
    linear_combination = df.dot(target_weights)

    if drift_type == 'model_drift':
        # Introduce a change in how target depends on features
        # For example, invert the importance of one feature or add a new interaction
        target_weights_drift = np.copy(target_weights)
        target_weights_drift[0] *= -1 # Invert weight of first feature
        linear_combination = df.dot(target_weights_drift) # Use new weights

    probability = 1 / (1 + np.exp(-(linear_combination + np.random.normal(0, 1, n_samples)))) # Sigmoid function
    df[state['target_name']] = (probability > 0.5).astype(int)

    logger.debug(f"Simulated data head:\n{df.head()}")
    return df

def train_model(state: dict, df: pd.DataFrame) -> dict:
    """
    Trains a RandomForestClassifier model and stores it in the state.

    Parameters
    ----------
    state : dict
        The current application state dictionary.
    df : pd.DataFrame
        The DataFrame containing features and target for training.

    Returns
    -------
    dict
        The updated state dictionary with the 'trained_model' key populated.
    """
    logger.info("Training a RandomForestClassifier model.")
    X = df[state['feature_names']]
    y = df[state['target_name']]

    # Convert categorical features to 'category' dtype for scikit-learn compatibility if needed
    # For RandomForest, one-hot encoding is often preferred for categorical features if they are not ordinal.
    # For simplicity, assuming numerical representation is sufficient for this simulation.
    # If proper handling is needed, use pd.get_dummies or sklearn.preprocessing.OneHotEncoder

    model = RandomForestClassifier(random_state=state['random_state'])
    model.fit(X, y)
    state['trained_model'] = model
    logger.info("Model training complete.")
    return state

def detect_drift_ks_test(state: dict, reference_data: pd.DataFrame, production_data: pd.DataFrame) -> dict:
    """
    Detects data drift in numerical features using the Kolmogorov-Smirnov (KS) test.

    Parameters
    ----------
    state : dict
        The current application state dictionary containing 'numerical_feature_names' and 'ks_alpha'.
    reference_data : pd.DataFrame
        The DataFrame of the reference dataset.
    production_data : pd.DataFrame
        The DataFrame of the current production dataset.

    Returns
    -------
    dict
        A dictionary where keys are numerical feature names and values are dictionaries
        containing 'drift_detected' (bool) and 'p_value' (float).
    """
    logger.info("Performing KS-test for numerical feature drift detection.")
    drift_results = {}
    for feature in state['numerical_feature_names']:
        if feature not in reference_data.columns or feature not in production_data.columns:
            logger.warning(f"Feature {feature} not found in one of the datasets. Skipping KS-test.")
            continue

        # Add a small random jitter to avoid identical values leading to p-value = 0 issues
        ref_values = reference_data[feature].values + np.random.normal(0, 1e-6, len(reference_data))
        prod_values = production_data[feature].values + np.random.normal(0, 1e-6, len(production_data))

        try:
            stat, p_value = stats.ks_2samp(ref_values, prod_values)
            drift_detected = p_value < state['ks_alpha']
            drift_results[feature] = {'drift_detected': drift_detected, 'p_value': p_value}
            logger.debug(f"KS-test for {feature}: p-value={p_value:.4f}, Drift detected: {drift_detected}")
        except ValueError as e:
            logger.error(f"Error during KS-test for feature {feature}: {e}")
            drift_results[feature] = {'drift_detected': False, 'p_value': np.nan}

    logger.info("KS-test drift detection complete.")
    return drift_results

def detect_drift_chi2_test(state: dict, reference_data: pd.DataFrame, production_data: pd.DataFrame) -> dict:
    """
    Detects data drift in categorical features using the Chi-squared test.

    Parameters
    ----------
    state : dict
        The current application state dictionary containing 'categorical_feature_names' and 'chi2_alpha'.
    reference_data : pd.DataFrame
        The DataFrame of the reference dataset.
    production_data : pd.DataFrame
        The DataFrame of the current production dataset.

    Returns
    -------
    dict
        A dictionary where keys are categorical feature names and values are dictionaries
        containing 'drift_detected' (bool) and 'p_value' (float).
    """
    logger.info("Performing Chi-squared test for categorical feature drift detection.")
    drift_results = {}
    for feature in state['categorical_feature_names']:
        if feature not in reference_data.columns or feature not in production_data.columns:
            logger.warning(f"Feature {feature} not found in one of the datasets. Skipping Chi-squared test.")
            continue

        # Create contingency table
        ref_counts = reference_data[feature].value_counts().sort_index()
        prod_counts = production_data[feature].value_counts().sort_index()

        # Ensure all categories are present in both, fill with 0 if not
        all_categories = sorted(list(set(ref_counts.index).union(set(prod_counts.index))))

        contingency_table = pd.DataFrame({
            'reference': ref_counts.reindex(all_categories, fill_value=0),
            'production': prod_counts.reindex(all_categories, fill_value=0)
        })

        # Chi-squared test requires at least one non-zero value per row/column
        if contingency_table.sum().sum() == 0 or (contingency_table == 0).all().any():
            logger.warning(f"Contingency table for {feature} has all zeros or empty categories after alignment. Skipping Chi-squared test.")
            drift_results[feature] = {'drift_detected': False, 'p_value': np.nan}
            continue

        try:
            stat, p_value, _, _ = stats.chi2_contingency(contingency_table)
            drift_detected = p_value < state['chi2_alpha']
            drift_results[feature] = {'drift_detected': drift_detected, 'p_value': p_value}
            logger.debug(f"Chi-squared test for {feature}: p-value={p_value:.4f}, Drift detected: {drift_detected}")
        except ValueError as e:
            logger.error(f"Error during Chi-squared test for feature {feature}: {e}")
            drift_results[feature] = {'drift_detected': False, 'p_value': np.nan}

    logger.info("Chi-squared test drift detection complete.")
    return drift_results

def calculate_model_performance(state: dict, data: pd.DataFrame) -> dict:
    """
    Calculates the F1-score and accuracy of the trained model on provided data.

    Parameters
    ----------
    state : dict
        The current application state dictionary containing 'trained_model', 'feature_names', and 'target_name'.
    data : pd.DataFrame
        The DataFrame on which to evaluate the model.

    Returns
    -------
    dict
        A dictionary containing 'accuracy' and 'f1_score'. Returns default values
        if the model is not trained.
    """
    logger.info("Calculating model performance.")
    if state['trained_model'] is None:
        logger.warning("No trained model found in state. Skipping performance calculation.")
        return {'accuracy': np.nan, 'f1_score': np.nan}

    X = data[state['feature_names']]
    y_true = data[state['target_name']]

    y_pred = state['trained_model'].predict(X)

    accuracy = accuracy_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred)

    performance_metrics = {'accuracy': accuracy, 'f1_score': f1}
    logger.debug(f"Model performance: {performance_metrics}")
    return performance_metrics

def monitor_production_stream(state: dict, num_batches: int = 20, batch_size: int = 100, drift_frequency: int = 5) -> dict:
    """
    Simulates monitoring a production data stream and performing drift detection.

    Parameters
    ----------
    state : dict
        The current application state dictionary.
    num_batches : int, optional
        The number of simulated data batches to process. Defaults to 20.
    batch_size : int, optional
        The number of samples in each batch. Defaults to 100.
    drift_frequency : int, optional
        How many batches before introducing drift. Defaults to 5.

    Returns
    -------
    dict
        The updated state dictionary, potentially with collected drift metrics and performance data.
    """
    logger.info(f"Starting production stream monitoring for {num_batches} batches.")
    state['drift_detection_history'] = []
    state['model_performance_history'] = []

    for i in range(num_batches):
        # Add random jitter to simulate varying processing times
        time.sleep(0.1 + random.uniform(0, 0.2))

        current_drift_type = 'none'
        current_drift_magnitude = 0.0

        if i == drift_frequency: # Introduce data drift after some batches
            current_drift_type = 'data_drift_numerical'
            current_drift_magnitude = state['drift_magnitude']
            logger.warning(f"Introducing {current_drift_type} at batch {i}.")
        elif i == drift_frequency * 2: # Introduce model drift later
            current_drift_type = 'model_drift'
            current_drift_magnitude = state['drift_magnitude'] * 0.5 # Model drift might be more subtle
            logger.warning(f"Introducing {current_drift_type} at batch {i}.")

        new_production_data = simulate_data(state, batch_size, current_drift_type, current_drift_magnitude)
        state['production_data_window'].extend([new_production_data])

        # Concatenate data in the deque for analysis
        current_production_df = pd.concat(list(state['production_data_window']))

        logger.info(f"Processing batch {i+1}/{num_batches}. Current window size: {len(current_production_df)} samples.")

        # Perform drift detection if enough data in window
        if len(current_production_df) >= state['n_samples']:
            ks_results = detect_drift_ks_test(state, state['reference_data'], current_production_df)
            chi2_results = detect_drift_chi2_test(state, state['reference_data'], current_production_df)
            state['drift_detection_history'].append({
                'batch': i,
                'ks_test': ks_results,
                'chi2_test': chi2_results,
                'drift_type_simulated': current_drift_type
            })

            # Calculate model performance
            model_perf = calculate_model_performance(state, current_production_df)
            state['model_performance_history'].append({'batch': i, **model_perf})

            if model_perf['f1_score'] < state['model_performance_threshold']:
                logger.critical(f"Model performance dropped below threshold at batch {i}: F1-score = {model_perf['f1_score']:.2f}")

    logger.info("Production stream monitoring finished.")
    return state

# Initialize the state
app_state = create_initial_state()

# Simulate reference data and train the model
reference_df = simulate_data(app_state, app_state['n_samples'], drift_type='none')
app_state['reference_data'] = reference_df
app_state = train_model(app_state, reference_df)

logger.info("Initial setup complete. Reference data simulated and model trained.")

# Display reference data head
print("\n--- Reference Data Head ---")
display(app_state['reference_data'].head())

# Simulate monitoring a production stream with potential drift
app_state = monitor_production_stream(app_state, num_batches=20, batch_size=100, drift_frequency=7)

logger.info("Monitoring simulation complete. Analyzing results.")

plt.figure(figsize=(12, 6))

# Identify a numerical feature to plot
num_feature = app_state['numerical_feature_names'][0]

sns.histplot(app_state['reference_data'][num_feature], color='blue', label='Reference Data', kde=True, stat='density', alpha=0.5)

# Find a production batch where numerical drift was introduced
drift_batch_index = next((i for i, d in enumerate(app_state['drift_detection_history']) if d['drift_type_simulated'] == 'data_drift_numerical'), -1)

if drift_batch_index != -1:
    # Reconstruct the production data at that batch
    prod_data_at_drift = pd.concat(list(app_state['production_data_window'])[max(0, drift_batch_index - 5):drift_batch_index + 1]) # Take a few batches around it
    sns.histplot(prod_data_at_drift[num_feature], color='red', label='Production Data (with drift)', kde=True, stat='density', alpha=0.5)
    plt.title(f'Distribution of {num_feature} (Reference vs. Production with Numerical Drift)')
else:
    # If no drift was simulated, just show the latest production data
    sns.histplot(pd.concat(list(app_state['production_data_window']))[num_feature], color='green', label='Latest Production Data (No Drift Simulated)', kde=True, stat='density', alpha=0.5)
    plt.title(f'Distribution of {num_feature} (Reference vs. Latest Production)')

plt.xlabel(num_feature)
plt.ylabel('Density')
plt.legend()
plt.grid(True, linestyle='--', alpha=0.7)
plt.show()

plt.figure(figsize=(12, 6))

# Identify a categorical feature to plot
cat_feature = app_state['categorical_feature_names'][0]

# Calculate counts for reference data
ref_counts = app_state['reference_data'][cat_feature].value_counts(normalize=True).sort_index()
ref_df = ref_counts.reset_index()
ref_df.columns = [cat_feature, 'proportion']
ref_df['dataset'] = 'Reference Data'

# Find a production batch where categorical drift was introduced
drift_batch_index_cat = next((i for i, d in enumerate(app_state['drift_detection_history']) if d['drift_type_simulated'] == 'data_drift_categorical'), -1)

if drift_batch_index_cat != -1:
    # Reconstruct the production data at that batch
    prod_data_at_drift_cat = pd.concat(list(app_state['production_data_window'])[max(0, drift_batch_index_cat - 5):drift_batch_index_cat + 1])
    prod_counts = prod_data_at_drift_cat[cat_feature].value_counts(normalize=True).sort_index()
    prod_df = prod_counts.reset_index()
    prod_df.columns = [cat_feature, 'proportion']
    prod_df['dataset'] = 'Production Data (with drift)'
    plot_df = pd.concat([ref_df, prod_df])
    title_str = f'Distribution of {cat_feature} (Reference vs. Production with Categorical Drift)'
else:
    # If no drift was simulated, just show the latest production data
    latest_prod_data = pd.concat(list(app_state['production_data_window']))
    prod_counts = latest_prod_data[cat_feature].value_counts(normalize=True).sort_index()
    prod_df = prod_counts.reset_index()
    prod_df.columns = [cat_feature, 'proportion']
    prod_df['dataset'] = 'Latest Production Data (No Drift Simulated)'
    plot_df = pd.concat([ref_df, prod_df])
    title_str = f'Distribution of {cat_feature} (Reference vs. Latest Production)'

sns.barplot(x=cat_feature, y='proportion', hue='dataset', data=plot_df, palette='viridis')
plt.title(title_str)
plt.xlabel(cat_feature)
plt.ylabel('Proportion')
plt.legend(title='Dataset')
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.show()

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

# Prepare data for plotting p-values
ks_p_values = []
chi2_p_values = []
batch_numbers = []

for entry in app_state['drift_detection_history']:
    batch_numbers.append(entry['batch'])

    # Collect KS test p-values for the first numerical feature
    ks_p = entry['ks_test'].get(app_state['numerical_feature_names'][0], {}).get('p_value', np.nan)
    ks_p_values.append(ks_p)

    # Collect Chi2 test p-values for the first categorical feature
    chi2_p = entry['chi2_test'].get(app_state['categorical_feature_names'][0], {}).get('p_value', np.nan)
    chi2_p_values.append(chi2_p)

plt.subplot(2, 1, 1)
sns.lineplot(x=batch_numbers, y=ks_p_values, marker='o', label=f'KS Test P-value ({app_state['numerical_feature_names'][0]})', color='blue')
plt.axhline(y=app_state['ks_alpha'], color='red', linestyle='--', label=f'Significance Level (alpha={app_state['ks_alpha']})')
plt.title('Numerical Feature Drift: KS-Test P-values Over Batches')
plt.xlabel('Batch Number')
plt.ylabel('P-value')
plt.ylim(0, 1.05)
plt.legend()
plt.grid(True, linestyle='--', alpha=0.7)

plt.subplot(2, 1, 2)
sns.lineplot(x=batch_numbers, y=chi2_p_values, marker='o', label=f'Chi-squared Test P-value ({app_state['categorical_feature_names'][0]})', color='green')
plt.axhline(y=app_state['chi2_alpha'], color='red', linestyle='--', label=f'Significance Level (alpha={app_state['chi2_alpha']})')
plt.title('Categorical Feature Drift: Chi-squared Test P-values Over Batches')
plt.xlabel('Batch Number')
plt.ylabel('P-value')
plt.ylim(0, 1.05)
plt.legend()
plt.grid(True, linestyle='--', alpha=0.7)

plt.tight_layout()
plt.show()

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

# Prepare data for plotting model performance
model_accuracy = []
model_f1_score = []
perf_batch_numbers = []

for entry in app_state['model_performance_history']:
    perf_batch_numbers.append(entry['batch'])
    model_accuracy.append(entry['accuracy'])
    model_f1_score.append(entry['f1_score'])

plt.subplot(2, 1, 1)
sns.lineplot(x=perf_batch_numbers, y=model_accuracy, marker='o', label='Accuracy', color='purple')
plt.axhline(y=app_state['model_performance_threshold'], color='orange', linestyle='--', label=f'F1-score Threshold ({app_state['model_performance_threshold']})')
plt.title('Model Accuracy Over Batches')
plt.xlabel('Batch Number')
plt.ylabel('Accuracy')
plt.ylim(0, 1.05)
plt.legend()
plt.grid(True, linestyle='--', alpha=0.7)

plt.subplot(2, 1, 2)
sns.lineplot(x=perf_batch_numbers, y=model_f1_score, marker='o', label='F1-score', color='brown')
plt.axhline(y=app_state['model_performance_threshold'], color='orange', linestyle='--', label=f'F1-score Threshold ({app_state['model_performance_threshold']})')
plt.title('Model F1-score Over Batches')
plt.xlabel('Batch Number')
plt.ylabel('F1-score')
plt.ylim(0, 1.05)
plt.legend()
plt.grid(True, linestyle='--', alpha=0.7)

plt.tight_layout()
plt.show()