%pip install pandas numpy scipy scikit-learn matplotlib seaborn faker

import logging
import datetime
from collections import deque
import time
import random
import math

import pandas as pd
import numpy as np
from scipy.stats import ks_2samp, chi2_contingency
from sklearn.preprocessing import LabelEncoder
import matplotlib.pyplot as plt
import seaborn as sns
from faker import Faker

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

def create_drift_monitor_state(window_size: int, alfa: float = 0.05) -> dict:
    """
    Initializes the state dictionary for the feature drift monitor.

    Parameters
    ----------
    window_size : int
        The number of recent data points to keep in the monitoring window.
    alfa : float, optional
        The significance level (alpha) for statistical hypothesis testing, defaults to 0.05.

    Returns
    -------
    dict
        An initialized state dictionary for the drift monitor.
    """
    logger.info(f"Initializing drift monitor state with window_size={window_size}, alfa={alfa}")
    state = {
        'window_size': window_size,
        'alfa': alfa,
        'reference_data': pd.DataFrame(),
        'current_window_data': deque(maxlen=window_size),
        'drift_history': [],
        'feature_types': {},
        'categorical_encoders': {}
    }
    logger.debug("Drift monitor state initialized successfully.")
    return state

def generate_synthetic_data(num_samples: int, num_features: int, num_categorical_features: int, drift_start_sample: int = None) -> pd.DataFrame:
    """
    Generates synthetic tabular data for demonstration purposes, with optional drift.

    Parameters
    ----------
    num_samples : int
        The total number of samples to generate.
    num_features : int
        The number of numerical features.
    num_categorical_features : int
        The number of categorical features.
    drift_start_sample : int, optional
        The sample index at which to start introducing drift, defaults to None.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing the synthetic data.
    """
    logger.info(f"Generating {num_samples} synthetic samples. Drift starts at {drift_start_sample}.")
    data = {}
    fake = Faker()

    # Generate numerical features
    for i in range(num_features):
        feature_name = f'numerical_feature_{i+1}'
        initial_mean = 10 + i * 5
        initial_std = 2 + i
        drift_mean_change = 0 if drift_start_sample is None else 5 + i

        samples = []
        for j in range(num_samples):
            if drift_start_sample is not None and j >= drift_start_sample:
                samples.append(np.random.normal(initial_mean + drift_mean_change, initial_std))
            else:
                samples.append(np.random.normal(initial_mean, initial_std))
        data[feature_name] = samples

    # Generate categorical features
    for i in range(num_categorical_features):
        feature_name = f'categorical_feature_{i+1}'
        categories_initial = [f'CatA_{i}', f'CatB_{i}', f'CatC_{i}']
        categories_drift = [f'CatA_{i}', f'CatB_{i}', f'CatD_{i}'] if i % 2 == 0 else [f'CatX_{i}', f'CatY_{i}']

        samples = []
        for j in range(num_samples):
            if drift_start_sample is not None and j >= drift_start_sample:
                # Introduce new category or change distribution
                samples.append(random.choice(categories_drift))
            else:
                samples.append(random.choice(categories_initial))
        data[feature_name] = samples

    df = pd.DataFrame(data)
    logger.debug(f"Generated DataFrame shape: {df.shape}")
    return df

def _calculate_numerical_drift_ks(reference_series: pd.Series, current_series: pd.Series, alfa: float) -> dict:
    """
    Calculates drift for numerical features using the Kolmogorov-Smirnov (KS) test.

    Parameters
    ----------
    reference_series : pd.Series
        The series of data points from the reference distribution.
    current_series : pd.Series
        The series of data points from the current monitoring window.
    alfa : float
        The significance level for the KS test.

    Returns
    -------
    dict
        A dictionary containing the KS statistic, p-value, and a boolean indicating if drift was detected.
    """
    if reference_series.empty or current_series.empty:
        logger.warning("One of the series is empty for KS test, cannot calculate drift.")
        return {'statistic': np.nan, 'p_value': np.nan, 'drift_detected': False}

    stat, p_value = ks_2samp(reference_series, current_series)
    drift_detected = p_value < alfa
    logger.debug(f"KS test: statistic={stat:.4f}, p_value={p_value:.4f}, drift_detected={drift_detected}")
    return {'statistic': stat, 'p_value': p_value, 'drift_detected': drift_detected}

