import os
os.environ["MLFLOW_ALLOW_FILE_STORE"] = "true"
print("MLFLOW_ALLOW_FILE_STORE environment variable set to 'true'.")

# Install MLflow and other necessary libraries
!pip install mlflow scikit-learn pandas matplotlib seaborn --quiet

import logging

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

# Standard library imports
import os
import random
import time
from collections import deque
import logging

# Set environment variable to allow MLflow filesystem backend
os.environ["MLFLOW_ALLOW_FILE_STORE"] = "true"

# Third-party library imports
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
import mlflow
import mlflow.sklearn

logging.info("All necessary libraries imported and MLFLOW_ALLOW_FILE_STORE set.")

def create_mlflow_experiment_state(
    experiment_name: str,
    run_name_prefix: str,
    artifact_location: str = "./mlruns",
    state: dict = None          # FIX (Bug 1): accept existing state so callers data is preserved
) -> dict:
    """
    Initializes (or updates) a state dictionary with MLflow experiment configuration.

    Parameters
    ----------
    experiment_name : str
        The name of the MLflow experiment.
    run_name_prefix : str
        A prefix for the MLflow run name to ensure uniqueness.
    artifact_location : str, optional
        The URI to which artifacts will be logged, defaults to './mlruns'.
    state : dict, optional
        An existing state dictionary to update.  If None a new dict is created.

    Returns
    -------
    dict
        Updated state dictionary containing MLflow configuration.

    Examples
    --------
    >>> state = create_mlflow_experiment_state("MyClassificationExperiment", "run_", "./mlruns_data")
    >>> print(state['mlflow_config']['experiment_name'])
    MyClassificationExperiment
    """
    # FIX (Bug 1): Start from the caller-supplied state (or a fresh dict) so that
    # keys already in state (X, y, X_train …) are NOT discarded.
    if state is None:
        state = {}

    # FIX (Bug 2): Normalise artifact_location to a proper file:// URI so that
    # mlflow.set_tracking_uri() works correctly across all MLflow versions.
    # Plain relative paths like './mlruns' are ambiguous; an explicit URI is not.
    import os as _os
    if not artifact_location.startswith(("file://", "http://", "https://", "sqlite")):
        artifact_location = "file://" + _os.path.abspath(artifact_location)

    state["mlflow_config"] = {
        "experiment_name": experiment_name,
        "run_name": f"{run_name_prefix}{int(time.time())}",
        "artifact_location": artifact_location,
    }
    logging.info(f"MLflow experiment state initialized for experiment '{experiment_name}'.")
    return state

def create_synthetic_data(
    state: dict,
    n_samples: int,
    n_features: int,
    n_informative: int,
    random_state: int = 42
) -> dict:
    """
    Generates synthetic classification data.

    Parameters
    ----------
    state : dict
        Current state dictionary.
    n_samples : int
        The number of samples to generate.
    n_features : int
        The total number of features.
    n_informative : int
        The number of informative features (must be <= n_features).
    random_state : int, optional
        Controls the reproducibility of the dataset generation, defaults to 42.

    Returns
    -------
    dict
        Updated state with 'X' (features) and 'y' (labels).

    Examples
    --------
    >>> state = {}
    >>> state = create_synthetic_data(state, 100, 5, 3)
    >>> print(state['X'].shape, state['y'].shape)
    (100, 5) (100,)
    """
    # FIX (Bug 8): Validate that n_informative does not exceed n_features.
    # sklearn raises a cryptic ValueError otherwise; a clear message is better.
    if n_informative > n_features:
        raise ValueError(
            f"n_informative ({n_informative}) must be <= n_features ({n_features})."
        )
    if n_samples <= 0:
        raise ValueError(f"n_samples must be a positive integer, got {n_samples}.")

    logging.info(f"Generating synthetic data with {n_samples} samples, {n_features} features, {n_informative} informative features.")
    X, y = make_classification(
        n_samples=n_samples,
        n_features=n_features,
        n_informative=n_informative,
        random_state=random_state,
    )
    state['X'] = X
    state['y'] = y
    logging.debug("Synthetic data generation complete.")
    return state

