#!/usr/bin/env python3
import os
import sys

from subprocess import check_call
import logging

# ============================================================================
# PACKAGE INSTALLATION CONFIGURATION
# ============================================================================

# Control which PyPI source to use via environment variable
USE_SECURE_PYPI = os.environ.get("USE_SECURE_PYPI", "true").lower() == "true"

# Logging setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def _get_secure_pypi_access_token() -> str:
    """Get CodeArtifact access token for secure PyPI."""
    import boto3

    try:
        os.environ["AWS_STS_REGIONAL_ENDPOINTS"] = "regional"
        sts = boto3.client("sts", region_name="us-east-1")
        caller_identity = sts.get_caller_identity()
        assumed_role_object = sts.assume_role(
            RoleArn="arn:aws:iam::675292366480:role/SecurePyPIReadRole_"
            + caller_identity["Account"],
            RoleSessionName="SecurePypiReadRole",
        )
        credentials = assumed_role_object["Credentials"]
        code_artifact_client = boto3.client(
            "codeartifact",
            aws_access_key_id=credentials["AccessKeyId"],
            aws_secret_access_key=credentials["SecretAccessKey"],
            aws_session_token=credentials["SessionToken"],
            region_name="us-west-2",
        )
        token = code_artifact_client.get_authorization_token(
            domain="amazon", domainOwner="149122183214"
        )["authorizationToken"]
        logger.info("Successfully retrieved secure PyPI access token")
        return token
    except Exception as e:
        logger.error(f"Failed to retrieve secure PyPI access token: {e}")
        raise


def install_packages_from_public_pypi(packages: list) -> None:
    """Install packages from standard public PyPI."""
    logger.info(f"Installing {len(packages)} packages from public PyPI")
    logger.info(f"Packages: {packages}")
    try:
        check_call([sys.executable, "-m", "pip", "install", *packages])
        logger.info("✓ Successfully installed packages from public PyPI")
    except Exception as e:
        logger.error(f"✗ Failed to install packages from public PyPI: {e}")
        raise


def install_packages_from_secure_pypi(packages: list) -> None:
    """Install packages from secure CodeArtifact PyPI."""
    logger.info(f"Installing {len(packages)} packages from secure PyPI")
    logger.info(f"Packages: {packages}")
    try:
        token = _get_secure_pypi_access_token()
        index_url = f"https://aws:{token}@amazon-149122183214.d.codeartifact.us-west-2.amazonaws.com/pypi/secure-pypi/simple/"
        check_call(
            [
                sys.executable,
                "-m",
                "pip",
                "install",
                "--index-url",
                index_url,
                *packages,
            ]
        )
        logger.info("✓ Successfully installed packages from secure PyPI")
    except Exception as e:
        logger.error(f"✗ Failed to install packages from secure PyPI: {e}")
        raise


def install_packages(packages: list, use_secure: bool = USE_SECURE_PYPI) -> None:
    """Install packages from PyPI source based on configuration."""
    logger.info("=" * 70)
    logger.info("PACKAGE INSTALLATION")
    logger.info("=" * 70)
    logger.info(f"PyPI Source: {'SECURE (CodeArtifact)' if use_secure else 'PUBLIC'}")
    logger.info(
        f"Environment Variable USE_SECURE_PYPI: {os.environ.get('USE_SECURE_PYPI', 'not set')}"
    )
    logger.info(f"Number of packages: {len(packages)}")
    logger.info("=" * 70)
    try:
        if use_secure:
            install_packages_from_secure_pypi(packages)
        else:
            install_packages_from_public_pypi(packages)
        logger.info("=" * 70)
        logger.info("✓ PACKAGE INSTALLATION COMPLETED SUCCESSFULLY")
        logger.info("=" * 70)
    except Exception as e:
        logger.error("=" * 70)
        logger.error("✗ PACKAGE INSTALLATION FAILED")
        logger.error("=" * 70)
        raise


# ============================================================================
# INSTALL REQUIRED PACKAGES
# ============================================================================

required_packages = [
    "lightgbm>=3.3.0,<4.0.0",
    "matplotlib>=3.3.0,<3.7.0",
]

install_packages(required_packages)

print("***********************Package Installation Complete*********************")

import json
import argparse
import pandas as pd
import numpy as np
import pickle as pkl
from pathlib import Path
from sklearn.metrics import (
    roc_auc_score,
    average_precision_score,
    precision_recall_curve,
    roc_curve,
    f1_score,
)
from scipy import stats
from scipy.stats import pearsonr, spearmanr
import lightgbm as lgb  # Import LightGBM instead of XGBoost
import matplotlib.pyplot as plt
import time
import sys
import tarfile
from datetime import datetime
from typing import Dict, Any, Optional, List, Tuple, Union

# Embedded processor classes to remove external dependencies


class RiskTableMappingProcessor:
    """
    A processor that performs risk-table-based mapping on a specified categorical variable.
    The 'process' method (called via __call__) handles single values.
    The 'transform' method handles pandas Series or DataFrames.
    """

    def __init__(
        self,
        column_name: str,
        label_name: str,
        smooth_factor: float = 0.0,
        count_threshold: int = 0,
        risk_tables: Optional[Dict] = None,
    ):
        """
        Initialize RiskTableMappingProcessor.

        Args:
            column_name: Name of the categorical column to be binned.
            label_name: Name of label/target variable (expected to be binary 0 or 1).
            smooth_factor: Smoothing factor for risk calculation (0 to 1).
            count_threshold: Minimum count for considering a category's calculated risk.
            risk_tables: Optional pre-computed risk tables.
        """
        self.processor_name = "risk_table_mapping_processor"
        self.function_name_list = ["process", "transform", "fit"]

        if not isinstance(column_name, str) or not column_name:
            raise ValueError("column_name must be a non-empty string.")
        self.column_name = column_name
        self.label_name = label_name
        self.smooth_factor = smooth_factor
        self.count_threshold = count_threshold

        self.is_fitted = False
        if risk_tables:
            self._validate_risk_tables(risk_tables)
            self.risk_tables = risk_tables
            self.is_fitted = True
        else:
            self.risk_tables = {}

    def get_name(self) -> str:
        return self.processor_name

    def _validate_risk_tables(self, risk_tables: Dict) -> None:
        if not isinstance(risk_tables, dict):
            raise ValueError("Risk tables must be a dictionary.")
        if "bins" not in risk_tables or "default_bin" not in risk_tables:
            raise ValueError("Risk tables must contain 'bins' and 'default_bin' keys.")
        if not isinstance(risk_tables["bins"], dict):
            raise ValueError("Risk tables 'bins' must be a dictionary.")
        if not isinstance(
            risk_tables["default_bin"], (int, float, np.floating, np.integer)
        ):
            raise ValueError(
                f"Risk tables 'default_bin' must be a number, got {type(risk_tables['default_bin'])}."
            )

    def set_risk_tables(self, risk_tables: Dict) -> None:
        self._validate_risk_tables(risk_tables)
        self.risk_tables = risk_tables
        self.is_fitted = True

    def fit(self, data: pd.DataFrame) -> "RiskTableMappingProcessor":
        if not isinstance(data, pd.DataFrame):
            raise TypeError("fit() requires a pandas DataFrame.")
        if self.label_name not in data.columns:
            raise ValueError(
                f"Label variable '{self.label_name}' not found in input data."
            )
        if self.column_name not in data.columns:
            raise ValueError(
                f"Column to bin '{self.column_name}' not found in input data."
            )

        filtered_data = data[data[self.label_name] != -1].dropna(
            subset=[self.label_name, self.column_name]
        )

        if filtered_data.empty:
            # Handle case with no valid data for fitting
            print(
                f"Warning: Filtered data for column '{self.column_name}' is empty during fit. "
                "Risk tables will be empty, default_bin will be 0.5 or NaN if no labels at all."
            )
            # Attempt to get a global mean if any data existed before filtering for column_name
            overall_label_mean = data[self.label_name][
                data[self.label_name] != -1
            ].mean()
            self.risk_tables = {
                "bins": {},
                "default_bin": (
                    0.5 if pd.isna(overall_label_mean) else float(overall_label_mean)
                ),
            }
            self.is_fitted = True
            return self

        default_risk = float(filtered_data[self.label_name].mean())
        smooth_samples = int(len(filtered_data) * self.smooth_factor)

        cross_tab_result = pd.crosstab(
            index=filtered_data[self.column_name].astype(str),
            columns=filtered_data[self.label_name].astype(int),
            margins=True,
            margins_name="_count_",
            dropna=False,
        )

        positive_label_col = 1
        negative_label_col = 0

        if positive_label_col not in cross_tab_result.columns:
            cross_tab_result[positive_label_col] = 0
        if negative_label_col not in cross_tab_result.columns:
            cross_tab_result[negative_label_col] = 0

        calc_df = cross_tab_result[cross_tab_result.index != "_count_"].copy()

        if calc_df.empty:
            self.risk_tables = {"bins": {}, "default_bin": default_risk}
            self.is_fitted = True
            return self

        calc_df["risk"] = calc_df.apply(
            lambda x: (
                x[positive_label_col] / (x[positive_label_col] + x[negative_label_col])
                if (x[positive_label_col] + x[negative_label_col]) > 0
                else 0.0
            ),
            axis=1,
        )

        calc_df["_category_count_"] = cross_tab_result.loc[calc_df.index, "_count_"]

        calc_df["smooth_risk"] = calc_df.apply(
            lambda x: (
                (x["_category_count_"] * x["risk"] + smooth_samples * default_risk)
                / (x["_category_count_"] + smooth_samples)
                if (
                    x["_category_count_"] >= self.count_threshold
                    and (x["_category_count_"] + smooth_samples) > 0
                )
                else default_risk
            ),
            axis=1,
        )

        self.risk_tables = {
            "bins": dict(zip(calc_df.index.astype(str), calc_df["smooth_risk"])),
            "default_bin": default_risk,
        }

        self.is_fitted = True
        return self

    def process(self, input_value: Any) -> float:
        """
        Process a single input value (for the configured 'column_name'),
        mapping it to its binned risk value.
        This method is called when the processor instance is called as a function.
        """
        if not self.is_fitted:
            raise RuntimeError(
                "RiskTableMappingProcessor must be fitted or initialized with risk tables before processing."
            )
        str_value = str(input_value)
        return self.risk_tables["bins"].get(str_value, self.risk_tables["default_bin"])

    def transform(
        self, data: Union[pd.DataFrame, pd.Series, Any]
    ) -> Union[pd.DataFrame, pd.Series, float]:
        """
        Transform data using the computed risk tables.
        - If data is a DataFrame, transforms the 'column_name' Series within it.
        - If data is a Series, transforms the Series (assumed to be the target column).
        - If data is a single value, uses the 'process' method.
        """
        if not self.is_fitted:
            raise RuntimeError(
                "RiskTableMappingProcessor must be fitted or initialized with risk tables before transforming."
            )

        if isinstance(data, pd.DataFrame):
            if self.column_name not in data.columns:
                raise ValueError(
                    f"Column '{self.column_name}' not found in input DataFrame for transform operation."
                )
            output_data = data.copy()
            output_data[self.column_name] = (
                data[self.column_name]
                .astype(str)
                .map(self.risk_tables["bins"])
                .fillna(self.risk_tables["default_bin"])
            )
            return output_data
        elif isinstance(data, pd.Series):
            return (
                data.astype(str)
                .map(self.risk_tables["bins"])
                .fillna(self.risk_tables["default_bin"])
            )
        else:
            return self.process(data)  # Consistent with __call__

    def get_risk_tables(self) -> Dict:
        if not self.is_fitted:
            raise RuntimeError(
                "RiskTableMappingProcessor has not been fitted or initialized with risk tables."
            )
        return self.risk_tables


