import logging
import io
import inspect
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from PIL import Image
from shap.explainers._explainer import Explainer
from coded_flows.types import DataFrame, Any, Tuple, MediaData, Union, PILImage
from coded_flows.utils import CodedFlowsLogger

matplotlib.use("Agg")
import shap

logger = CodedFlowsLogger(name="SHAP Explanation", level=logging.INFO)
logging.getLogger("shap").setLevel(logging.ERROR)


def _call_explainer(explainer, data, **base_kwargs):
    sig = inspect.signature(explainer.__call__)
    if "silent" in sig.parameters:
        return explainer(data, silent=True, **base_kwargs)
    return explainer(data, **base_kwargs)


def _to_numpy(df):
    """Convert pandas or polars DataFrame to numpy array."""
    if hasattr(df, "__class__") and "polars" in df.__class__.__module__:
        return df.to_numpy()
    elif hasattr(df, "values"):
        return df.values
    elif hasattr(df, "to_numpy"):
        return df.to_numpy()
    else:
        return np.array(df)


def _get_columns(df):
    """Get column names from pandas or polars DataFrame."""
    if hasattr(df, "columns"):
        columns = df.columns
        return list(columns) if not isinstance(columns, list) else columns
    return None


def shap_explanation(
    explainer: Any,
    data: DataFrame,
    original_data: DataFrame = None,
    options: dict = None,
) -> Tuple[Union[MediaData, PILImage], DataFrame]:
    options = options or {}
    verbose = options.get("verbose", True)
    row_index = options.get("row_index", 0)
    class_index = options.get("class_index", 0)
    max_display = options.get("max_display", 20)
    image_format = options.get("image_format", "pil")
    image = None
    explanation = pd.DataFrame()
    try:
        verbose and logger.info(
            f"Starting SHAP Individual Waterfall for row {row_index}."
        )
        if not isinstance(explainer, Explainer):
            verbose and logger.error("Expects a Shap Explainer as an input.")
            raise ValueError("Expects a Shap Explainer as an input.")
        shap_explanation = _call_explainer(explainer, data)
        if original_data is not None:
            verbose and logger.info(
                "Original data provided. Overwriting explanation data."
            )
            columns = _get_columns(original_data)
            if columns is not None:
                shap_explanation.feature_names = columns
            shap_explanation.data = _to_numpy(original_data)
        if row_index >= shap_explanation.shape[0] or row_index < 0:
            raise ValueError(
                f"Row index {row_index} is out of bounds (0-{shap_explanation.shape[0] - 1})."
            )
        ndim = len(shap_explanation.shape)
        single_explanation = None
        if ndim == 2:
            verbose and logger.info("Detected Regression or Binary Classification.")
            single_explanation = shap_explanation[row_index]
        elif ndim == 3:
            num_classes = shap_explanation.shape[2]
            verbose and logger.info(
                f"Detected Multiclass output with {num_classes} classes."
            )
            if class_index >= num_classes or class_index < 0:
                verbose and logger.warning(
                    f"Class index {class_index} out of bounds. Defaulting to 0."
                )
                class_index = 0
            verbose and logger.info(f"Selecting class index: {class_index}")
            single_explanation = shap_explanation[row_index, :, class_index]
        else:
            raise ValueError(f"Unexpected SHAP explanation dimensionality: {ndim}")
        feature_names = single_explanation.feature_names
        if not feature_names:
            feature_names = [f"Feature {i}" for i in range(single_explanation.shape[0])]
        values = np.array(single_explanation.data).flatten()
        contributions = np.array(single_explanation.values).flatten()
        rows = []
        for f, v, c in zip(feature_names, values, contributions):
            rows.append(
                {
                    "feature": f,
                    "value": v,
                    "contribution": c,
                    "abs_contribution": abs(c),
                }
            )
        rows.sort(key=lambda x: x["abs_contribution"], reverse=True)
        for r in rows:
            del r["abs_contribution"]
        base_value = single_explanation.base_values
        if isinstance(base_value, (np.ndarray, list)):
            base_value = base_value if np.ndim(base_value) == 0 else base_value[0]
        rows.insert(
            0,
            {
                "feature": "Base Value (Expected)",
                "value": np.nan,
                "contribution": base_value,
            },
        )
        explanation = pd.DataFrame(rows)
        verbose and logger.info(
            "Constructed sorted contribution DataFrame (Base Value + Features)."
        )
        verbose and logger.info("Generating Waterfall Plot.")
        fig = plt.figure()
        shap.plots.waterfall(single_explanation, max_display=max_display, show=False)
        buf = io.BytesIO()
        plt.savefig(buf, format="png", bbox_inches="tight", dpi=300)
        buf.seek(0)
        plt.close(plt.gcf())
        plt.close(fig)
        if image_format == "pil":
            image = Image.open(buf)
        elif image_format == "array":
            img = Image.open(buf)
            image = np.array(img)
        elif image_format == "bytes":
            image = buf.getvalue()
        verbose and logger.info("Image generation complete.")
    except Exception as e:
        verbose and logger.error(f"Error computing SHAP waterfall: {str(e)}")
        raise e
    return (image, explanation)