def split_data(
    state: dict,
    test_size: float = 0.2,
    random_state: int = 42
) -> dict:
    """
    Splits the data into training and testing sets.

    Parameters
    ----------
    state : dict
        Current state dictionary containing 'X' and 'y'.
    test_size : float, optional
        The proportion of the dataset to include in the test split, defaults to 0.2.
    random_state : int, optional
        Controls the shuffling applied to the data before splitting, defaults to 42.

    Returns
    -------
    dict
        Updated state with 'X_train', 'X_test', 'y_train', 'y_test'.

    Examples
    --------
    >>> state = {'X': np.array([[1,2],[3,4],[5,6],[7,8]]), 'y': np.array([0,1,0,1])}
    >>> state = split_data(state, test_size=0.5, random_state=1)
    >>> print(state['X_train'].shape, state['X_test'].shape)
    (2, 2) (2, 2)
    """
    logging.info(f"Splitting data into training and test sets with test_size={test_size}.")
    X_train, X_test, y_train, y_test = train_test_split(state['X'], state['y'], test_size=test_size, random_state=random_state)
    state['X_train'] = X_train
    state['X_test'] = X_test
    state['y_train'] = y_train
    state['y_test'] = y_test
    logging.debug("Data splitting complete.")
    return state

def train_model(
    state: dict,
    model_params: dict,
    max_retries: int = 3
) -> dict:
    """
    Trains a Logistic Regression model and logs to MLflow.

    Parameters
    ----------
    state : dict
        Current state dictionary containing 'X_train', 'y_train', and MLflow config.
    model_params : dict
        A dictionary of parameters for the Logistic Regression model.
    max_retries : int, optional
        Maximum number of retries for MLflow logging operations, defaults to 3.

    Returns
    -------
    dict
        Updated state including the trained 'model'.
    """
    # FIX (Bug 10): use .get() so a missing key raises a clear message instead
    # of an opaque KeyError propagating from deep inside the retry loop.
    X_train = state.get('X_train')
    y_train = state.get('y_train')
    if X_train is None or y_train is None:
        raise ValueError("state must contain 'X_train' and 'y_train'. Run split_data() first.")

    logging.info("Starting model training and MLflow logging.")

    # FIX (Bug 3): Track whether params have already been logged so that a retry
    # does not attempt to re-log them (which raises MlflowException if a value
    # changes, or produces duplicate entries).  Train the model FIRST so that
    # log_params is only reached after a successful fit.
    params_logged = False
    last_exc = None

    for attempt in range(max_retries):
        try:
            model = LogisticRegression(**model_params, random_state=state.get('random_state', 42))
            model.fit(X_train, y_train)
            state['model'] = model

            # Log params and model only after a successful fit.
            if not params_logged:
                mlflow.log_params(model_params)
                params_logged = True
                logging.info(f"Logged model parameters: {model_params}")

            mlflow.sklearn.log_model(sk_model=model, artifact_path="model")
            logging.info("Logged model artifact.")
            break  # success — exit retry loop

        except Exception as e:
            last_exc = e
            # FIX (Bug 7): skip sleep on the last attempt to avoid pointless delay.
            if attempt < max_retries - 1:
                jitter = random.uniform(0.1, 0.5)
                wait_time = (2 ** attempt) + jitter
                logging.warning(
                    f"Attempt {attempt + 1} failed: {e}. Retrying in {wait_time:.2f} seconds..."
                )
                time.sleep(wait_time)
            else:
                logging.warning(f"Attempt {attempt + 1} failed: {e}. No more retries.")
    else:
        logging.error(f"Failed to train and log model after {max_retries} attempts.")
        raise Exception(
            f"Failed to train and log model after {max_retries} attempts."
        ) from last_exc

    logging.debug("Model training and logging complete.")
    return state