def _calculate_categorical_drift_chi2(reference_series: pd.Series, current_series: pd.Series, alfa: float) -> dict:
    """
    Calculates drift for categorical features using the Chi-squared test.

    Parameters
    ----------
    reference_series : pd.Series
        The series of data points from the reference distribution.
    current_series : pd.Series
        The series of data points from the current monitoring window.
    alfa : float
        The significance level for the Chi-squared test.

    Returns
    -------
    dict
        A dictionary containing the Chi-squared statistic, p-value, and a boolean indicating if drift was detected.
    """
    if reference_series.empty or current_series.empty:
        logger.warning("One of the series is empty for Chi-squared test, cannot calculate drift.")
        return {'statistic': np.nan, 'p_value': np.nan, 'drift_detected': False}

    # Combine categories to ensure all levels are considered
    all_categories = pd.concat([reference_series, current_series]).unique()

    ref_counts = reference_series.value_counts().reindex(all_categories, fill_value=0)
    current_counts = current_series.value_counts().reindex(all_categories, fill_value=0)

    # Ensure there are no zero sums for chi2_contingency
    # If all counts are zero for a category, it might cause issues, but reindex should handle it.
    # chi2_contingency expects a 2D array: rows are distributions, columns are categories
    contingency_table = pd.DataFrame({'reference': ref_counts, 'current': current_counts})

    # Remove categories where both reference and current counts are zero
    contingency_table = contingency_table[(contingency_table['reference'] > 0) | (contingency_table['current'] > 0)]

    if contingency_table.empty:
        logger.warning("Contingency table is empty after filtering, cannot calculate Chi-squared drift.")
        return {'statistic': np.nan, 'p_value': np.nan, 'drift_detected': False}

    try:
        stat, p_value, _, _ = chi2_contingency(contingency_table.values)
        drift_detected = p_value < alfa
        logger.debug(f"Chi-squared test: statistic={stat:.4f}, p_value={p_value:.4f}, drift_detected={drift_detected}")
        return {'statistic': stat, 'p_value': p_value, 'drift_detected': drift_detected}
    except ValueError as e:
        logger.error(f"Error during Chi-squared test: {e}. Contingency table: {contingency_table}")
        return {'statistic': np.nan, 'p_value': np.nan, 'drift_detected': False}

def set_reference_data(state: dict, data: pd.DataFrame) -> dict:
    """
    Sets the baseline (reference) data distribution for drift detection.
    Infers feature types and initializes categorical encoders.

    Parameters
    ----------
    state : dict
        Current state dictionary.
    data : pd.DataFrame
        The DataFrame to be used as the new reference data.

    Returns
    -------
    dict
        Updated state dictionary with the new reference data and feature types.
    """
    logger.info("Setting new reference data.")
    state['reference_data'] = data.copy()
    state['feature_types'] = {}
    state['categorical_encoders'] = {}

    for col in data.columns:
        if pd.api.types.is_numeric_dtype(data[col]):
            state['feature_types'][col] = 'numerical'
        else:
            state['feature_types'][col] = 'categorical'
            # Initialize LabelEncoder for categorical features if needed
            le = LabelEncoder()
            # Fit on all unique values from reference data to avoid issues later
            le.fit(data[col].astype(str).unique())
            state['categorical_encoders'][col] = le

    logger.debug(f"Reference data set. Features and types: {state['feature_types']}")
    return state

def detect_drift(state: dict) -> dict:
    """
    Detects drift for all features in the current window against the reference data.

    Parameters
    ----------
    state : dict
        Current state dictionary.

    Returns
    -------
    dict
        Updated state dictionary with detected drift events.
    """
    logger.info("Starting drift detection for current window.")
    if state['reference_data'].empty:
        logger.warning("Reference data is not set. Cannot detect drift.")
        return state

    if not state['current_window_data']:
        logger.warning("Current window data is empty. Cannot detect drift.")
        return state

    current_window_df = pd.DataFrame(list(state['current_window_data']))

    drift_results = {
        'timestamp': datetime.datetime.now(),
        'features': {}
    }

    for feature_name, feature_type in state['feature_types'].items():
        if feature_name not in current_window_df.columns:
            logger.warning(f"Feature '{feature_name}' not found in current window data. Skipping.")
            continue

        reference_series = state['reference_data'][feature_name]
        current_series = current_window_df[feature_name]

        # Handle potential mismatch in categorical levels between reference and current
        if feature_type == 'categorical':
            # Ensure both series are treated as strings before comparison
            reference_series = reference_series.astype(str)
            current_series = current_series.astype(str)

            # Add jitter before retry attempts
            jitter = random.uniform(0.01, 0.1)
            time.sleep(jitter)

            try:
                result = _calculate_categorical_drift_chi2(reference_series, current_series, state['alfa'])
            except Exception as e:
                logger.error(f"Attempt failed for {feature_name} (categorical): {e}")
                # Simple retry logic without full exponential backoff for brevity here
                time.sleep(1 + jitter)
                try:
                    result = _calculate_categorical_drift_chi2(reference_series, current_series, state['alfa'])
                except Exception as e_retry:
                    logger.error(f"Second attempt failed for {feature_name} (categorical): {e_retry}. Skipping feature.")
                    result = {'statistic': np.nan, 'p_value': np.nan, 'drift_detected': False}

        elif feature_type == 'numerical':
            # Add jitter before retry attempts
            jitter = random.uniform(0.01, 0.1)
            time.sleep(jitter)

            try:
                result = _calculate_numerical_drift_ks(reference_series, current_series, state['alfa'])
            except Exception as e:
                logger.error(f"Attempt failed for {feature_name} (numerical): {e}")
                time.sleep(1 + jitter)
                try:
                    result = _calculate_numerical_drift_ks(reference_series, current_series, state['alfa'])
                except Exception as e_retry:
                    logger.error(f"Second attempt failed for {feature_name} (numerical): {e_retry}. Skipping feature.")
                    result = {'statistic': np.nan, 'p_value': np.nan, 'drift_detected': False}

        else:
            logger.warning(f"Unknown feature type '{feature_type}' for feature '{feature_name}'. Skipping.")
            continue

        drift_results['features'][feature_name] = result
        if result['drift_detected']:
            logger.warning(f"Drift detected for feature '{feature_name}' (p-value: {result['p_value']:.4f})")

    state['drift_history'].append(drift_results)
    logger.debug("Drift detection completed for current window.")
    return state

