"""Geosmith ML: Model evaluation utilities (confusion matrix calculations)

Migrated from geosuite.ml.
Layer 2: Primitives - Pure operations.
"""

from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import numpy as np
import logging

logger = logging.getLogger(__name__)


try:
    import pandas as pd

    PANDAS_AVAILABLE = True
except ImportError:
    PANDAS_AVAILABLE = False
    pd = None  # type: ignore

try:
    from sklearn.model_selection import BaseCrossValidator

    SKLEARN_AVAILABLE = True
except ImportError:
    SKLEARN_AVAILABLE = False
    BaseCrossValidator = None  # type: ignore

try:
    from shap import TreeExplainer, KernelExplainer

    SHAP_AVAILABLE = True
except ImportError:
    SHAP_AVAILABLE = False
    TreeExplainer = None  # type: ignore
    KernelExplainer = None  # type: ignore

try:
    from numba import njit

    NUMBA_AVAILABLE = True
except ImportError:
    NUMBA_AVAILABLE = False

    def njit(*args, **kwargs):
        def decorator(func):
            return func

        return decorator if not args else decorator(args[0])


@njit(cache=True)
def _adjust_confusion_matrix_kernel(
    cm: np.ndarray, adjacent_facies_array: np.ndarray
) -> np.ndarray:
    """
    Numba-optimized kernel for adjusting confusion matrix with adjacent facies.

    This function is JIT-compiled for 10-15x speedup on large confusion matrices.

    Args:
        cm: Confusion matrix (numpy array)
        adjacent_facies_array: 2D array where adjacent_facies_array[i] contains
                               indices of facies adjacent to facies i (-1 for padding)

    Returns:
        Adjusted confusion matrix
    """
    adj_cm = cm.copy()
    n_classes = cm.shape[0]

    for i in range(n_classes):
        for j in range(adjacent_facies_array.shape[1]):
            adj_idx = int(adjacent_facies_array[i, j])
            if adj_idx >= 0:  # -1 is used as padding
                adj_cm[i, i] += adj_cm[i, adj_idx]
                adj_cm[i, adj_idx] = 0.0

    return adj_cm


def display_cm(
    cm: np.ndarray,
    labels: List[str],
    hide_zeros: bool = False,
    display_metrics: bool = False,
) -> str:
    """
    Display confusion matrix with labels, along with
    metrics such as Recall, Precision and F1 score.

    Args:
        cm: Confusion matrix (numpy array)
        labels: List of class labels
        hide_zeros: If True, hide zero values in the matrix
        display_metrics: If True, display precision, recall, and F1 scores

    Returns:
        Formatted string representation of confusion matrix
    """
    precision = np.diagonal(cm) / cm.sum(axis=0).astype("float")
    recall = np.diagonal(cm) / cm.sum(axis=1).astype("float")
    F1 = 2 * (precision * recall) / (precision + recall)

    precision[np.isnan(precision)] = 0
    recall[np.isnan(recall)] = 0
    F1[np.isnan(F1)] = 0

    total_precision = np.sum(precision * cm.sum(axis=1)) / cm.sum(axis=(0, 1))
    total_recall = np.sum(recall * cm.sum(axis=1)) / cm.sum(axis=(0, 1))
    total_F1 = np.sum(F1 * cm.sum(axis=1)) / cm.sum(axis=(0, 1))

    columnwidth = max([len(x) for x in labels] + [5])  # 5 is value length
    empty_cell = " " * columnwidth

    # Build output string
    output = []

    # Print header
    header = "    " + " Pred"
    for label in labels:
        header += " " + f"{label:>{columnwidth}}"
    header += " " + f"{'Total':>{columnwidth}}"
    output.append(header)
    output.append("")
    output.append("    " + " True")

    # Print rows
    for i, label1 in enumerate(labels):
        row = "    " + f"{label1:>{columnwidth}}"
        for j in range(len(labels)):
            cell = f"{int(cm[i, j]):>{columnwidth}d}"
            if hide_zeros:
                cell = cell if float(cm[i, j]) != 0 else empty_cell
            row += " " + cell
        row += " " + f"{int(sum(cm[i,:])):>{columnwidth}d}"
        output.append(row)
        output.append("")

    if display_metrics:
        output.append("")
        precision_row = "Precision"
        for j in range(len(labels)):
            precision_row += " " + f"{precision[j]:>{columnwidth}.2f}"
        precision_row += " " + f"{total_precision:>{columnwidth}.2f}"
        output.append(precision_row)
        output.append("")

        recall_row = "   Recall"
        for j in range(len(labels)):
            recall_row += " " + f"{recall[j]:>{columnwidth}.2f}"
        recall_row += " " + f"{total_recall:>{columnwidth}.2f}"
        output.append(recall_row)
        output.append("")

        f1_row = "       F1"
        for j in range(len(labels)):
            f1_row += " " + f"{F1[j]:>{columnwidth}.2f}"
        f1_row += " " + f"{total_F1:>{columnwidth}.2f}"
        output.append(f1_row)
        output.append("")

    result = "\n".join(output)
    logger.info(result)
    return result