def evaluate_model(
    state: dict,
    max_retries: int = 3
) -> dict:
    """
    Evaluates the trained model and logs metrics to MLflow.

    Parameters
    ----------
    state : dict
        Current state dictionary containing 'model', 'X_test', 'y_test', and MLflow config.
    max_retries : int, optional
        Maximum number of retries for MLflow logging operations, defaults to 3.

    Returns
    -------
    dict
        Updated state with 'metrics' (a dictionary of calculated metrics).
    """
    logging.info("Starting model evaluation and MLflow metric logging.")

    # FIX (Bug 10): use .get() for safe access with a clear error message.
    model  = state.get('model')
    X_test = state.get('X_test')
    y_test = state.get('y_test')
    if model is None or X_test is None or y_test is None:
        raise ValueError("state must contain 'model', 'X_test', and 'y_test'.")

    y_pred  = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:, 1]

    metrics = {
        "accuracy":  accuracy_score(y_test, y_pred),
        "precision": precision_score(y_test, y_pred, zero_division=0),
        "recall":    recall_score(y_test, y_pred, zero_division=0),
        "f1_score":  f1_score(y_test, y_pred, zero_division=0),
        "roc_auc":   roc_auc_score(y_test, y_proba),
    }
    state['metrics'] = metrics

    last_exc = None
    for attempt in range(max_retries):
        try:
            mlflow.log_metrics(metrics)
            logging.info(f"Logged evaluation metrics: {metrics}")
            break
        except Exception as e:
            last_exc = e
            # FIX (Bug 7): no sleep on the last attempt.
            if attempt < max_retries - 1:
                jitter = random.uniform(0.1, 0.5)
                wait_time = (2 ** attempt) + jitter
                logging.warning(
                    f"Attempt {attempt + 1} failed: {e}. Retrying in {wait_time:.2f} seconds..."
                )
                time.sleep(wait_time)
            else:
                logging.warning(f"Attempt {attempt + 1} failed: {e}. No more retries.")
    else:
        logging.error(f"Failed to log metrics after {max_retries} attempts.")
        raise Exception(f"Failed to log metrics after {max_retries} attempts.") from last_exc

    logging.debug("Model evaluation and metric logging complete.")
    return state

def log_feature_importance_plot(
    state: dict,
    feature_names: list,
    plot_filename: str = 'feature_importance.png',
    max_retries: int = 3
) -> dict:
    """
    Creates and logs a feature importance plot as an MLflow artifact.

    Parameters
    ----------
    state : dict
        Current state dictionary containing 'model'.
    feature_names : list[str]
        A list of names for the features.  Must match the number of model coefficients.
    plot_filename : str, optional
        The filename for the plot artifact, defaults to 'feature_importance.png'.
    max_retries : int, optional
        Maximum number of retries for MLflow logging operations, defaults to 3.

    Returns
    -------
    dict
        The unmodified state dictionary.
    """
    logging.info("Generating and logging feature importance plot.")

    model = state.get('model')
    if model is None:
        raise ValueError("state must contain a trained 'model'.")

    coefficients = model.coef_[0]

    # FIX (Bug 11): Validate that feature_names length matches the number of
    # coefficients.  A mismatch would silently truncate the DataFrame in older
    # pandas or raise a confusing ValueError in newer pandas.
    if len(feature_names) != len(coefficients):
        raise ValueError(
            f"feature_names length ({len(feature_names)}) does not match "
            f"model coefficient length ({len(coefficients)})."
        )

    feature_importance_df = pd.DataFrame({
        'Feature':    feature_names,
        'Importance': np.abs(coefficients),
    }).sort_values(by='Importance', ascending=False)

    plt.figure(figsize=(10, 6))
    sns.barplot(x='Importance', y='Feature', data=feature_importance_df)
    plt.title('Feature Importances (Absolute Coefficients)')
    plt.xlabel('Absolute Coefficient Value')
    plt.ylabel('Feature')
    plt.tight_layout()
    plt.savefig(plot_filename)
    plt.close()

    last_exc = None
    try:
        for attempt in range(max_retries):
            try:
                mlflow.log_artifact(plot_filename)
                logging.info(f"Logged plot '{plot_filename}' as artifact.")
                break
            except Exception as e:
                last_exc = e
                # FIX (Bug 7): skip sleep on the last attempt.
                if attempt < max_retries - 1:
                    jitter = random.uniform(0.1, 0.5)
                    wait_time = (2 ** attempt) + jitter
                    logging.warning(
                        f"Attempt {attempt + 1} failed: {e}. Retrying in {wait_time:.2f} seconds..."
                    )
                    time.sleep(wait_time)
                else:
                    logging.warning(f"Attempt {attempt + 1} failed: {e}. No more retries.")
        else:
            logging.error(f"Failed to log plot after {max_retries} attempts.")
            raise Exception(f"Failed to log plot after {max_retries} attempts.") from last_exc
    finally:
        # FIX (Bug 5): Always remove the local temp file after attempting to log
        # it, whether logging succeeded or failed.  This prevents stale PNG files
        # from accumulating in the working directory across runs.
        import os as _os
        if _os.path.exists(plot_filename):
            _os.remove(plot_filename)
            logging.debug(f"Removed local temp file '{plot_filename}'.")

    logging.debug("Feature importance plot logging complete.")
    return state