class NumericalVariableImputationProcessor:
    """
    A processor that performs imputation on numerical variables using predefined or computed values.
    Supports mean, median, and mode imputation strategies.
    """

    def __init__(
        self,
        variables: Optional[List[str]] = None,
        imputation_dict: Optional[Dict[str, Union[int, float]]] = None,
        strategy: str = "mean",
    ):
        self.processor_name = "numerical_variable_imputation_processor"
        self.function_name_list = ["fit", "process", "transform"]

        self.variables = variables
        self.strategy = strategy
        self.is_fitted = False

        if imputation_dict:
            self._validate_imputation_dict(imputation_dict)
            self.imputation_dict = imputation_dict
            self.is_fitted = True
        else:
            self.imputation_dict = None

    def get_name(self) -> str:
        return self.processor_name

    def __call__(self, input_data):
        return self.process(input_data)

    def _validate_imputation_dict(self, imputation_dict: Dict[str, Any]) -> None:
        if not isinstance(imputation_dict, dict):
            raise ValueError("imputation_dict must be a dictionary")
        if not imputation_dict:
            raise ValueError("imputation_dict cannot be empty")
        for k, v in imputation_dict.items():
            if not isinstance(k, str):
                raise ValueError(f"All keys must be strings, got {type(k)} for key {k}")
            if not isinstance(v, (int, float, np.number)):
                raise ValueError(
                    f"All values must be numeric, got {type(v)} for key {k}"
                )

    def fit(
        self, X: pd.DataFrame, y: Optional[pd.Series] = None
    ) -> "NumericalVariableImputationProcessor":
        if self.imputation_dict is None:
            self.imputation_dict = {}

            if self.variables is None:
                self.variables = X.select_dtypes(include=np.number).columns.tolist()

            for var in self.variables:
                if var not in X.columns:
                    raise ValueError(f"Variable {var} not found in the input data")

                if X[var].isna().all():
                    self.imputation_dict[var] = np.nan
                    continue

                if self.strategy == "mean":
                    self.imputation_dict[var] = float(X[var].mean())
                elif self.strategy == "median":
                    self.imputation_dict[var] = float(X[var].median())
                elif self.strategy == "mode":
                    self.imputation_dict[var] = float(X[var].mode()[0])
                else:
                    raise ValueError(f"Unknown strategy: {self.strategy}")

        self._validate_imputation_dict(self.imputation_dict)
        self.is_fitted = True
        return self

    def process(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        if not self.is_fitted:
            raise RuntimeError(
                "Processor is not fitted. Call 'fit' with appropriate arguments before using this method."
            )

        output_data = input_data.copy()
        for var, value in input_data.items():
            if var in self.imputation_dict and pd.isna(value):
                output_data[var] = self.imputation_dict[var]
        return output_data

    def transform(
        self, X: Union[pd.DataFrame, pd.Series]
    ) -> Union[pd.DataFrame, pd.Series]:
        """
        Transform input data by imputing missing values.

        Args:
            X: Input DataFrame or Series

        Returns:
            Transformed DataFrame or Series with imputed values
        """
        if not self.is_fitted:
            raise RuntimeError(
                "Processor is not fitted. Call 'fit' with appropriate arguments before using this method."
            )

        # Handle Series input
        if isinstance(X, pd.Series):
            if X.name not in self.imputation_dict:
                raise ValueError(f"No imputation value found for series name: {X.name}")
            return X.fillna(self.imputation_dict[X.name])

        # Handle DataFrame input
        if not isinstance(X, pd.DataFrame):
            raise TypeError("Input must be pandas Series or DataFrame")

        # Make a copy to avoid modifying the input
        df = X.copy()

        # Apply imputation only to variables in imputation_dict and only to NaN values
        for var, impute_value in self.imputation_dict.items():
            if var in df.columns:
                # Create mask for NaN values
                nan_mask = df[var].isna()
                # Only replace NaN values
                df.loc[nan_mask, var] = impute_value

        return df

    def get_params(self) -> Dict[str, Any]:
        return {
            "variables": self.variables,
            "imputation_dict": self.imputation_dict,
            "strategy": self.strategy,
        }


import logging

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

# Container path constants - aligned with script contract
CONTAINER_PATHS = {
    "MODEL_DIR": "/opt/ml/processing/input/model",
    "EVAL_DATA_DIR": "/opt/ml/processing/input/eval_data",
    "OUTPUT_EVAL_DIR": "/opt/ml/processing/output/eval",
    "OUTPUT_METRICS_DIR": "/opt/ml/processing/output/metrics",
}


# ============================================================================
# FILE I/O HELPER FUNCTIONS WITH FORMAT PRESERVATION
# ============================================================================


def _detect_file_format(file_path: Path) -> str:
    """
    Detect the format of a data file based on its extension.

    Args:
        file_path: Path to the file

    Returns:
        Format string: 'csv', 'tsv', or 'parquet'
    """
    suffix = file_path.suffix.lower()

    if suffix == ".csv":
        return "csv"
    elif suffix == ".tsv":
        return "tsv"
    elif suffix == ".parquet":
        return "parquet"
    else:
        raise RuntimeError(f"Unsupported file format: {suffix}")


def load_dataframe_with_format(file_path: Path) -> Tuple[pd.DataFrame, str]:
    """
    Load DataFrame and detect its format.

    Args:
        file_path: Path to the file

    Returns:
        Tuple of (DataFrame, format_string)
    """
    detected_format = _detect_file_format(file_path)

    if detected_format == "csv":
        df = pd.read_csv(file_path)
    elif detected_format == "tsv":
        df = pd.read_csv(file_path, sep="\t")
    elif detected_format == "parquet":
        df = pd.read_parquet(file_path)
    else:
        raise RuntimeError(f"Unsupported format: {detected_format}")

    return df, detected_format


def save_dataframe_with_format(
    df: pd.DataFrame, output_path: Path, format_str: str
) -> Path:
    """
    Save DataFrame in specified format.

    Args:
        df: DataFrame to save
        output_path: Base output path (without extension)
        format_str: Format to save in ('csv', 'tsv', or 'parquet')

    Returns:
        Path to saved file
    """
    if format_str == "csv":
        file_path = output_path.with_suffix(".csv")
        df.to_csv(file_path, index=False)
    elif format_str == "tsv":
        file_path = output_path.with_suffix(".tsv")
        df.to_csv(file_path, sep="\t", index=False)
    elif format_str == "parquet":
        file_path = output_path.with_suffix(".parquet")
        df.to_parquet(file_path, index=False)
    else:
        raise RuntimeError(f"Unsupported output format: {format_str}")

    return file_path


def decompress_model_artifacts(model_dir: str):
    """
    Checks for a model.tar.gz file in the model directory and extracts it.
    """
    model_tar_path = Path(model_dir) / "model.tar.gz"
    if model_tar_path.exists():
        logger.info(f"Found model.tar.gz at {model_tar_path}. Extracting...")
        with tarfile.open(model_tar_path, "r:gz") as tar:
            tar.extractall(path=model_dir)
        logger.info("Extraction complete.")
    else:
        logger.info("No model.tar.gz found. Assuming artifacts are directly available.")


def load_model_artifacts(
    model_dir: str,
) -> Tuple[lgb.Booster, Dict[str, Any], Dict[str, Any], List[str], Dict[str, Any]]:
    """
    Load the trained LightGBM model and all preprocessing artifacts from the specified directory.
    Returns model, risk_tables, impute_dict, feature_columns, and hyperparameters.
    """
    logger.info(f"Loading model artifacts from {model_dir}")

    # Decompress the model tarball if it exists
    logger.info("Checking for model.tar.gz and decompressing if present")
    decompress_model_artifacts(model_dir)

    model_file = os.path.join(model_dir, "lightgbm_model.txt")
    model = lgb.Booster(model_file=model_file)
    logger.info(f"Loaded lightgbm_model.txt")

    with open(os.path.join(model_dir, "risk_table_map.pkl"), "rb") as f:
        risk_tables = pkl.load(f)
    logger.info("Loaded risk_table_map.pkl")
    with open(os.path.join(model_dir, "impute_dict.pkl"), "rb") as f:
        impute_dict = pkl.load(f)
    logger.info("Loaded impute_dict.pkl")
    with open(os.path.join(model_dir, "feature_columns.txt"), "r") as f:
        feature_columns = [
            line.strip().split(",")[1] for line in f if not line.startswith("#")
        ]
    logger.info(f"Loaded feature_columns.txt: {feature_columns}")
    with open(os.path.join(model_dir, "hyperparameters.json"), "r") as f:
        hyperparams = json.load(f)
    logger.info("Loaded hyperparameters.json")

    return model, risk_tables, impute_dict, feature_columns, hyperparams


def preprocess_eval_data(
    df: pd.DataFrame,
    feature_columns: List[str],
    risk_tables: Dict[str, Any],
    impute_dict: Dict[str, Any],
) -> pd.DataFrame:
    """
    Apply risk table mapping and numerical imputation to the evaluation DataFrame.
    Ensures all features are numeric and columns are ordered as required by the model.
    Preserves any non-feature columns like id and label.
    """
    # Make a copy of the input dataframe to avoid modifying the original
    result_df = df.copy()

    # Get available feature columns (features that exist in the input data)
    available_features = [col for col in feature_columns if col in df.columns]
    logger.info(
        f"Found {len(available_features)} out of {len(feature_columns)} expected feature columns"
    )

    # Process only feature columns
    logger.info("Starting risk table mapping for categorical features")
    for feature, risk_table in risk_tables.items():
        if feature in available_features:
            logger.info(f"Applying risk table mapping for feature: {feature}")
            proc = RiskTableMappingProcessor(
                column_name=feature, label_name="label", risk_tables=risk_table
            )
            result_df[feature] = proc.transform(df[feature])
    logger.info("Risk table mapping complete")

    # For numerical imputation, only process features, not all columns
    logger.info("Starting numerical imputation")
    feature_df = result_df[available_features].copy()
    imputer = NumericalVariableImputationProcessor(imputation_dict=impute_dict)
    imputed_df = imputer.transform(feature_df)
    # Update only the feature columns in the result dataframe
    for col in available_features:
        if col in imputed_df:
            result_df[col] = imputed_df[col]
    logger.info("Numerical imputation complete")

    # Convert feature columns to numeric, leaving other columns unchanged
    logger.info("Ensuring feature columns are numeric")
    result_df[available_features] = (
        result_df[available_features].apply(pd.to_numeric, errors="coerce").fillna(0)
    )

    logger.info(
        f"Preprocessed data shape: {result_df.shape} (preserving all original columns)"
    )

    # Return the result with all original columns preserved
    return result_df


def log_metrics_summary(
    metrics: Dict[str, Union[int, float, str]], is_binary: bool = True
) -> None:
    """
    Log a nicely formatted summary of metrics for easy visibility in logs.

    Args:
        metrics: Dictionary of metrics to log
        is_binary: Whether these are binary classification metrics
    """
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    logger.info("=" * 80)
    logger.info(f"METRICS SUMMARY - {timestamp}")
    logger.info("=" * 80)

    # Log each metric with a consistent format
    for name, value in metrics.items():
        # Format numeric values to 4 decimal places
        if isinstance(value, (int, float)):
            formatted_value = f"{value:.4f}"
        else:
            formatted_value = str(value)

        # Add a special prefix for easy searching in logs
        logger.info(f"METRIC: {name.ljust(25)} = {formatted_value}")

    # Highlight key metrics based on task type
    logger.info("=" * 80)
    logger.info("KEY PERFORMANCE METRICS")
    logger.info("=" * 80)

    if is_binary:
        logger.info(
            f"METRIC_KEY: AUC-ROC               = {metrics.get('auc_roc', 'N/A'):.4f}"
        )
        logger.info(
            f"METRIC_KEY: Average Precision     = {metrics.get('average_precision', 'N/A'):.4f}"
        )
        logger.info(
            f"METRIC_KEY: F1 Score              = {metrics.get('f1_score', 'N/A'):.4f}"
        )
    else:
        logger.info(
            f"METRIC_KEY: Macro AUC-ROC         = {metrics.get('auc_roc_macro', 'N/A'):.4f}"
        )
        logger.info(
            f"METRIC_KEY: Micro AUC-ROC         = {metrics.get('auc_roc_micro', 'N/A'):.4f}"
        )
        ap_macro = metrics.get("average_precision_macro", "N/A")
        if isinstance(ap_macro, (int, float)):
            logger.info(f"METRIC_KEY: Macro Average Precision = {ap_macro:.4f}")
        else:
            logger.info(f"METRIC_KEY: Macro Average Precision = {ap_macro}")
        logger.info(
            f"METRIC_KEY: Macro F1              = {metrics.get('f1_score_macro', 'N/A'):.4f}"
        )
        logger.info(
            f"METRIC_KEY: Micro F1              = {metrics.get('f1_score_micro', 'N/A'):.4f}"
        )

    # Add a summary section with pass/fail criteria if defined
    logger.info("=" * 80)


def compute_metrics_binary(y_true: np.ndarray, y_prob: np.ndarray) -> Dict[str, float]:
    """
    Compute binary classification metrics: AUC-ROC, average precision, and F1 score.
    """
    logger.info("Computing binary classification metrics")
    y_score = y_prob[:, 1]
    metrics = {
        "auc_roc": roc_auc_score(y_true, y_score),
        "average_precision": average_precision_score(y_true, y_score),
        "f1_score": f1_score(y_true, y_score > 0.5),
    }

    # Add more detailed metrics
    precision, recall, _ = precision_recall_curve(y_true, y_score)
    metrics["precision_at_threshold_0.5"] = precision[0]
    metrics["recall_at_threshold_0.5"] = recall[0]

    # Thresholds at different operating points
    for threshold in [0.3, 0.5, 0.7]:
        y_pred = (y_score >= threshold).astype(int)
        metrics[f"f1_score_at_{threshold}"] = f1_score(y_true, y_pred)

    # Log basic summary and detailed formatted metrics
    logger.info(
        f"Binary metrics computed: AUC={metrics['auc_roc']:.4f}, AP={metrics['average_precision']:.4f}, F1={metrics['f1_score']:.4f}"
    )
    log_metrics_summary(metrics, is_binary=True)

    return metrics


def compute_metrics_multiclass(
    y_true: np.ndarray, y_prob: np.ndarray, n_classes: int
) -> Dict[str, Union[int, float]]:
    """
    Compute multiclass metrics: one-vs-rest AUC-ROC, average precision, F1 for each class,
    and micro/macro averages for all metrics.
    """
    logger.info("Computing multiclass metrics")
    metrics = {}

    # Per-class metrics
    for i in range(n_classes):
        y_true_bin = (y_true == i).astype(int)
        y_score = y_prob[:, i]
        metrics[f"auc_roc_class_{i}"] = roc_auc_score(y_true_bin, y_score)
        metrics[f"average_precision_class_{i}"] = average_precision_score(
            y_true_bin, y_score
        )
        metrics[f"f1_score_class_{i}"] = f1_score(y_true_bin, y_score > 0.5)

    # Micro and macro averages
    metrics["auc_roc_micro"] = roc_auc_score(
        y_true, y_prob, multi_class="ovr", average="micro"
    )
    metrics["auc_roc_macro"] = roc_auc_score(
        y_true, y_prob, multi_class="ovr", average="macro"
    )
    metrics["average_precision_micro"] = average_precision_score(
        y_true, y_prob, average="micro"
    )
    metrics["average_precision_macro"] = average_precision_score(
        y_true, y_prob, average="macro"
    )

    y_pred = np.argmax(y_prob, axis=1)
    metrics["f1_score_micro"] = f1_score(y_true, y_pred, average="micro")
    metrics["f1_score_macro"] = f1_score(y_true, y_pred, average="macro")

    # Class distribution metrics
    unique, counts = np.unique(y_true, return_counts=True)
    for cls, count in zip(unique, counts):
        metrics[f"class_{cls}_count"] = int(count)
        metrics[f"class_{cls}_ratio"] = float(count) / len(y_true)

    # Log basic summary and detailed formatted metrics
    logger.info(
        f"Multiclass metrics computed: Macro AUC={metrics['auc_roc_macro']:.4f}, Micro AUC={metrics['auc_roc_micro']:.4f}"
    )
    log_metrics_summary(metrics, is_binary=False)

    return metrics


def compute_comparison_metrics(
    y_true: np.ndarray,
    y_new_score: np.ndarray,
    y_prev_score: np.ndarray,
    is_binary: bool = True,
) -> Dict[str, float]:
    """
    Compute comparison metrics between new model and previous model scores.

    Args:
        y_true: True labels
        y_new_score: New model prediction scores
        y_prev_score: Previous model prediction scores
        is_binary: Whether this is binary classification

    Returns:
        Dictionary of comparison metrics
    """
    logger.info("Computing model comparison metrics")

    comparison_metrics = {}

    # Basic correlation metrics - with error handling for scipy compatibility
    try:
        pearson_corr, pearson_p = pearsonr(y_new_score, y_prev_score)
        spearman_corr, spearman_p = spearmanr(y_new_score, y_prev_score)
    except (TypeError, AttributeError) as e:
        logger.warning(
            f"SciPy correlation computation failed: {e}. Using fallback numpy correlation."
        )
        # Fallback to numpy correlation
        pearson_corr = float(np.corrcoef(y_new_score, y_prev_score)[0, 1])
        pearson_p = 0.0  # p-value not available with numpy
        spearman_corr = pearson_corr  # Use Pearson as fallback
        spearman_p = 0.0

    comparison_metrics.update(
        {
            "pearson_correlation": pearson_corr,
            "pearson_p_value": pearson_p,
            "spearman_correlation": spearman_corr,
            "spearman_p_value": spearman_p,
        }
    )

    # Performance comparison metrics
    if is_binary:
        # Binary classification comparison
        new_auc = roc_auc_score(y_true, y_new_score)
        prev_auc = roc_auc_score(y_true, y_prev_score)
        new_ap = average_precision_score(y_true, y_new_score)
        prev_ap = average_precision_score(y_true, y_prev_score)

        # Delta metrics
        comparison_metrics.update(
            {
                "new_model_auc": new_auc,
                "previous_model_auc": prev_auc,
                "auc_delta": new_auc - prev_auc,
                "auc_lift_percent": ((new_auc - prev_auc) / prev_auc) * 100
                if prev_auc > 0
                else 0,
                "new_model_ap": new_ap,
                "previous_model_ap": prev_ap,
                "ap_delta": new_ap - prev_ap,
                "ap_lift_percent": ((new_ap - prev_ap) / prev_ap) * 100
                if prev_ap > 0
                else 0,
            }
        )

        # F1 score comparison at different thresholds
        for threshold in [0.3, 0.5, 0.7]:
            new_f1 = f1_score(y_true, (y_new_score >= threshold).astype(int))
            prev_f1 = f1_score(y_true, (y_prev_score >= threshold).astype(int))
            comparison_metrics[f"new_model_f1_at_{threshold}"] = new_f1
            comparison_metrics[f"previous_model_f1_at_{threshold}"] = prev_f1
            comparison_metrics[f"f1_delta_at_{threshold}"] = new_f1 - prev_f1

    # Score distribution comparison
    comparison_metrics.update(
        {
            "new_score_mean": float(np.mean(y_new_score)),
            "previous_score_mean": float(np.mean(y_prev_score)),
            "new_score_std": float(np.std(y_new_score)),
            "previous_score_std": float(np.std(y_prev_score)),
            "score_mean_delta": float(np.mean(y_new_score) - np.mean(y_prev_score)),
        }
    )

    # Agreement metrics
    # For binary classification, compute agreement at different thresholds
    if is_binary:
        for threshold in [0.3, 0.5, 0.7]:
            new_pred = (y_new_score >= threshold).astype(int)
            prev_pred = (y_prev_score >= threshold).astype(int)
            agreement = np.mean(new_pred == prev_pred)
            comparison_metrics[f"prediction_agreement_at_{threshold}"] = agreement

    logger.info(
        f"Comparison metrics computed: AUC delta={comparison_metrics.get('auc_delta', 'N/A'):.4f}, "
        f"Correlation={comparison_metrics.get('pearson_correlation', 'N/A'):.4f}"
    )

    return comparison_metrics


def perform_statistical_tests(
    y_true: np.ndarray,
    y_new_score: np.ndarray,
    y_prev_score: np.ndarray,
    is_binary: bool = True,
) -> Dict[str, float]:
    """
    Perform statistical significance tests comparing model performances.

    Args:
        y_true: True labels
        y_new_score: New model prediction scores
        y_prev_score: Previous model prediction scores
        is_binary: Whether this is binary classification

    Returns:
        Dictionary of statistical test results
    """
    logger.info("Performing statistical significance tests")

    test_results = {}

    if is_binary:
        # McNemar's test for binary classification
        # Compare predictions at 0.5 threshold
        new_pred = (y_new_score >= 0.5).astype(int)
        prev_pred = (y_prev_score >= 0.5).astype(int)

        # Create contingency table for McNemar's test
        # [correct_both, new_correct_prev_wrong]
        # [new_wrong_prev_correct, wrong_both]
        correct_both = np.sum((new_pred == y_true) & (prev_pred == y_true))
        new_correct_prev_wrong = np.sum((new_pred == y_true) & (prev_pred != y_true))
        new_wrong_prev_correct = np.sum((new_pred != y_true) & (prev_pred == y_true))
        wrong_both = np.sum((new_pred != y_true) & (prev_pred != y_true))

        # McNemar's test statistic
        if (new_correct_prev_wrong + new_wrong_prev_correct) > 0:
            mcnemar_stat = (
                (abs(new_correct_prev_wrong - new_wrong_prev_correct) - 1) ** 2
            ) / (new_correct_prev_wrong + new_wrong_prev_correct)
            mcnemar_p_value = 1 - stats.chi2.cdf(mcnemar_stat, 1)
        else:
            mcnemar_stat = 0.0
            mcnemar_p_value = 1.0

        test_results.update(
            {
                "mcnemar_statistic": mcnemar_stat,
                "mcnemar_p_value": mcnemar_p_value,
                "mcnemar_significant": bool(mcnemar_p_value < 0.05),
                "correct_both": int(correct_both),
                "new_correct_prev_wrong": int(new_correct_prev_wrong),
                "new_wrong_prev_correct": int(new_wrong_prev_correct),
                "wrong_both": int(wrong_both),
            }
        )

    # Paired t-test on prediction scores
    t_stat, t_p_value = stats.ttest_rel(y_new_score, y_prev_score)
    test_results.update(
        {
            "paired_t_statistic": t_stat,
            "paired_t_p_value": t_p_value,
            "paired_t_significant": bool(t_p_value < 0.05),
        }
    )

    # Wilcoxon signed-rank test (non-parametric alternative)
    try:
        wilcoxon_stat, wilcoxon_p = stats.wilcoxon(y_new_score, y_prev_score)
        test_results.update(
            {
                "wilcoxon_statistic": wilcoxon_stat,
                "wilcoxon_p_value": wilcoxon_p,
                "wilcoxon_significant": bool(wilcoxon_p < 0.05),
            }
        )
    except ValueError as e:
        logger.warning(f"Could not perform Wilcoxon test: {e}")
        test_results.update(
            {
                "wilcoxon_statistic": np.nan,
                "wilcoxon_p_value": np.nan,
                "wilcoxon_significant": False,
            }
        )

    logger.info(
        f"Statistical tests completed: McNemar p={test_results.get('mcnemar_p_value', 'N/A'):.4f}, "
        f"Paired t-test p={test_results.get('paired_t_p_value', 'N/A'):.4f}"
    )

    return test_results


def load_eval_data(eval_data_dir: str) -> Tuple[pd.DataFrame, str]:
    """
    Load the first data file found in the evaluation data directory.
    Returns a pandas DataFrame and the detected format.
    """
    logger.info(f"Loading eval data from {eval_data_dir}")
    eval_files = sorted(
        [
            f
            for f in Path(eval_data_dir).glob("**/*")
            if f.suffix in [".csv", ".tsv", ".parquet"]
        ]
    )
    if not eval_files:
        logger.error("No eval data file found in eval_data input.")
        raise RuntimeError("No eval data file found in eval_data input.")
    eval_file = eval_files[0]
    logger.info(f"Using eval data file: {eval_file}")

    df, input_format = load_dataframe_with_format(eval_file)
    logger.info(f"Loaded eval data shape: {df.shape}, format: {input_format}")
    return df, input_format


def get_id_label_columns(
    df: pd.DataFrame, id_field: str, label_field: str
) -> Tuple[str, str]:
    """
    Determine the ID and label columns in the DataFrame.
    Falls back to the first and second columns if not found.
    """
    id_col = id_field if id_field in df.columns else df.columns[0]
    label_col = label_field if label_field in df.columns else df.columns[1]
    logger.info(f"Using id_col: {id_col}, label_col: {label_col}")
    return id_col, label_col


def save_predictions(
    ids: np.ndarray,
    y_true: np.ndarray,
    y_prob: np.ndarray,
    id_col: str,
    label_col: str,
    output_eval_dir: str,
    input_format: str = "csv",
) -> None:
    """
    Save predictions preserving input format, including id, true label, and class probabilities.
    """
    logger.info(f"Saving predictions to {output_eval_dir} in {input_format} format")
    prob_cols = [f"prob_class_{i}" for i in range(y_prob.shape[1])]
    out_df = pd.DataFrame({id_col: ids, label_col: y_true})
    for i, col in enumerate(prob_cols):
        out_df[col] = y_prob[:, i]

    output_base = Path(output_eval_dir) / "eval_predictions"
    output_path = save_dataframe_with_format(out_df, output_base, input_format)
    logger.info(f"Saved predictions (format={input_format}): {output_path}")


def save_metrics(
    metrics: Dict[str, Union[int, float, str]], output_metrics_dir: str
) -> None:
    """
    Save computed metrics as a JSON file.
    """
    out_path = os.path.join(output_metrics_dir, "metrics.json")
    with open(out_path, "w") as f:
        json.dump(metrics, f, indent=2)
    logger.info(f"Saved metrics to {out_path}")

    # Also create a plain text summary for easy viewing
    summary_path = os.path.join(output_metrics_dir, "metrics_summary.txt")
    with open(summary_path, "w") as f:
        f.write("METRICS SUMMARY\n")
        f.write("=" * 50 + "\n")

        # Write key metrics at the top
        if "auc_roc" in metrics:  # Binary classification
            f.write(f"AUC-ROC:           {metrics['auc_roc']:.4f}\n")
            if "average_precision" in metrics:
                f.write(f"Average Precision: {metrics['average_precision']:.4f}\n")
            if "f1_score" in metrics:
                f.write(f"F1 Score:          {metrics['f1_score']:.4f}\n")
        else:  # Multiclass classification
            f.write(f"AUC-ROC (Macro):   {metrics.get('auc_roc_macro', 'N/A'):.4f}\n")
            f.write(f"AUC-ROC (Micro):   {metrics.get('auc_roc_micro', 'N/A'):.4f}\n")
            f.write(f"F1 Score (Macro):  {metrics.get('f1_score_macro', 'N/A'):.4f}\n")

        f.write("=" * 50 + "\n\n")

        # Write all metrics
        f.write("ALL METRICS\n")
        f.write("=" * 50 + "\n")
        for name, value in sorted(metrics.items()):
            if isinstance(value, (int, float)):
                f.write(f"{name}: {value:.6f}\n")
            else:
                f.write(f"{name}: {value}\n")

    logger.info(f"Saved metrics summary to {summary_path}")


def plot_and_save_roc_curve(
    y_true: np.ndarray, y_score: np.ndarray, output_dir: str, prefix: str = ""
) -> None:
    """
    Plot ROC curve and save as JPG.
    """
    fpr, tpr, _ = roc_curve(y_true, y_score)
    auc = roc_auc_score(y_true, y_score)
    plt.figure()
    plt.plot(fpr, tpr, label=f"ROC curve (AUC = {auc:.2f})")
    plt.plot([0, 1], [0, 1], "k--", label="Random")
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title("ROC Curve")
    plt.legend(loc="lower right")
    out_path = os.path.join(output_dir, f"{prefix}roc_curve.jpg")
    plt.savefig(out_path, format="jpg")
    plt.close()
    logger.info(f"Saved ROC curve to {out_path}")


def plot_and_save_pr_curve(
    y_true: np.ndarray, y_score: np.ndarray, output_dir: str, prefix: str = ""
) -> None:
    """
    Plot Precision-Recall curve and save as JPG.
    """
    precision, recall, _ = precision_recall_curve(y_true, y_score)
    ap = average_precision_score(y_true, y_score)
    plt.figure()
    plt.plot(recall, precision, label=f"PR curve (AP = {ap:.2f})")
    plt.xlabel("Recall")
    plt.ylabel("Precision")
    plt.title("Precision-Recall Curve")
    plt.legend(loc="lower left")
    out_path = os.path.join(output_dir, f"{prefix}pr_curve.jpg")
    plt.savefig(out_path, format="jpg")
    plt.close()
    logger.info(f"Saved PR curve to {out_path}")


def plot_comparison_roc_curves(
    y_true: np.ndarray,
    y_new_score: np.ndarray,
    y_prev_score: np.ndarray,
    output_dir: str,
) -> None:
    """
    Plot side-by-side ROC curves comparing new and previous models.
    """
    logger.info("Creating comparison ROC curves")

    # Calculate ROC curves for both models
    fpr_new, tpr_new, _ = roc_curve(y_true, y_new_score)
    fpr_prev, tpr_prev, _ = roc_curve(y_true, y_prev_score)

    auc_new = roc_auc_score(y_true, y_new_score)
    auc_prev = roc_auc_score(y_true, y_prev_score)

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

    # Plot both ROC curves
    plt.plot(
        fpr_new, tpr_new, "b-", linewidth=2, label=f"New Model (AUC = {auc_new:.3f})"
    )
    plt.plot(
        fpr_prev,
        tpr_prev,
        "r--",
        linewidth=2,
        label=f"Previous Model (AUC = {auc_prev:.3f})",
    )
    plt.plot([0, 1], [0, 1], "k:", alpha=0.6, label="Random")

    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title(f"ROC Curve Comparison (Δ AUC = {auc_new - auc_prev:+.3f})")
    plt.legend(loc="lower right")
    plt.grid(True, alpha=0.3)

    out_path = os.path.join(output_dir, "comparison_roc_curves.jpg")
    plt.savefig(out_path, format="jpg", dpi=150, bbox_inches="tight")
    plt.close()
    logger.info(f"Saved comparison ROC curves to {out_path}")


def plot_comparison_pr_curves(
    y_true: np.ndarray,
    y_new_score: np.ndarray,
    y_prev_score: np.ndarray,
    output_dir: str,
) -> None:
    """
    Plot side-by-side Precision-Recall curves comparing new and previous models.
    """
    logger.info("Creating comparison PR curves")

    # Calculate PR curves for both models
    precision_new, recall_new, _ = precision_recall_curve(y_true, y_new_score)
    precision_prev, recall_prev, _ = precision_recall_curve(y_true, y_prev_score)

    ap_new = average_precision_score(y_true, y_new_score)
    ap_prev = average_precision_score(y_true, y_prev_score)

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

    # Plot both PR curves
    plt.plot(
        recall_new,
        precision_new,
        "b-",
        linewidth=2,
        label=f"New Model (AP = {ap_new:.3f})",
    )
    plt.plot(
        recall_prev,
        precision_prev,
        "r--",
        linewidth=2,
        label=f"Previous Model (AP = {ap_prev:.3f})",
    )

    # Add baseline (random classifier)
    baseline = np.mean(y_true)
    plt.axhline(
        y=baseline,
        color="k",
        linestyle=":",
        alpha=0.6,
        label=f"Random (AP = {baseline:.3f})",
    )

    plt.xlabel("Recall")
    plt.ylabel("Precision")
    plt.title(f"Precision-Recall Curve Comparison (Δ AP = {ap_new - ap_prev:+.3f})")
    plt.legend(loc="lower left")
    plt.grid(True, alpha=0.3)

    out_path = os.path.join(output_dir, "comparison_pr_curves.jpg")
    plt.savefig(out_path, format="jpg", dpi=150, bbox_inches="tight")
    plt.close()
    logger.info(f"Saved comparison PR curves to {out_path}")


def plot_score_scatter(
    y_new_score: np.ndarray,
    y_prev_score: np.ndarray,
    y_true: np.ndarray,
    output_dir: str,
) -> None:
    """
    Plot scatter plot of new vs previous model scores, colored by true labels.
    """
    logger.info("Creating score scatter plot")

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

    # Separate positive and negative examples
    pos_mask = y_true == 1
    neg_mask = y_true == 0

    # Plot negative examples
    plt.scatter(
        y_prev_score[neg_mask],
        y_new_score[neg_mask],
        c="lightcoral",
        alpha=0.6,
        s=20,
        label="Negative (0)",
    )

    # Plot positive examples
    plt.scatter(
        y_prev_score[pos_mask],
        y_new_score[pos_mask],
        c="lightblue",
        alpha=0.6,
        s=20,
        label="Positive (1)",
    )

    # Add diagonal line (perfect correlation)
    min_score = min(np.min(y_prev_score), np.min(y_new_score))
    max_score = max(np.max(y_prev_score), np.max(y_new_score))
    plt.plot(
        [min_score, max_score],
        [min_score, max_score],
        "k--",
        alpha=0.8,
        label="Perfect Correlation",
    )

    # Calculate and display correlation with error handling for SciPy compatibility
    try:
        correlation = pearsonr(y_new_score, y_prev_score)[0]
    except (TypeError, AttributeError) as e:
        logger.warning(f"SciPy pearsonr failed: {e}. Using numpy correlation.")
        correlation = float(np.corrcoef(y_new_score, y_prev_score)[0, 1])

    plt.xlabel("Previous Model Score")
    plt.ylabel("New Model Score")
    plt.title(f"Model Score Comparison (Correlation = {correlation:.3f})")
    plt.legend()
    plt.grid(True, alpha=0.3)

    # Add correlation text box
    textstr = f"Pearson r = {correlation:.3f}"
    props = dict(boxstyle="round", facecolor="wheat", alpha=0.8)
    plt.text(
        0.05,
        0.95,
        textstr,
        transform=plt.gca().transAxes,
        fontsize=12,
        verticalalignment="top",
        bbox=props,
    )

    out_path = os.path.join(output_dir, "score_scatter_plot.jpg")
    plt.savefig(out_path, format="jpg", dpi=150, bbox_inches="tight")
    plt.close()
    logger.info(f"Saved score scatter plot to {out_path}")


def plot_score_distributions(
    y_new_score: np.ndarray,
    y_prev_score: np.ndarray,
    y_true: np.ndarray,
    output_dir: str,
) -> None:
    """
    Plot score distributions for both models, separated by true labels.
    """
    logger.info("Creating score distribution plots")

    # Set matplotlib backend for headless environments
    import matplotlib

    matplotlib.use("Agg")

    # Create figure and axes with comprehensive error handling
    fig = None
    axes = None

    try:
        # First attempt: standard subplots
        result = plt.subplots(2, 2, figsize=(15, 10))
        if isinstance(result, tuple) and len(result) == 2:
            fig, axes = result
            # Ensure axes is always a 2D array
            if hasattr(axes, "ndim") and axes.ndim == 1:
                axes = axes.reshape(2, 2)
        else:
            raise ValueError("subplots returned unexpected format")
    except Exception as e:
        logger.warning(f"Standard subplots failed: {e}. Using fallback approach.")

        # Fallback approach: create figure and individual subplots
        try:
            fig = plt.figure(figsize=(15, 10))
            axes = []
            for i in range(4):
                ax = fig.add_subplot(2, 2, i + 1)
                axes.append(ax)
            axes = np.array(axes).reshape(2, 2)
        except Exception as e2:
            logger.error(f"Fallback subplot creation also failed: {e2}")
            raise RuntimeError(
                f"Failed to create matplotlib plots after multiple attempts: {e2}"
            ) from e2

    # Separate positive and negative examples
    pos_mask = y_true == 1
    neg_mask = y_true == 0

    # Plot 1: New model score distributions
    axes[0, 0].hist(
        y_new_score[neg_mask],
        bins=30,
        alpha=0.7,
        color="lightcoral",
        label="Negative (0)",
        density=True,
    )
    axes[0, 0].hist(
        y_new_score[pos_mask],
        bins=30,
        alpha=0.7,
        color="lightblue",
        label="Positive (1)",
        density=True,
    )
    axes[0, 0].set_title("New Model Score Distribution")
    axes[0, 0].set_xlabel("Score")
    axes[0, 0].set_ylabel("Density")
    axes[0, 0].legend()
    axes[0, 0].grid(True, alpha=0.3)

    # Plot 2: Previous model score distributions
    axes[0, 1].hist(
        y_prev_score[neg_mask],
        bins=30,
        alpha=0.7,
        color="lightcoral",
        label="Negative (0)",
        density=True,
    )
    axes[0, 1].hist(
        y_prev_score[pos_mask],
        bins=30,
        alpha=0.7,
        color="lightblue",
        label="Positive (1)",
        density=True,
    )
    axes[0, 1].set_title("Previous Model Score Distribution")
    axes[0, 1].set_xlabel("Score")
    axes[0, 1].set_ylabel("Density")
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)

    # Plot 3: Score difference distribution
    score_diff = y_new_score - y_prev_score
    axes[1, 0].hist(
        score_diff[neg_mask],
        bins=30,
        alpha=0.7,
        color="lightcoral",
        label="Negative (0)",
        density=True,
    )
    axes[1, 0].hist(
        score_diff[pos_mask],
        bins=30,
        alpha=0.7,
        color="lightblue",
        label="Positive (1)",
        density=True,
    )
    axes[1, 0].axvline(x=0, color="black", linestyle="--", alpha=0.8)
    axes[1, 0].set_title("Score Difference Distribution (New - Previous)")
    axes[1, 0].set_xlabel("Score Difference")
    axes[1, 0].set_ylabel("Density")
    axes[1, 0].legend()
    axes[1, 0].grid(True, alpha=0.3)

    # Plot 4: Box plots comparing both models
    box_data = [
        y_prev_score[neg_mask],
        y_new_score[neg_mask],
        y_prev_score[pos_mask],
        y_new_score[pos_mask],
    ]
    box_labels = ["Prev (Neg)", "New (Neg)", "Prev (Pos)", "New (Pos)"]
    box_colors = ["lightcoral", "lightcoral", "lightblue", "lightblue"]

    bp = axes[1, 1].boxplot(box_data, labels=box_labels, patch_artist=True)
    for patch, color in zip(bp["boxes"], box_colors):
        patch.set_facecolor(color)
        patch.set_alpha(0.7)

    axes[1, 1].set_title("Score Distribution Comparison")
    axes[1, 1].set_ylabel("Score")
    axes[1, 1].grid(True, alpha=0.3)

    plt.tight_layout()
    out_path = os.path.join(output_dir, "score_distributions.jpg")
    plt.savefig(out_path, format="jpg", dpi=150, bbox_inches="tight")
    plt.close()
    logger.info(f"Saved score distribution plots to {out_path}")


