"""Shapley attributions via shapiq + TabPFN imputation explainer (cloud-safe)."""

from __future__ import annotations

import io
import warnings
from typing import Any, Dict, List, Optional, Tuple

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from PIL import Image
from tabpfn_client import TabPFNClassifier
from tabpfn_extensions.interpretability.shapiq import get_tabpfn_imputation_explainer

from model import FEATURES, FEATURE_DISPLAY_NAMES

# One explainer per predicted class — background is fixed, no per-coalition refit.
_explainer_cache: Dict[int, object] = {}
_BACKGROUND_ROWS = 256
_SHAPIQ_BUDGET = 64

_POS_COLOR = "#e85d75"
_NEG_COLOR = "#4a9eff"
_BG_COLOR = "#1e0f14"
_TEXT_COLOR = "#f5e6d3"
_GRID_COLOR = "#722f37"


def clear_explainer_cache() -> None:
    """Drop cached explainers after reconnecting with a new TabPFN token."""
    _explainer_cache.clear()


def _get_explainer(
    clf: TabPFNClassifier,
    train_df: pd.DataFrame,
    class_index: int,
):
    if class_index not in _explainer_cache:
        n = min(_BACKGROUND_ROWS, len(train_df))
        background = train_df[FEATURES].sample(n=n, random_state=42).values
        _explainer_cache[class_index] = get_tabpfn_imputation_explainer(
            clf,
            background,
            index="SV",
            max_order=1,
            imputer="baseline",
            class_index=class_index,
            approximator="permutation",
            random_state=42,
        )
    return _explainer_cache[class_index]


def _shapley_pairs(sv) -> List[Tuple[str, float]]:
    rows = []
    for interaction, value in sv.dict_values.items():
        if len(interaction) != 1:
            continue
        key = FEATURES[interaction[0]]
        rows.append((FEATURE_DISPLAY_NAMES[key], float(value)))
    rows.sort(key=lambda item: abs(item[1]), reverse=True)
    return rows


def _shapley_bar_image(
    pairs: List[Tuple[str, float]],
    baseline: float,
    final: float,
) -> Image.Image:
    """Diverging bar chart — rasterized for reliable display in Gradio on HF."""
    if not pairs:
        fig, ax = plt.subplots(figsize=(10, 2), facecolor=_BG_COLOR)
        ax.set_facecolor(_BG_COLOR)
        ax.text(0.5, 0.5, "No attributions available", ha="center", va="center", color=_TEXT_COLOR)
        ax.axis("off")
    else:
        names = [name for name, _ in pairs]
        values = [value for _, value in pairs]
        height = max(5.0, len(names) * 0.42)
        fig, ax = plt.subplots(figsize=(10, height), facecolor=_BG_COLOR)
        ax.set_facecolor(_BG_COLOR)
        colors = [_POS_COLOR if v > 0 else _NEG_COLOR for v in values]
        y_pos = np.arange(len(names))
        ax.barh(y_pos, values, color=colors, height=0.72)
        ax.set_yticks(y_pos)
        ax.set_yticklabels(names, color=_TEXT_COLOR, fontsize=10)
        ax.axvline(0, color=_GRID_COLOR, linewidth=1)
        ax.tick_params(axis="x", colors=_TEXT_COLOR, labelcolor=_TEXT_COLOR)
        ax.set_xlabel("Shapley contribution (predicted-class probability)", color="#c9a87c")
        ax.set_title(
            f"Baseline {baseline:.3f}  →  prediction {final:.3f}",
            color=_TEXT_COLOR,
            fontsize=11,
            pad=10,
        )
        for spine in ax.spines.values():
            spine.set_color(_GRID_COLOR)
        ax.grid(axis="x", color=_GRID_COLOR, alpha=0.35, linestyle="--", linewidth=0.6)
        ax.invert_yaxis()
        fig.tight_layout()

    buf = io.BytesIO()
    fig.savefig(
        buf,
        format="png",
        dpi=150,
        bbox_inches="tight",
        facecolor=_BG_COLOR,
        edgecolor="none",
    )
    plt.close(fig)
    buf.seek(0)
    return Image.open(buf)


def _shapley_table(pairs: List[Tuple[str, float]]) -> pd.DataFrame:
    rows = []
    for name, value in pairs:
        magnitude = abs(value)
        bar_len = max(1, round(magnitude * 40)) if magnitude > 0 else 0
        rows.append({
            "Feature": name,
            "Shapley": round(value, 4),
            "Effect": (
                "supports predicted class" if value > 0
                else "works against predicted class" if value < 0
                else "neutral"
            ),
            "Bar": "█" * bar_len,
        })
    return pd.DataFrame(rows)


def explain_prediction(
    clf: TabPFNClassifier,
    train_df: pd.DataFrame,
    features: dict,
    class_index: int,
) -> Tuple[Optional[Image.Image], pd.DataFrame, Dict[str, Any]]:
    """Run shapiq SV attribution and return a bar-chart image, table, and metadata."""
    explainer = _get_explainer(clf, train_df, class_index)
    x = np.array([[features[f] for f in FEATURES]], dtype=float)
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            category=RuntimeWarning,
            module=r"shapiq\.approximator\.regression",
        )
        sv = explainer.explain(x, budget=_SHAPIQ_BUDGET)

    pairs = _shapley_pairs(sv)
    shapley_sum = sum(value for _, value in pairs)
    baseline = float(sv.baseline_value)
    final = baseline + shapley_sum

    image = _shapley_bar_image(pairs, baseline, final)
    meta = {
        "baseline_probability": round(baseline, 4),
        "final_probability": round(final, 4),
        "predicted_class_index": class_index,
        "attribution_method": "shapiq SV via permutation sampling (first-order)",
    }
    return image, _shapley_table(pairs), meta