def update_monitor_with_new_data(state: dict, new_data_batch: pd.DataFrame) -> dict:
    """
    Adds new data to the monitoring window and triggers drift detection.

    Parameters
    ----------
    state : dict
        Current state dictionary.
    new_data_batch : pd.DataFrame
        A DataFrame containing new data points to add to the monitor.

    Returns
    -------
    dict
        Updated state dictionary after processing the new data and detecting drift.
    """
    logger.info(f"Updating monitor with {len(new_data_batch)} new data points.")
    for _, row in new_data_batch.iterrows():
        state['current_window_data'].append(row.to_dict())

    if len(state['current_window_data']) == state['window_size']:
        logger.debug("Current window is full, performing drift detection.")
        state = detect_drift(state)
    else:
        logger.info(f"Current window not yet full ({len(state['current_window_data'])}/{state['window_size']}). Skipping drift detection.")

    return state

def summarize_drift_history(state: dict) -> pd.DataFrame:
    """
    Summarizes the drift detection history into a pandas DataFrame.

    Parameters
    ----------
    state : dict
        Current state dictionary.

    Returns
    -------
    pd.DataFrame
        A DataFrame summarizing the drift detection history.
    """
    logger.info("Summarizing drift detection history.")
    summary_records = []
    for entry in state['drift_history']:
        timestamp = entry['timestamp']
        for feature_name, results in entry['features'].items():
            summary_records.append({
                'timestamp': timestamp,
                'feature': feature_name,
                'statistic': results['statistic'],
                'p_value': results['p_value'],
                'drift_detected': results['drift_detected']
            })

    if not summary_records:
        logger.warning("No drift history records to summarize.")
        return pd.DataFrame()

    df_summary = pd.DataFrame(summary_records)
    logger.debug(f"Drift history summarized into DataFrame of shape: {df_summary.shape}")
    return df_summary

# 1. Generate Synthetic Data with Drift
TOTAL_SAMPLES = 1000
NUM_NUMERICAL_FEATURES = 3
NUM_CATEGORICAL_FEATURES = 2
DRIFT_START_SAMPLE = 600 # Drift starts after 600 samples

logger.info("Generating synthetic dataset for demonstration.")
synthetic_data = generate_synthetic_data(
    num_samples=TOTAL_SAMPLES,
    num_features=NUM_NUMERICAL_FEATURES,
    num_categorical_features=NUM_CATEGORICAL_FEATURES,
    drift_start_sample=DRIFT_START_SAMPLE
)
display(synthetic_data.head())
display(synthetic_data.tail())

logger.info(f"Total samples generated: {len(synthetic_data)}")

# 2. Initialize the Drift Monitor
WINDOW_SIZE = 100 # Compare last 100 samples against reference
ALFA_LEVEL = 0.01 # Significance level

logger.info("Initializing drift monitor.")
drift_monitor_state = create_drift_monitor_state(window_size=WINDOW_SIZE, alfa=ALFA_LEVEL)

# Set the first 500 samples as the reference data (before drift)
REFERENCE_DATA_SIZE = 500
initial_reference_data = synthetic_data.iloc[:REFERENCE_DATA_SIZE]
drift_monitor_state = set_reference_data(drift_monitor_state, initial_reference_data)

logger.info(f"Reference data set with {len(initial_reference_data)} samples.")
print("Reference Data Features and Types:")
display(pd.DataFrame([drift_monitor_state['feature_types']]).T.rename(columns={0: 'Type'}))