def evaluate_model(
    model: lgb.Booster,
    df: pd.DataFrame,
    feature_columns: List[str],
    id_col: str,
    label_col: str,
    hyperparams: Dict[str, Any],
    output_eval_dir: str,
    output_metrics_dir: str,
    input_format: str = "csv",
) -> None:
    """
    Run model prediction and evaluation, then save predictions and metrics preserving format.
    Also generate and save ROC and PR curves as JPG.
    """
    logger.info("Evaluating model")
    y_true = df[label_col].values
    ids = df[id_col].values
    X = df[feature_columns].values

    y_prob = model.predict(X)
    logger.info(f"Model prediction shape: {y_prob.shape}")
    if len(y_prob.shape) == 1:
        y_prob = np.column_stack([1 - y_prob, y_prob])
        logger.info("Converted binary prediction to two-column probabilities")

    # Determine the classification type from the model's saved hyperparameters,
    # which is the definitive source of truth.
    is_binary_model = hyperparams.get("is_binary", True)

    if is_binary_model:
        logger.info(
            "Detected binary classification task based on model hyperparameters."
        )
        # Ensure y_true is also binary (0 or 1) for consistent metric calculation
        y_true = (y_true > 0).astype(int)
        metrics = compute_metrics_binary(y_true, y_prob)
        plot_and_save_roc_curve(y_true, y_prob[:, 1], output_metrics_dir)
        plot_and_save_pr_curve(y_true, y_prob[:, 1], output_metrics_dir)
    else:
        n_classes = y_prob.shape[1]
        logger.info(
            f"Detected multiclass classification task with {n_classes} classes."
        )
        metrics = compute_metrics_multiclass(y_true, y_prob, n_classes)
        for i in range(n_classes):
            y_true_bin = (y_true == i).astype(int)
            if len(np.unique(y_true_bin)) > 1:
                plot_and_save_roc_curve(
                    y_true_bin, y_prob[:, i], output_metrics_dir, prefix=f"class_{i}_"
                )
                plot_and_save_pr_curve(
                    y_true_bin, y_prob[:, i], output_metrics_dir, prefix=f"class_{i}_"
                )

    save_predictions(
        ids, y_true, y_prob, id_col, label_col, output_eval_dir, input_format
    )
    save_metrics(metrics, output_metrics_dir)
    logger.info("Evaluation complete")