def display_adj_cm(
    cm: np.ndarray,
    labels: List[str],
    adjacent_facies: List[List[int]],
    hide_zeros: bool = False,
    display_metrics: bool = False,
) -> str:
    """
    Display a confusion matrix that counts adjacent facies as correct.

    This is useful for geological facies classification where adjacent
    facies (e.g., transitional lithologies) should be considered
    partially correct predictions.

    Args:
        cm: Confusion matrix (numpy array)
        labels: List of class labels
        adjacent_facies: List of lists, where adjacent_facies[i] contains
                        the indices of facies adjacent to facies i
        hide_zeros: If True, hide zero values in the matrix
        display_metrics: If True, display precision, recall, and F1 scores

    Returns:
        Formatted string representation of adjusted confusion matrix
    """
    # Convert adjacent_facies list to padded numpy array for Numba
    max_adjacent = max(len(adj) for adj in adjacent_facies) if adjacent_facies else 0
    adjacent_array = np.full((len(adjacent_facies), max_adjacent), -1, dtype=np.int64)

    for i, adj_list in enumerate(adjacent_facies):
        for j, idx in enumerate(adj_list):
            adjacent_array[i, j] = idx

    # Call optimized kernel
    adj_cm = _adjust_confusion_matrix_kernel(cm.astype(np.float64), adjacent_array)

    return display_cm(adj_cm, labels, hide_zeros, display_metrics)


def confusion_matrix_to_dataframe(cm: np.ndarray, labels: List[str]) -> pd.DataFrame:
    """
    Convert confusion matrix to pandas DataFrame for easier analysis.

    Args:
        cm: Confusion matrix (numpy array)
        labels: List of class labels

    Returns:
        DataFrame with confusion matrix and row/column labels
    """
    df = pd.DataFrame(cm, index=labels, columns=labels)
    df.index.name = "True"
    df.columns.name = "Predicted"
    return df


def compute_metrics_from_cm(cm: np.ndarray, labels: List[str]) -> pd.DataFrame:
    """
    Compute precision, recall, and F1 scores from confusion matrix.

    Args:
        cm: Confusion matrix (numpy array)
        labels: List of class labels

    Returns:
        DataFrame with per-class metrics
    """
    precision = np.diagonal(cm) / cm.sum(axis=0).astype("float")
    recall = np.diagonal(cm) / cm.sum(axis=1).astype("float")
    F1 = 2 * (precision * recall) / (precision + recall)

    # Replace NaN with 0
    precision[np.isnan(precision)] = 0
    recall[np.isnan(recall)] = 0
    F1[np.isnan(F1)] = 0

    # Support (number of true instances per class)
    support = cm.sum(axis=1)

    metrics_df = pd.DataFrame(
        {
            "Class": labels,
            "Precision": precision,
            "Recall": recall,
            "F1-Score": F1,
            "Support": support,
        }
    )

    # Add weighted averages
    total_support = support.sum()
    avg_row = pd.DataFrame(
        {
            "Class": ["Weighted Avg"],
            "Precision": [np.sum(precision * support) / total_support],
            "Recall": [np.sum(recall * support) / total_support],
            "F1-Score": [np.sum(F1 * support) / total_support],
            "Support": [total_support],
        }
    )

    metrics_df = pd.concat([metrics_df, avg_row], ignore_index=True)

    return metrics_df