def create_mlflow_run(
    state: dict,
    run_callback: callable,
    feature_names: list
) -> dict:
    """
    Orchestrates an MLflow run by setting the tracking URI, creating/setting
    an experiment, and executing a callback function within the run context.

    Parameters
    ----------
    state : dict
        Current state dictionary containing MLflow configuration.
    run_callback : callable
        A function that takes the current state and `feature_names` as arguments.
    feature_names : list[str]
        A list of feature names to be passed to the `run_callback`.

    Returns
    -------
    dict
        The state dictionary after the run_callback has been executed.
    """
    mlflow_config = state['mlflow_config']

    mlflow.set_tracking_uri(mlflow_config['artifact_location'])
    logging.info(f"MLflow Tracking URI set to: {mlflow_config['artifact_location']}")

    mlflow.set_experiment(mlflow_config['experiment_name'])
    logging.info(f"MLflow Experiment set to: {mlflow_config['experiment_name']}")

    with mlflow.start_run(run_name=mlflow_config['run_name']) as run:
        # FIX (Bug 4): Capture run_id from the context-manager `run` object, NOT
        # from mlflow.active_run().  After the `with` block ends, active_run()
        # returns None and calling .info.run_id raises AttributeError.
        run_id = run.info.run_id
        logging.info(f"MLflow Run started: {run_id} (Name: {mlflow_config['run_name']})")
        state['mlflow_run_id'] = run_id
        state = run_callback(state, feature_names)
        # Use the already-captured run_id — do NOT call mlflow.active_run() here.
        logging.info(f"MLflow Run finished: {run_id}")

    return state

state = {}
state['random_state'] = 42

# 1. Initialize MLflow experiment state
# FIX (Bug 1): pass the existing state so subsequent keys (X, y, …) are preserved.
state = create_mlflow_experiment_state(
    experiment_name="LogisticRegression_Classification",
    run_name_prefix="lr_run_",
    artifact_location="./mlruns",
    state=state,
)

# 2. Generate synthetic data
n_samples    = 1000
n_features   = 10
n_informative = 5
state = create_synthetic_data(state, n_samples, n_features, n_informative, state['random_state'])

# Prepare feature names for logging plots
feature_names = [f'feature_{i}' for i in range(n_features)]

# 3. Split the data
state = split_data(state, test_size=0.2, random_state=state['random_state'])

logging.info("Data prepared for ML experiments.")
display(pd.DataFrame(state['X']).head())
display(pd.Series(state['y']).value_counts().to_frame('Label Count'))