def evaluate_model_with_comparison(
    model: lgb.Booster,
    df: pd.DataFrame,
    feature_columns: List[str],
    id_col: str,
    label_col: str,
    previous_scores: np.ndarray,
    hyperparams: Dict[str, Any],
    output_eval_dir: str,
    output_metrics_dir: str,
    comparison_metrics: str,
    statistical_tests: bool,
    comparison_plots: bool,
    input_format: str = "csv",
) -> None:
    """
    Run model prediction and evaluation with comparison to previous model scores preserving format.
    Generates comprehensive comparison metrics, statistical tests, and visualizations.
    """
    logger.info("Evaluating model with comparison mode enabled")

    # Get basic data
    y_true = df[label_col].values
    ids = df[id_col].values
    X = df[feature_columns].values

    # Generate new model predictions
    y_prob = model.predict(X)
    logger.info(f"Model prediction shape: {y_prob.shape}")

    if len(y_prob.shape) == 1:
        y_prob = np.column_stack([1 - y_prob, y_prob])
        logger.info("Converted binary prediction to two-column probabilities")

    # Determine the classification type
    is_binary_model = hyperparams.get("is_binary", True)

    if is_binary_model:
        logger.info(
            "Detected binary classification task based on model hyperparameters."
        )
        # Ensure y_true is also binary (0 or 1) for consistent metric calculation
        y_true = (y_true > 0).astype(int)

        # Get new model scores (probability of positive class)
        y_new_score = y_prob[:, 1]

        # Compute standard metrics for new model
        new_metrics = compute_metrics_binary(y_true, y_prob)

        # Compute comparison metrics
        if comparison_metrics in ["all", "basic"]:
            comp_metrics = compute_comparison_metrics(
                y_true, y_new_score, previous_scores, is_binary=True
            )
            new_metrics.update(comp_metrics)

        # Perform statistical tests
        if statistical_tests:
            stat_results = perform_statistical_tests(
                y_true, y_new_score, previous_scores, is_binary=True
            )
            new_metrics.update(stat_results)

        # Generate comparison plots
        if comparison_plots:
            plot_comparison_roc_curves(
                y_true, y_new_score, previous_scores, output_metrics_dir
            )
            plot_comparison_pr_curves(
                y_true, y_new_score, previous_scores, output_metrics_dir
            )
            plot_score_scatter(y_new_score, previous_scores, y_true, output_metrics_dir)
            plot_score_distributions(
                y_new_score, previous_scores, y_true, output_metrics_dir
            )

        # Generate standard plots for new model
        plot_and_save_roc_curve(
            y_true, y_new_score, output_metrics_dir, prefix="new_model_"
        )
        plot_and_save_pr_curve(
            y_true, y_new_score, output_metrics_dir, prefix="new_model_"
        )

        # Generate plots for previous model
        plot_and_save_roc_curve(
            y_true, previous_scores, output_metrics_dir, prefix="previous_model_"
        )
        plot_and_save_pr_curve(
            y_true, previous_scores, output_metrics_dir, prefix="previous_model_"
        )

    else:
        # Multiclass comparison (simplified for now)
        n_classes = y_prob.shape[1]
        logger.info(
            f"Detected multiclass classification task with {n_classes} classes."
        )
        logger.warning(
            "Multiclass comparison mode has limited functionality compared to binary classification."
        )

        # Compute standard multiclass metrics
        new_metrics = compute_metrics_multiclass(y_true, y_prob, n_classes)

        # For multiclass, we can still do some basic comparison if previous_scores represent class probabilities
        # This is a simplified implementation
        if len(previous_scores.shape) == 1:
            # If previous_scores is 1D, assume it's the score for the positive class in a binary-like scenario
            logger.warning(
                "Previous scores appear to be 1D for multiclass problem. Limited comparison available."
            )

        # Generate standard plots for multiclass
        for i in range(n_classes):
            y_true_bin = (y_true == i).astype(int)
            if len(np.unique(y_true_bin)) > 1:
                plot_and_save_roc_curve(
                    y_true_bin,
                    y_prob[:, i],
                    output_metrics_dir,
                    prefix=f"new_model_class_{i}_",
                )
                plot_and_save_pr_curve(
                    y_true_bin,
                    y_prob[:, i],
                    output_metrics_dir,
                    prefix=f"new_model_class_{i}_",
                )

    # Save enhanced predictions with previous scores
    save_predictions_with_comparison(
        ids,
        y_true,
        y_prob,
        previous_scores,
        id_col,
        label_col,
        output_eval_dir,
        input_format,
    )

    # Save comprehensive metrics
    save_metrics(new_metrics, output_metrics_dir)

    # Create comparison summary report
    create_comparison_report(new_metrics, output_metrics_dir, is_binary_model)

    logger.info("Enhanced evaluation with comparison complete")


