# Install necessary libraries
!pip install pandas numpy scikit-learn tensorflow google-cloud-storage google-cloud-aiplatform matplotlib seaborn --quiet
!pip install loguru --quiet # For structured logging
!pip install tenacity --quiet # For retry mechanisms

# Standard library imports
import os
import sys
import time
import logging
import random
from collections import deque
from typing import Dict, Any, List, Tuple

# Third-party library imports
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
from sklearn.preprocessing import StandardScaler
import tensorflow as tf # Example for a deep learning framework
from google.cloud import storage
from google.cloud import aiplatform
from loguru import logger
from tenacity import retry, wait_exponential, stop_after_attempt, Retrying, wait_fixed, retry_if_exception_type

# Visualization libraries
import matplotlib.pyplot as plt
import seaborn as sns

# Configure logging with loguru
logger.remove()
logger.add(sys.stderr, level="INFO", format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>")

# Suppress TensorFlow warnings (optional, adjust as needed)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

# Configure AI Platform logging (optional)
# aiplatform.init(project='YOUR_GCP_PROJECT_ID', location='YOUR_GCP_REGION')

logger.info("All necessary libraries imported and logging configured.")

def create_inference_state(
    project_id: str,
    bucket_name: str,
    model_name: str,
    model_version: str,
    target_variable: str,
    feature_columns: List[str],
    batch_size: int = 1000
) -> Dict[str, Any]:
    """
    Initializes the state dictionary for the batch inference pipeline.

    This function sets up fundamental configurations such as GCP project details,
    GCS bucket and model information, target variable, feature columns, and
    batch processing size. It ensures all necessary parameters are available
    for subsequent steps in the pipeline.

    Parameters
    ----------
    project_id : str
        The Google Cloud Project ID.
    bucket_name : str
        The name of the GCS bucket for model artifacts and data.
    model_name : str
        The name of the model to be used for inference.
    model_version : str
        The version of the model to load.
    target_variable : str
        The name of the target variable (label) in the dataset.
    feature_columns : List[str]
        A list of column names used as features for the model.
    batch_size : int, optional
        The number of records to process in each batch during inference,
        defaults to 1000.

    Returns
    -------
    Dict[str, Any]
        An initialized state dictionary containing all necessary configuration
        for the inference pipeline.

    Examples
    --------
    >>> state = create_inference_state(
    ...     project_id="my-gcp-project",
    ...     bucket_name="my-ml-bucket",
    ...     model_name="fraud-detection-model",
    ...     model_version="v1.0.0",
    ...     target_variable="is_fraud",
    ...     feature_columns=["transaction_amount", "user_age"],
    ...     batch_size=500
    ... )
    >>> assert state['project_id'] == "my-gcp-project"
    >>> assert state['batch_size'] == 500
    """
    logger.info("Initializing inference state.")
    state = {
        "project_id": project_id,
        "bucket_name": bucket_name,
        "model_name": model_name,
        "model_version": model_version,
        "target_variable": target_variable,
        "feature_columns": feature_columns,
        "batch_size": batch_size,
        "model_path": f"gs://{bucket_name}/models/{model_name}/{model_version}/",
        "raw_data_path": f"gs://{bucket_name}/data/raw/",
        "processed_data_path": f"gs://{bucket_name}/data/processed/",
        "inference_results_path": f"gs://{bucket_name}/inference_results/",
        "inference_start_time": None,
        "inference_end_time": None,
        "total_records_processed": 0,
        "predictions": [],
        "metrics": {},
        "errors": []
    }
    logger.info(f"Inference state initialized for model '{model_name}' version '{model_version}'.")
    logger.debug(f"Initial state: {state}")
    return state

def load_data_for_inference(
    state: Dict[str, Any],
    file_pattern: str = "*.csv"
) -> Tuple[Dict[str, Any], pd.DataFrame]:
    """
    Loads raw data for inference from Google Cloud Storage (GCS).

    This function reads data from the specified `raw_data_path` in the state
    dictionary, supporting retry logic for resilient data retrieval.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary, must contain 'raw_data_path',
        'bucket_name', and 'project_id'.
    file_pattern : str, optional
        Pattern to filter files within the raw_data_path, e.g., '*.csv',
        defaults to "*.csv".

    Returns
    -------
    Tuple[Dict[str, Any], pd.DataFrame]
        An updated state dictionary and the loaded raw data as a pandas DataFrame.
        Returns an empty DataFrame if no data is found or an error occurs after retries.

    Examples
    --------
    >>> current_state = {
    ...     "project_id": "test-project",
    ...     "bucket_name": "test-bucket",
    ...     "raw_data_path": "gs://test-bucket/data/raw/inference_data/"
    ... }
    >>> # Assume 'test-bucket/data/raw/inference_data/batch_1.csv' exists in GCS
    >>> # with some data.
    >>> # state, df = load_data_for_inference(current_state, file_pattern="batch_*.csv")
    >>> # if not df.empty: print(f"Loaded {len(df)} records.")
    """
    raw_data_path = state.get("raw_data_path")
    bucket_name = state.get("bucket_name")
    project_id = state.get("project_id")

    if not all([raw_data_path, bucket_name, project_id]):
        logger.error("Missing raw_data_path, bucket_name, or project_id in state for data loading.")
        state["errors"].append("Missing state parameters for data loading.")
        return state, pd.DataFrame()

    logger.info(f"Attempting to load data from GCS path: {raw_data_path} with pattern: {file_pattern}")
    data_df = pd.DataFrame()

    @retry(
        wait=wait_exponential(multiplier=1, min=4, max=60),
        stop=stop_after_attempt(5),
        retry=retry_if_exception_type((IOError, ConnectionError, Exception)),
        reraise=True
    )
    def _load_from_gcs():
        logger.debug(f"Retrying data load for {raw_data_path}...")
        client = storage.Client(project=project_id)
        # Extract bucket name from path if not explicitly provided, or use state's
        path_parts = raw_data_path.replace(f"gs://{bucket_name}/", "").split('/')
        prefix = '/'.join(path_parts) if path_parts[-1] else '/'.join(path_parts[:-1])
        if prefix and not prefix.endswith('/'):
            prefix += '/'

        blobs = list(client.bucket(bucket_name).list_blobs(prefix=prefix))
        csv_files = [blob for blob in blobs if blob.name.endswith('.csv') and glob_match(blob.name.split('/')[-1], file_pattern)]

        if not csv_files:
            logger.warning(f"No CSV files found matching '{file_pattern}' in '{raw_data_path}'.")
            return pd.DataFrame()

        all_dfs = []
        for blob in csv_files:
            blob_path = f"gs://{bucket_name}/{blob.name}"
            logger.info(f"Reading file: {blob_path}")
            # Add a small random jitter to prevent thundering herd problem on GCS
            time.sleep(random.uniform(0.1, 0.5))
            try:
                df = pd.read_csv(blob_path)
                all_dfs.append(df)
            except Exception as e:
                logger.warning(f"Could not read {blob_path}: {e}")
                # If one file fails, try to continue with others or raise if critical.
                # For now, we'll log and continue.

        if not all_dfs:
            logger.error(f"Failed to load any data from {raw_data_path} after filtering.")
            return pd.DataFrame()

        return pd.concat(all_dfs, ignore_index=True)

    def glob_match(filename: str, pattern: str) -> bool:
        # Simple glob-like matching for file patterns
        # More robust solution would use fnmatch.fnmatch
        pattern_parts = pattern.replace('.', '\\.').replace('*', '.*').split('/')
        filename_parts = filename.split('/')
        if len(pattern_parts) != len(filename_parts):
            return False
        for p, f in zip(pattern_parts, filename_parts):
            if not re.fullmatch(p, f):
                return False
        return True

    import re # Imported locally for glob_match

    try:
        data_df = _load_from_gcs()
        logger.info(f"Successfully loaded {len(data_df)} records for inference.")
    except Exception as e:
        logger.error(f"Failed to load data after multiple retries from {raw_data_path}: {e}")
        state["errors"].append(f"Data loading failed: {e}")
        data_df = pd.DataFrame() # Ensure data_df is empty on failure

    state["total_records_processed"] = len(data_df)
    return state, data_df

def preprocess_data_for_inference(
    state: Dict[str, Any],
    raw_df: pd.DataFrame
) -> Tuple[Dict[str, Any], pd.DataFrame]:
    """
    Preprocesses raw data for model inference.

    Performs data cleaning, feature engineering, and scaling to prepare
    the input data for the machine learning model. This includes handling
    missing values, scaling numerical features, and potentially encoding
    categorical features.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary, must contain 'feature_columns'.
        Can also contain 'scaler' for pre-fitted scalers.
    raw_df : pd.DataFrame
        The raw input data to be preprocessed.

    Returns
    -------
    Tuple[Dict[str, Any], pd.DataFrame]
        An updated state dictionary and the preprocessed data as a pandas DataFrame.
        Returns an empty DataFrame if preprocessing fails.

    Examples
    --------
    >>> current_state = {
    ...     "feature_columns": ["numerical_feature", "categorical_feature"],
    ...     "target_variable": "label"
    ... }
    >>> data = pd.DataFrame({
    ...     "numerical_feature": [10, 20, np.nan, 40, 50],
    ...     "categorical_feature": ["A", "B", "A", "C", "B"],
    ...     "another_feature": [1,2,3,4,5]
    ... })
    >>> # state, processed_df = preprocess_data_for_inference(current_state, data)
    >>> # if not processed_df.empty: print(processed_df.head())
    """
    logger.info("Starting data preprocessing for inference.")

    if raw_df.empty:
        logger.warning("Raw DataFrame is empty, skipping preprocessing.")
        return state, pd.DataFrame()

    feature_cols = state.get("feature_columns", [])
    if not feature_cols:
        logger.error("Feature columns not defined in state. Cannot preprocess.")
        state["errors"].append("Feature columns missing for preprocessing.")
        return state, pd.DataFrame()

    processed_df = raw_df.copy()

    try:
        # Identify numerical and categorical features that are in the raw_df
        available_feature_cols = [col for col in feature_cols if col in processed_df.columns]

        numerical_cols = processed_df[available_feature_cols].select_dtypes(include=np.number).columns.tolist()
        categorical_cols = processed_df[available_feature_cols].select_dtypes(include='object').columns.tolist()

        logger.debug(f"Numerical columns identified: {numerical_cols}")
        logger.debug(f"Categorical columns identified: {categorical_cols}")

        # Handle missing values with a small random jitter to avoid perfect colinearity if values are constant
        for col in numerical_cols:
            if processed_df[col].isnull().any():
                median_val = processed_df[col].median()
                logger.info(f"Filling missing numerical values in '{col}' with median: {median_val}")
                processed_df[col] = processed_df[col].fillna(median_val + random.uniform(-0.01, 0.01))

        for col in categorical_cols:
            if processed_df[col].isnull().any():
                mode_val = processed_df[col].mode()[0] # mode() can return multiple if tied
                logger.info(f"Filling missing categorical values in '{col}' with mode: {mode_val}")
                processed_df[col] = processed_df[col].fillna(mode_val)

        # Scaling numerical features
        if numerical_cols:
            logger.info(f"Scaling numerical features: {numerical_cols}")
            scaler = state.get("scaler")
            if scaler is None:
                logger.warning("No pre-fitted scaler found in state. Fitting new StandardScaler. \n In a production scenario, a pre-fitted scaler should be loaded.")
                scaler = StandardScaler()
                processed_df[numerical_cols] = scaler.fit_transform(processed_df[numerical_cols])
                state["scaler"] = scaler # Store for potential future use or consistency
            else:
                logger.info("Using pre-fitted scaler from state to transform numerical features.")
                processed_df[numerical_cols] = scaler.transform(processed_df[numerical_cols])

        # One-hot encode categorical features (if any)
        if categorical_cols:
            logger.info(f"One-hot encoding categorical features: {categorical_cols}")
            # For simplicity, we'll use pandas get_dummies. In a real MLOps setting,
            # a consistent encoder (e.g., from sklearn or a stored mapping) should be used.
            processed_df = pd.get_dummies(processed_df, columns=categorical_cols, drop_first=True, dtype=int)

        # Update feature_columns in state to reflect the actual columns in processed_df
        # This is crucial for subsequent steps like run_batch_inference
        state["feature_columns"] = processed_df.columns.tolist()

        logger.info(f"Data preprocessing complete. Processed {len(processed_df)} records.")
        logger.debug(f"Processed DataFrame head:\n{processed_df.head()}")

    except Exception as e:
        logger.error(f"Error during data preprocessing: {e}", exc_info=True)
        state["errors"].append(f"Preprocessing failed: {e}")
        return state, pd.DataFrame()

    return state, processed_df

import joblib # For loading scikit-learn models

def load_model_for_inference(
    state: Dict[str, Any]
) -> Tuple[Dict[str, Any], Any]: # Any for the model type, as it can be tf.Model, sklearn model etc.
    """
    Loads a trained machine learning model for inference from GCS.

    This function retrieves the model based on the 'model_path' specified in
    the state dictionary. It supports different model formats and includes
    retry logic for resilient model loading.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary, must contain 'model_path' and 'model_name'.

    Returns
    -------
    Tuple[Dict[str, Any], Any]
        An updated state dictionary and the loaded model object. Returns `None`
        as the model object if loading fails after retries.

    Examples
    --------
    >>> current_state = {
    ...     "project_id": "my-gcp-project",
    ...     "bucket_name": "my-ml-bucket",
    ...     "model_name": "my-tf-model",
    ...     "model_version": "1",
    ...     "model_path": "gs://my-ml-bucket/models/my-tf-model/1/"
    ... }
    >>> # Example: model_object = tf.keras.models.Sequential([...])
    >>> # model_object.save('my-tf-model/1/')
    >>> # state, model = load_model_for_inference(current_state)
    >>> # if model: print("Model loaded successfully!")
    """
    model_path = state.get("model_path")
    model_name = state.get("model_name")

    if not all([model_path, model_name]):
        logger.error("Missing model_path or model_name in state for model loading.")
        state["errors"].append("Missing state parameters for model loading.")
        return state, None

    logger.info(f"Attempting to load model '{model_name}' from: {model_path}")
    loaded_model = None

    @retry(
        wait=wait_exponential(multiplier=1, min=4, max=60),
        stop=stop_after_attempt(5),
        retry=retry_if_exception_type((IOError, ConnectionError, Exception)),
        reraise=True
    )
    def _load_model_with_retry():
        logger.debug(f"Retrying model load for {model_name} from {model_path}...")
        # Add a small random jitter to prevent thundering herd problem
        time.sleep(random.uniform(0.1, 0.5))
        try:
            # Heuristic to determine model type - can be made more robust with metadata
            if "tensorflow" in model_path or "tf" in model_name.lower():
                # Assuming TensorFlow SavedModel format
                logger.info(f"Loading TensorFlow model from {model_path}")
                model = tf.keras.models.load_model(model_path)
            elif "sklearn" in model_name.lower() or model_path.endswith('.joblib'):
                # Assuming Scikit-learn model saved with joblib
                logger.info(f"Loading Scikit-learn model from {model_path}")
                # For joblib, typically load a single file, not a directory
                # This part might need adjustment based on how the sklearn model is saved
                # For simplicity, assuming model_path points directly to the .joblib file
                client = storage.Client(project=state.get("project_id"))
                bucket = client.get_bucket(state.get("bucket_name"))
                blob_name = model_path.replace(f"gs://{state.get('bucket_name')}/", "")
                blob = bucket.blob(blob_name)
                # Download to a temporary file, then load
                temp_model_file = f"/tmp/{model_name}.joblib"
                blob.download_to_filename(temp_model_file)
                model = joblib.load(temp_model_file)
                os.remove(temp_model_file)
            else:
                logger.warning(f"Unknown model type for '{model_name}'. Attempting generic load.")
                # Fallback for other model types if applicable
                # For example, for PyTorch, you might need to define a model architecture first
                raise ValueError("Unsupported model type or format.")

            return model
        except Exception as e:
            logger.error(f"Error loading model '{model_name}' from {model_path}: {e}")
            raise # Re-raise for tenacity to catch

    try:
        loaded_model = _load_model_with_retry()
        state["model"] = loaded_model
        logger.info(f"Model '{model_name}' loaded successfully.")
    except Exception as e:
        logger.error(f"Failed to load model '{model_name}' after multiple retries: {e}")
        state["errors"].append(f"Model loading failed: {e}")
        state["model"] = None # Ensure model is None on failure

    return state, loaded_model

def run_batch_inference(
    state: Dict[str, Any],
    processed_df: pd.DataFrame
) -> Tuple[Dict[str, Any], pd.Series]:
    """
    Executes batch inference using the loaded model on preprocessed data.

    Processes the input DataFrame in batches, makes predictions, and stores
    the results in the state dictionary. Supports TensorFlow and Scikit-learn models.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary, must contain the loaded 'model',
        'batch_size', and 'feature_columns'.
    processed_df : pd.DataFrame
        The preprocessed data to run inference on.

    Returns
    -------
    Tuple[Dict[str, Any], pd.Series]
        An updated state dictionary and a pandas Series containing the generated
        predictions. Returns an empty Series if inference fails.

    Examples
    --------
    >>> # Assuming 'state' contains a loaded model and 'processed_df' has features
    >>> # state, predictions = run_batch_inference(state, processed_df)
    >>> # if not predictions.empty: print(predictions.head())
    """
    logger.info("Starting batch inference.")

    model = state.get("model")
    batch_size = state.get("batch_size")
    feature_cols = state.get("feature_columns")

    if model is None:
        logger.error("No model found in state. Cannot run inference.")
        state["errors"].append("Model missing for inference.")
        return state, pd.Series()

    if processed_df.empty:
        logger.warning("Processed DataFrame is empty, skipping inference.")
        return state, pd.Series()

    if not feature_cols:
        logger.error("Feature columns not defined in state. Cannot run inference.")
        state["errors"].append("Feature columns missing for inference.")
        return state, pd.Series()

    predictions_list = []
    state["inference_start_time"] = time.time()

    num_batches = int(np.ceil(len(processed_df) / batch_size))
    logger.info(f"Running inference on {len(processed_df)} records in {num_batches} batches of size {batch_size}.")

    try:
        for i in range(num_batches):
            start_idx = i * batch_size
            end_idx = min((i + 1) * batch_size, len(processed_df))
            batch_data = processed_df.iloc[start_idx:end_idx]

            # Ensure batch_data has only the feature columns and in correct order
            # This step is crucial if get_dummies added new columns or changed order
            # A robust solution would involve storing the exact feature order from training
            # For now, we'll assume processed_df is already aligned from the preprocess step
            input_features = batch_data[feature_cols]

            # Add random jitter to simulate variable processing times, for backoff scenario
            time.sleep(random.uniform(0.01, 0.05))

            batch_predictions = None
            # Heuristic to determine model type based on its method attributes
            if hasattr(model, 'predict') and callable(getattr(model, 'predict')):
                if hasattr(model, 'call') or isinstance(model, tf.keras.Model): # More specific for TF Keras
                    logger.debug(f"Making TensorFlow Keras predictions for batch {i+1}/{num_batches}")
                    # Keras predict expects numpy array or tf.Tensor
                    batch_predictions = model.predict(input_features.values).flatten()
                elif hasattr(model, 'predict_proba') and callable(getattr(model, 'predict_proba')):
                    logger.debug(f"Making Scikit-learn predict_proba predictions for batch {i+1}/{num_batches}")
                    # Scikit-learn models usually have predict_proba for classification
                    batch_predictions = model.predict_proba(input_features.values)[:, 1] # Probability of the positive class
                else:
                    logger.debug(f"Making Scikit-learn predict predictions for batch {i+1}/{num_batches}")
                    batch_predictions = model.predict(input_features.values)
            else:
                logger.error("Loaded model does not have a 'predict' method. Cannot perform inference.")
                raise AttributeError("Model lacks predict method.")

            if batch_predictions is not None:
                predictions_list.extend(batch_predictions)
            logger.info(f"Processed batch {i+1}/{num_batches}. Records: {end_idx - start_idx}")

        state["inference_end_time"] = time.time()
        total_inference_time = state["inference_end_time"] - state["inference_start_time"]
        logger.info(f"Batch inference completed in {total_inference_time:.2f} seconds.")
        state["total_records_processed"] = len(processed_df)
        state["predictions"] = predictions_list

        return state, pd.Series(predictions_list, index=processed_df.index)

    except Exception as e:
        logger.error(f"Error during batch inference: {e}", exc_info=True)
        state["errors"].append(f"Inference failed: {e}")
        state["inference_end_time"] = time.time() # End time even on failure
        return state, pd.Series()

def postprocess_predictions(
    state: Dict[str, Any],
    original_indices: pd.Index,
    prediction_threshold: float = 0.5
) -> Tuple[Dict[str, Any], pd.DataFrame]:
    """
    Post-processes raw model predictions to make them actionable.

    This typically involves converting probabilities to class labels based on a
    threshold, and combining predictions with original data identifiers.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary, must contain 'predictions' (list of raw predictions)
        and 'target_variable' (for naming the prediction column).
    original_indices : pd.Index
        The index from the original (or processed) DataFrame to align predictions.
    prediction_threshold : float, optional
        The threshold to convert continuous predictions (probabilities) into
        binary class labels, defaults to 0.5.

    Returns
    -------
    Tuple[Dict[str, Any], pd.DataFrame]
        An updated state dictionary and a pandas DataFrame containing the
        post-processed predictions, with 'prediction' and 'confidence' columns.
        Returns an empty DataFrame if post-processing fails or no predictions.

    Examples
    --------
    >>> current_state = {
    ...     "predictions": [0.1, 0.9, 0.3, 0.7, 0.05],
    ...     "target_variable": "is_fraud"
    ... }
    >>> original_idx = pd.Index([101, 102, 103, 104, 105])
    >>> # state, final_preds_df = postprocess_predictions(current_state, original_idx)
    >>> # if not final_preds_df.empty: print(final_preds_df.head())
    """
    logger.info("Starting prediction post-processing.")

    raw_predictions = state.get("predictions")
    target_variable = state.get("target_variable", "prediction")

    if not raw_predictions:
        logger.warning("No raw predictions found in state, skipping post-processing.")
        return state, pd.DataFrame()

    try:
        # Convert list of predictions to a Series for easier handling
        predictions_series = pd.Series(raw_predictions, index=original_indices)

        # Initialize post-processed DataFrame with index
        postprocessed_df = pd.DataFrame(index=original_indices)

        # Assume raw predictions are probabilities if between 0 and 1
        if all(0.0 <= p <= 1.0 for p in raw_predictions):
            postprocessed_df[f"{target_variable}_confidence"] = predictions_series
            postprocessed_df[f"{target_variable}_label"] = (predictions_series >= prediction_threshold).astype(int)
            logger.info(f"Converted probabilities to labels using threshold {prediction_threshold}.")
        else:
            # If not probabilities, treat as direct predictions (e.g., regression output or discrete classes)
            postprocessed_df[f"{target_variable}"] = predictions_series
            logger.info("Treated raw predictions as direct labels/values.")

        state["postprocessed_predictions"] = postprocessed_df # Store the DataFrame itself
        logger.info(f"Prediction post-processing complete. Generated {len(postprocessed_df)} post-processed records.")
        logger.debug(f"Post-processed predictions head:\n{postprocessed_df.head()}")

    except Exception as e:
        logger.error(f"Error during prediction post-processing: {e}", exc_info=True)
        state["errors"].append(f"Post-processing failed: {e}")
        return state, pd.DataFrame()

    return state, postprocessed_df

def store_inference_results(
    state: Dict[str, Any],
    postprocessed_df: pd.DataFrame
) -> Dict[str, Any]:
    """
    Stores the post-processed inference results to Google Cloud Storage.

    Saves the DataFrame of predictions to a specified GCS location, typically
    in Parquet or CSV format, with retry logic for robust storage.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary, must contain 'inference_results_path'
        and 'bucket_name'.
    postprocessed_df : pd.DataFrame
        The DataFrame containing the final post-processed predictions.

    Returns
    -------
    Dict[str, Any]
        An updated state dictionary, including the path where results were
        stored or any errors encountered.

    Examples
    --------
    >>> current_state = {
    ...     "bucket_name": "my-ml-bucket",
    ...     "inference_results_path": "gs://my-ml-bucket/inference_results/"
    ... }
    >>> final_predictions = pd.DataFrame({
    ...     "id": [1,2,3], "prediction_label": [0,1,0], "prediction_confidence": [0.1, 0.9, 0.3]
    ... })
    >>> # updated_state = store_inference_results(current_state, final_predictions)
    >>> # if 'stored_results_path' in updated_state: print("Results saved.")
    """
    inference_results_path = state.get("inference_results_path")
    bucket_name = state.get("bucket_name")

    if not all([inference_results_path, bucket_name]):
        logger.error("Missing inference_results_path or bucket_name in state for storing results.")
        state["errors"].append("Missing state parameters for storing results.")
        return state

    if postprocessed_df.empty:
        logger.warning("Post-processed DataFrame is empty, skipping result storage.")
        return state

    # Generate a unique filename for the results
    timestamp = time.strftime("%Y%m%d-%H%M%S")
    output_filename = f"inference_results_{timestamp}.parquet"
    full_output_path = os.path.join(inference_results_path, output_filename)

    logger.info(f"Attempting to store inference results to GCS: {full_output_path}")

    @retry(
        wait=wait_exponential(multiplier=1, min=4, max=60),
        stop=stop_after_attempt(5),
        retry=retry_if_exception_type((IOError, ConnectionError, Exception)),
        reraise=True
    )
    def _save_to_gcs():
        logger.debug(f"Retrying saving results to {full_output_path}...")
        time.sleep(random.uniform(0.1, 0.5))
        try:
            # Using pandas to_parquet for efficiency and schema preservation
            postprocessed_df.to_parquet(full_output_path, index=False)
            logger.info(f"Successfully saved inference results to {full_output_path}")
        except Exception as e:
            logger.error(f"Error saving results to GCS: {e}")
            raise # Re-raise for tenacity to catch

    try:
        _save_to_gcs()
        state["stored_results_path"] = full_output_path
    except Exception as e:
        logger.error(f"Failed to store inference results after multiple retries: {e}")
        state["errors"].append(f"Storing results failed: {e}")

    return state

# --- Demonstration Setup ---

logger.info("--- Starting Demonstration ---")

# 1. Simulate raw data for inference
# Let's create a dummy dataset for demonstration purposes
np.random.seed(42)
num_samples = 1000

dummy_data = pd.DataFrame({
    'feature_1': np.random.rand(num_samples) * 100,
    'feature_2': np.random.normal(loc=50, scale=10, size=num_samples),
    'feature_3': np.random.randint(0, 5, num_samples).astype(str), # Categorical feature
    'feature_4': np.random.rand(num_samples) * 10,
    'id': range(num_samples)
})

# Introduce some missing values for demonstration of preprocessing
dummy_data.loc[np.random.choice(num_samples, 50, replace=False), 'feature_1'] = np.nan
dummy_data.loc[np.random.choice(num_samples, 30, replace=False), 'feature_3'] = np.nan

logger.info(f"Simulated raw data with {len(dummy_data)} samples and some missing values.")

# 2. Define placeholder state variables
PROJECT_ID = "gcp-project-demo"
BUCKET_NAME = "mlops-inference-bucket"
MODEL_NAME = "classification-model"
MODEL_VERSION = "v1.0.0"
TARGET_VARIABLE = "risk_score"
FEATURE_COLUMNS = ['feature_1', 'feature_2', 'feature_3', 'feature_4']
BATCH_SIZE = 100

# 3. Initialize inference state
initial_state = create_inference_state(
    project_id=PROJECT_ID,
    bucket_name=BUCKET_NAME,
    model_name=MODEL_NAME,
    model_version=MODEL_VERSION,
    target_variable=TARGET_VARIABLE,
    feature_columns=FEATURE_COLUMNS,
    batch_size=BATCH_SIZE
)

logger.info("Inference state initialized for demonstration.")
print("Initial State Snippet:")
print(pd.Series({k: v for k, v in initial_state.items() if not isinstance(v, list) and not isinstance(v, dict)}).to_frame(name="Value"))

# Store the original dataframe to simulate loading it later
original_data_for_demo = dummy_data.copy()

# Simulate loading data (using our dummy data)
state, loaded_data_df = initial_state, original_data_for_demo
logger.info(f"Simulated loading {len(loaded_data_df)} records.")
print("\nLoaded Data Head (before preprocessing):")
print(loaded_data_df.head())
print("\nMissing values before preprocessing:")
print(loaded_data_df.isnull().sum())

# Preprocess the loaded data
state, processed_data_df = preprocess_data_for_inference(state, loaded_data_df)

if not processed_data_df.empty:
    logger.info("Data preprocessing completed.")
    print("\nProcessed Data Head:")
    print(processed_data_df.head())
    print("\nProcessed Data Info:")
    print(processed_data_df.info())
else:
    logger.error("Data preprocessing failed or resulted in an empty DataFrame.")

from sklearn.linear_model import LogisticRegression

# Simulate a dummy model (e.g., a simple Logistic Regression)
logger.info("Simulating a dummy model for demonstration.")

# Create some dummy target values for fitting the dummy model
# In a real scenario, the model would already be trained.
# Here, we're just creating a placeholder for the `model` object.

dummy_X = processed_data_df.values # Features for dummy model
dummy_y = np.random.randint(0, 2, size=len(processed_data_df)) # Binary target

dummy_model = LogisticRegression(random_state=42)
dummy_model.fit(dummy_X, dummy_y)

# Add the dummy model to the state for demonstration purposes
# In a real scenario, load_model_for_inference would fetch it from GCS.
state['model'] = dummy_model
logger.info("Dummy model created and added to state.")

# Run batch inference
state, predictions_series = run_batch_inference(state, processed_data_df)

if not predictions_series.empty:
    logger.info(f"Inference completed. Generated {len(predictions_series)} predictions.")
    print("\nPredictions Snippet (raw output):")
    print(predictions_series.head())
else:
    logger.error("Batch inference failed or resulted in no predictions.")

# Post-process predictions
state, final_predictions_df = postprocess_predictions(state, original_data_for_demo.index, prediction_threshold=0.5)

if not final_predictions_df.empty:
    logger.info("Predictions post-processing completed.")
    print("\nFinal Post-processed Predictions Head:")
    print(final_predictions_df.head())

    # Add original 'id' column back for context in results
    final_predictions_df = pd.concat([original_data_for_demo['id'], final_predictions_df], axis=1)
    print("\nFinal Post-processed Predictions with Original ID Head:")
    print(final_predictions_df.head())

    # Display some basic statistics of predictions
    print("\nPrediction Label Distribution:")
    print(final_predictions_df[f'{TARGET_VARIABLE}_label'].value_counts())

    print("\nPrediction Confidence Statistics:")
    print(final_predictions_df[f'{TARGET_VARIABLE}_confidence'].describe().to_frame(name="Value"))

else:
    logger.error("Prediction post-processing failed or resulted in an empty DataFrame.")

# Simulate storing inference results
# We'll override the GCS path in state to a local path for demonstration
state["inference_results_path"] = "/tmp/inference_results/"

# Ensure the local directory exists
os.makedirs(state["inference_results_path"], exist_ok=True)

updated_state = store_inference_results(state, final_predictions_df)

if "stored_results_path" in updated_state:
    logger.info(f"Simulated storage completed. Results saved to: {updated_state['stored_results_path']}")
    print(f"\nResults are theoretically stored at: {updated_state['stored_results_path']}")
    # Verify by loading the locally saved file
    loaded_results = pd.read_parquet(updated_state['stored_results_path'])
    print("\nLoaded results from simulated storage head:")
    print(loaded_results.head())
else:
    logger.error("Simulated storage failed.")

logger.info("Generating visualizations for inference results.")

if not final_predictions_df.empty:
    plt.figure(figsize=(14, 6))

    # Plot 1: Distribution of Prediction Confidence
    plt.subplot(1, 2, 1)
    sns.histplot(final_predictions_df[f'{TARGET_VARIABLE}_confidence'], bins=20, kde=True)
    plt.title('Distribution of Prediction Confidence')
    plt.xlabel('Confidence Score')
    plt.ylabel('Frequency')
    plt.grid(axis='y', alpha=0.75)

    # Plot 2: Count of Predicted Labels
    plt.subplot(1, 2, 2)
    sns.countplot(x=f'{TARGET_VARIABLE}_label', data=final_predictions_df, palette='viridis')
    plt.title('Count of Predicted Labels')
    plt.xlabel('Predicted Label')
    plt.ylabel('Count')
    plt.xticks([0, 1], ['Class 0', 'Class 1'])
    plt.grid(axis='y', alpha=0.75)

    plt.tight_layout()
    plt.show()

    # Display summary statistics of the original features
    print("\nSummary Statistics of Original Features:")
    print(original_data_for_demo[FEATURE_COLUMNS].describe().T)

    # Display comparison: feature_1 vs confidence for different labels
    plt.figure(figsize=(10, 6))
    sns.scatterplot(
        x='feature_1', y=f'{TARGET_VARIABLE}_confidence', hue=f'{TARGET_VARIABLE}_label',
        data=pd.concat([original_data_for_demo, final_predictions_df], axis=1).dropna(subset=['feature_1']),
        palette='coolwarm', alpha=0.6
    )
    plt.title('Feature 1 vs. Prediction Confidence by Label')
    plt.xlabel('Feature 1 Value')
    plt.ylabel('Prediction Confidence')
    plt.legend(title='Predicted Label')
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.show()

    logger.info("Visualizations generated.")
else:
    logger.warning("No data to visualize as final predictions DataFrame is empty.")

logger.info("--- Demonstration Finished ---")