# 3. Continuously Feed New Data Batches and Detect Drift
BATCH_SIZE = 10

logger.info(f"Simulating data stream with batch size {BATCH_SIZE}.")

# Start monitoring from where reference data ends
for i in range(REFERENCE_DATA_SIZE, TOTAL_SAMPLES, BATCH_SIZE):
    end_idx = min(i + BATCH_SIZE, TOTAL_SAMPLES)
    current_batch = synthetic_data.iloc[i:end_idx]

    # Simulate some processing time with random jitter
    jitter_delay = random.uniform(0.01, 0.1)
    time.sleep(jitter_delay)

    logger.debug(f"Processing batch from index {i} to {end_idx-1}.")
    drift_monitor_state = update_monitor_with_new_data(drift_monitor_state, current_batch)

logger.info("Data stream simulation complete.")

# Summarize the drift history
drift_summary_df = summarize_drift_history(drift_monitor_state)

if not drift_summary_df.empty:
    print("Drift Detection Summary:")
    display(drift_summary_df.head())
    display(drift_summary_df.tail())
else:
    print("No drift events recorded or insufficient data to summarize.")

# 4. Visualize Drift Over Time
if not drift_summary_df.empty:
    plt.figure(figsize=(15, 7))
    sns.lineplot(data=drift_summary_df, x='timestamp', y='p_value', hue='feature', marker='o')
    plt.axhline(y=ALFA_LEVEL, color='r', linestyle='--', label=f'Significance Level (alpha={ALFA_LEVEL})')
    plt.title('Feature Drift p-values Over Time')
    plt.xlabel('Time')
    plt.ylabel('P-value')
    plt.xticks(rotation=45)
    plt.legend(title='Feature')
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()

    # Plot just the drift detected points
    drift_detected_df = drift_summary_df[drift_summary_df['drift_detected'] == True]
    if not drift_detected_df.empty:
        plt.figure(figsize=(15, 5))
        sns.scatterplot(data=drift_detected_df, x='timestamp', y='p_value', hue='feature', style='feature', s=100, palette='deep')
        plt.axhline(y=ALFA_LEVEL, color='r', linestyle='--', label=f'Significance Level (alpha={ALFA_LEVEL})')
        plt.title('Detected Feature Drift Events (p-value < alpha)')
        plt.xlabel('Time')
        plt.ylabel('P-value')
        plt.xticks(rotation=45)
        plt.legend(title='Drifting Feature')
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.show()
    else:
        logger.info("No drift detected to visualize.")
else:
    logger.info("No drift summary data available for visualization.")

# 5. Compare Feature Distributions Before and After Drift

# Define the 'before drift' data (reference data)
before_drift_df = initial_reference_data

# Define the 'after drift' data (e.g., a window well after drift_start_sample)
after_drift_start_idx = max(DRIFT_START_SAMPLE + WINDOW_SIZE, REFERENCE_DATA_SIZE + WINDOW_SIZE)
after_drift_end_idx = min(after_drift_start_idx + WINDOW_SIZE, TOTAL_SAMPLES)
after_drift_df = synthetic_data.iloc[after_drift_start_idx:after_drift_end_idx]

if after_drift_df.empty:
    logger.warning("Insufficient 'after drift' data to perform comparison visualizations.")
else:
    print(f"Comparing distributions using {len(before_drift_df)} 'before drift' samples and {len(after_drift_df)} 'after drift' samples.")
    for feature_name, f_type in drift_monitor_state['feature_types'].items():
        plt.figure(figsize=(12, 5))
        if f_type == 'numerical':
            sns.histplot(before_drift_df[feature_name], color='blue', label='Before Drift', kde=True, stat='density', alpha=0.6)
            sns.histplot(after_drift_df[feature_name], color='red', label='After Drift', kde=True, stat='density', alpha=0.6)
            plt.title(f'Distribution of {feature_name} (Numerical)')
            plt.xlabel(feature_name)
            plt.ylabel('Density')
        else: # Categorical
            # Combine and normalize counts for comparison
            before_counts = before_drift_df[feature_name].value_counts(normalize=True).sort_index()
            after_counts = after_drift_df[feature_name].value_counts(normalize=True).sort_index()

            combined_index = sorted(list(set(before_counts.index).union(set(after_counts.index))))
            plot_df = pd.DataFrame({
                'Before Drift': before_counts.reindex(combined_index, fill_value=0),
                'After Drift': after_counts.reindex(combined_index, fill_value=0)
            })
            plot_df.plot(kind='bar', ax=plt.gca(), width=0.8)
            plt.title(f'Distribution of {feature_name} (Categorical)')
            plt.xlabel(feature_name)
            plt.ylabel('Proportion')
            plt.xticks(rotation=45, ha='right')

        plt.legend()
        plt.grid(axis='y', linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.show()