def save_predictions_with_comparison(
    ids: np.ndarray,
    y_true: np.ndarray,
    y_prob: np.ndarray,
    previous_scores: np.ndarray,
    id_col: str,
    label_col: str,
    output_eval_dir: str,
    input_format: str = "csv",
) -> None:
    """
    Save predictions preserving input format, including id, true label, new model probabilities, and previous model scores.
    """
    logger.info(
        f"Saving predictions with comparison to {output_eval_dir} in {input_format} format"
    )

    # Create base dataframe
    prob_cols = [f"new_model_prob_class_{i}" for i in range(y_prob.shape[1])]
    out_df = pd.DataFrame({id_col: ids, label_col: y_true})

    # Add new model probabilities
    for i, col in enumerate(prob_cols):
        out_df[col] = y_prob[:, i]

    # Add previous model scores
    out_df["previous_model_score"] = previous_scores

    # Add score difference (for binary classification)
    if y_prob.shape[1] == 2:
        out_df["score_difference"] = y_prob[:, 1] - previous_scores

    output_base = Path(output_eval_dir) / "eval_predictions_with_comparison"
    output_path = save_dataframe_with_format(out_df, output_base, input_format)
    logger.info(
        f"Saved predictions with comparison (format={input_format}): {output_path}"
    )