def train_evaluate_log_callback(current_state: dict, feature_names: list[str]) -> dict:
    """
    Callback function to encapsulate the training, evaluation, and logging logic.
    """
    model_params = {"C": 1.0, "solver": "liblinear"}
    current_state = train_model(current_state, model_params)
    current_state = evaluate_model(current_state)
    current_state = log_feature_importance_plot(current_state, feature_names, 'feature_importance_run1.png')
    return current_state

logging.info("Starting Experiment Run 1 with default Logistic Regression parameters.")
state['mlflow_config']['run_name'] = f"lr_run_{int(time.time())}_default_C_1.0"
state = create_mlflow_run(state, train_evaluate_log_callback, feature_names)
logging.info(f"Experiment Run 1 completed. MLflow Run ID: {state['mlflow_run_id']}")
print(f"Run 1 Metrics: {state['metrics']}")

def train_evaluate_log_callback_tuned(current_state: dict, feature_names: list[str]) -> dict:
    """
    Callback function for the second run with tuned parameters.
    """
    model_params = {"C": 0.1, "solver": "liblinear"} # Lower C for stronger regularization
    current_state = train_model(current_state, model_params)
    current_state = evaluate_model(current_state)
    current_state = log_feature_importance_plot(current_state, feature_names, 'feature_importance_run2.png')
    return current_state

logging.info("Starting Experiment Run 2 with tuned Logistic Regression parameters (C=0.1).")
state['mlflow_config']['run_name'] = f"lr_run_{int(time.time())}_tuned_C_0.1"
state = create_mlflow_run(state, train_evaluate_log_callback_tuned, feature_names)
logging.info(f"Experiment Run 2 completed. MLflow Run ID: {state['mlflow_run_id']}")
print(f"Run 2 Metrics: {state['metrics']}")

import pandas as pd

def query_mlflow_runs(experiment_name: str, artifact_location: str) -> pd.DataFrame:
    """
    Queries MLflow runs for a given experiment and returns a DataFrame of results.

    Parameters
    ----------
    experiment_name : str
        The name of the MLflow experiment to query.
    artifact_location : str
        The MLflow tracking URI.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing parameters, metrics, and run IDs of the experiment runs.
    """
    # FIX (Bug 9): Warn if an active run exists before switching the tracking URI.
    # Changing the URI mid-run would silently redirect subsequent log_* calls.
    if mlflow.active_run() is not None:
        logging.warning(
            "query_mlflow_runs called while an MLflow run is active. "
            "Changing the tracking URI now may redirect active logging."
        )

    mlflow.set_tracking_uri(artifact_location)
    client = mlflow.tracking.MlflowClient()

    experiment = client.get_experiment_by_name(experiment_name)
    if experiment is None:
        logging.warning(f"Experiment '{experiment_name}' not found.")
        return pd.DataFrame()

    runs_info = []
    for run in client.search_runs(experiment_ids=[experiment.experiment_id]):
        run_data = {
            "run_id":   run.info.run_id,
            "run_name": run.info.run_name,
            **run.data.params,   # logged parameters
            **run.data.metrics,  # logged metrics
        }
        runs_info.append(run_data)

    df_runs = pd.DataFrame(runs_info)
    logging.info(f"Queried {len(df_runs)} runs for experiment '{experiment_name}'.")
    return df_runs


# Query and display the results
experiment_name   = state['mlflow_config']['experiment_name']
artifact_location = state['mlflow_config']['artifact_location']

all_runs_df = query_mlflow_runs(experiment_name, artifact_location)

# FIX (Bug 6): Guard against empty DataFrame or missing 'accuracy' column
# (e.g. if all runs failed before logging metrics) to avoid a KeyError.
if not all_runs_df.empty and 'accuracy' in all_runs_df.columns:
    display(all_runs_df.sort_values(by='accuracy', ascending=False))
else:
    logging.warning("No runs with 'accuracy' metric found. Displaying unsorted results.")
    display(all_runs_df)