def create_comparison_report(
    metrics: Dict[str, Union[int, float, str]], output_metrics_dir: str, is_binary: bool
) -> None:
    """
    Create a comprehensive comparison report summarizing model performance differences.
    """
    logger.info("Creating comparison report")

    report_path = os.path.join(output_metrics_dir, "comparison_report.txt")

    with open(report_path, "w") as f:
        f.write("MODEL COMPARISON REPORT\n")
        f.write("=" * 60 + "\n")
        f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")

        # Performance Summary
        f.write("PERFORMANCE SUMMARY\n")
        f.write("-" * 30 + "\n")

        if is_binary:
            new_auc = metrics.get("new_model_auc", "N/A")
            prev_auc = metrics.get("previous_model_auc", "N/A")
            auc_delta = metrics.get("auc_delta", "N/A")
            auc_lift = metrics.get("auc_lift_percent", "N/A")

            f.write(f"New Model AUC-ROC:      {new_auc:.4f}\n")
            f.write(f"Previous Model AUC-ROC: {prev_auc:.4f}\n")
            f.write(f"AUC Delta:              {auc_delta:+.4f}\n")
            f.write(f"AUC Lift:               {auc_lift:+.2f}%\n\n")

            new_ap = metrics.get("new_model_ap", "N/A")
            prev_ap = metrics.get("previous_model_ap", "N/A")
            ap_delta = metrics.get("ap_delta", "N/A")
            ap_lift = metrics.get("ap_lift_percent", "N/A")

            f.write(f"New Model Avg Precision: {new_ap:.4f}\n")
            f.write(f"Previous Model Avg Precision: {prev_ap:.4f}\n")
            f.write(f"AP Delta:               {ap_delta:+.4f}\n")
            f.write(f"AP Lift:                {ap_lift:+.2f}%\n\n")

        # Correlation Analysis
        f.write("CORRELATION ANALYSIS\n")
        f.write("-" * 30 + "\n")
        pearson_corr = metrics.get("pearson_correlation", "N/A")
        spearman_corr = metrics.get("spearman_correlation", "N/A")
        f.write(f"Pearson Correlation:    {pearson_corr:.4f}\n")
        f.write(f"Spearman Correlation:   {spearman_corr:.4f}\n\n")

        # Statistical Tests
        f.write("STATISTICAL SIGNIFICANCE\n")
        f.write("-" * 30 + "\n")

        if is_binary:
            mcnemar_p = metrics.get("mcnemar_p_value", "N/A")
            mcnemar_sig = metrics.get("mcnemar_significant", False)
            f.write(f"McNemar's Test p-value: {mcnemar_p:.4f}\n")
            f.write(
                f"McNemar's Test Result:  {'Significant' if mcnemar_sig else 'Not Significant'}\n\n"
            )

        paired_t_p = metrics.get("paired_t_p_value", "N/A")
        paired_t_sig = metrics.get("paired_t_significant", False)
        f.write(f"Paired t-test p-value:  {paired_t_p:.4f}\n")
        f.write(
            f"Paired t-test Result:   {'Significant' if paired_t_sig else 'Not Significant'}\n\n"
        )

        # Recommendations
        f.write("RECOMMENDATIONS\n")
        f.write("-" * 30 + "\n")

        if is_binary and isinstance(auc_delta, (int, float)):
            if auc_delta > 0.01:
                f.write(
                    "✓ NEW MODEL RECOMMENDED: Significant AUC improvement detected.\n"
                )
            elif auc_delta > 0.005:
                f.write(
                    "? MARGINAL IMPROVEMENT: Small AUC gain. Consider business impact.\n"
                )
            elif auc_delta > -0.005:
                f.write("≈ SIMILAR PERFORMANCE: Models perform similarly.\n")
            else:
                f.write(
                    "✗ PREVIOUS MODEL PREFERRED: New model shows performance degradation.\n"
                )

        f.write("\nFor detailed metrics, see metrics.json\n")
        f.write("For visualizations, see generated plot files\n")

    logger.info(f"Saved comparison report to {report_path}")


def create_health_check_file(output_path: str) -> str:
    """Create a health check file to signal script completion."""
    health_path = output_path
    with open(health_path, "w") as f:
        f.write(f"healthy: {datetime.now().isoformat()}")
    return health_path


def main(
    input_paths: Dict[str, str],
    output_paths: Dict[str, str],
    environ_vars: Dict[str, str],
    job_args: argparse.Namespace,
) -> None:
    """
    Main entry point for LightGBM model evaluation script.
    Loads model and data, runs evaluation, and saves results.

    Args:
        input_paths (Dict[str, str]): Dictionary of input paths
        output_paths (Dict[str, str]): Dictionary of output paths
        environ_vars (Dict[str, str]): Dictionary of environment variables
        job_args (argparse.Namespace): Command line arguments
    """
    # Extract paths from parameters - using contract-defined logical names
    model_dir = input_paths.get("model_input", input_paths.get("model_dir"))
    eval_data_dir = input_paths.get("processed_data", input_paths.get("eval_data_dir"))
    output_eval_dir = output_paths.get(
        "eval_output", output_paths.get("output_eval_dir")
    )
    output_metrics_dir = output_paths.get(
        "metrics_output", output_paths.get("output_metrics_dir")
    )

    # Extract environment variables
    id_field = environ_vars.get("ID_FIELD", "id")
    label_field = environ_vars.get("LABEL_FIELD", "label")

    # Extract comparison mode environment variables
    comparison_mode = environ_vars.get("COMPARISON_MODE", "false").lower() == "true"
    previous_score_field = environ_vars.get("PREVIOUS_SCORE_FIELD", "")
    comparison_metrics = environ_vars.get("COMPARISON_METRICS", "all")
    statistical_tests = environ_vars.get("STATISTICAL_TESTS", "true").lower() == "true"
    comparison_plots = environ_vars.get("COMPARISON_PLOTS", "true").lower() == "true"

    # Guard rail: If PREVIOUS_SCORE_FIELD is empty, disable comparison mode
    if comparison_mode and (
        not previous_score_field or previous_score_field.strip() == ""
    ):
        logger.warning(
            "COMPARISON_MODE is enabled but PREVIOUS_SCORE_FIELD is empty. Disabling comparison mode."
        )
        comparison_mode = False

    logger.info(f"Comparison mode: {comparison_mode}")
    if comparison_mode:
        logger.info(f"Previous score field: {previous_score_field}")
        logger.info(f"Comparison metrics: {comparison_metrics}")
        logger.info(f"Statistical tests: {statistical_tests}")
        logger.info(f"Comparison plots: {comparison_plots}")

    # Log job info
    job_type = job_args.job_type
    logger.info(f"Running model evaluation with job_type: {job_type}")

    # Ensure output directories exist
    os.makedirs(output_eval_dir, exist_ok=True)
    os.makedirs(output_metrics_dir, exist_ok=True)

    logger.info("Starting model evaluation script")

    # Load model artifacts
    model, risk_tables, impute_dict, feature_columns, hyperparams = (
        load_model_artifacts(model_dir)
    )

    # Load and preprocess data with format detection
    df, input_format = load_eval_data(eval_data_dir)

    # Get ID and label columns before preprocessing
    id_col, label_col = get_id_label_columns(df, id_field, label_field)

    # Process the data - our updated preprocess_eval_data preserves all columns including id and label
    df = preprocess_eval_data(df, feature_columns, risk_tables, impute_dict)

    # No need to filter or re-add id and label columns since they're already preserved
    logger.info(f"Final evaluation DataFrame shape: {df.shape}")

    # Get the available features (those that exist in the DataFrame)
    available_features = [col for col in feature_columns if col in df.columns]

    # Log inference strategy
    logger.info("INFERENCE STRATEGY:")
    logger.info(f"  → Using {len(available_features)} features for model inference")
    logger.info(f"  → These are the exact features the model was trained on")
    logger.info(f"  → Input data will be filtered to these features only")

    # Check for comparison mode
    previous_scores = None
    if comparison_mode:
        if previous_score_field in df.columns:
            previous_scores = df[previous_score_field].values
            logger.info(
                f"Found previous model scores in column '{previous_score_field}' with {len(previous_scores)} values"
            )
        else:
            logger.warning(
                f"Comparison mode enabled but column '{previous_score_field}' not found in data. Proceeding with standard evaluation."
            )
            comparison_mode = False

    # Evaluate model using the final DataFrame with both features and ID/label columns
    if comparison_mode and previous_scores is not None:
        # Enhanced evaluation with comparison
        evaluate_model_with_comparison(
            model,
            df,
            available_features,
            id_col,
            label_col,
            previous_scores,
            hyperparams,
            output_eval_dir,
            output_metrics_dir,
            comparison_metrics,
            statistical_tests,
            comparison_plots,
            input_format,
        )
    else:
        # Standard evaluation
        evaluate_model(
            model,
            df,
            available_features,  # Only use available feature columns for prediction
            id_col,
            label_col,
            hyperparams,
            output_eval_dir,
            output_metrics_dir,
            input_format,
        )

    logger.info("Model evaluation script complete")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--job_type", type=str, required=True)
    args = parser.parse_args()

    # Set up paths using contract-defined paths only
    input_paths = {
        "model_input": CONTAINER_PATHS["MODEL_DIR"],
        "processed_data": CONTAINER_PATHS["EVAL_DATA_DIR"],
    }

    output_paths = {
        "eval_output": CONTAINER_PATHS["OUTPUT_EVAL_DIR"],
        "metrics_output": CONTAINER_PATHS["OUTPUT_METRICS_DIR"],
    }

    # Collect environment variables - ID_FIELD and LABEL_FIELD are required per contract
    environ_vars = {
        "ID_FIELD": os.environ.get("ID_FIELD", "id"),  # Fallback for testing
        "LABEL_FIELD": os.environ.get("LABEL_FIELD", "label"),  # Fallback for testing
        "COMPARISON_MODE": os.environ.get("COMPARISON_MODE", "false"),
        "PREVIOUS_SCORE_FIELD": os.environ.get("PREVIOUS_SCORE_FIELD", ""),
        "COMPARISON_METRICS": os.environ.get("COMPARISON_METRICS", "all"),
        "STATISTICAL_TESTS": os.environ.get("STATISTICAL_TESTS", "true"),
        "COMPARISON_PLOTS": os.environ.get("COMPARISON_PLOTS", "true"),
    }

    try:
        # Call main function with testability parameters
        main(input_paths, output_paths, environ_vars, args)

        # Signal success
        success_path = os.path.join(output_paths["metrics_output"], "_SUCCESS")
        Path(success_path).touch()
        logger.info(f"Created success marker: {success_path}")

        # Create health check file
        health_path = os.path.join(output_paths["metrics_output"], "_HEALTH")
        create_health_check_file(health_path)
        logger.info(f"Created health check file: {health_path}")

        sys.exit(0)
    except Exception as e:
        # Log error and create failure marker
        logger.exception(f"Script failed with error: {e}")
        failure_path = os.path.join(
            output_paths.get("metrics_output", "/tmp"), "_FAILURE"
        )
        os.makedirs(os.path.dirname(failure_path), exist_ok=True)
        with open(failure_path, "w") as f:
            f.write(f"Error: {str(e)}")
        sys.exit(1)
