import numpy as np


class EvaluationResult:
    """
    Wrapper for evaluation metrics.

    - Renders as an HTML table in Jupyter environments.
    - Prints cleanly in terminals / scripts.
    - Supports dict-like access for backward compatibility.
    """

    def __init__(
        self,
        metrics: dict,
        model_type: str = "",
        confusion_matrix=None,
        class_labels=None,
        y_true=None,
        y_pred=None,
        y_proba=None,
    ):
        self._metrics = dict(metrics)
        self._model_type = model_type
        self.confusion_matrix = confusion_matrix
        self.class_labels = class_labels
        self._y_true = y_true
        self._y_pred = y_pred
        self._y_proba = y_proba

    # ── Programmatic access ─────────────────────────────

    def to_dict(self) -> dict:
        """Return the raw metrics dictionary."""
        return dict(self._metrics)

    # ── Dict-like interface ─────────────────────────────

    def __getitem__(self, key):
        return self._metrics[key]

    def __contains__(self, key):
        return key in self._metrics

    def __iter__(self):
        return iter(self._metrics)

    def __len__(self):
        return len(self._metrics)

    def keys(self):
        return self._metrics.keys()

    def values(self):
        return self._metrics.values()

    def items(self):
        return self._metrics.items()

    def get(self, key, default=None):
        return self._metrics.get(key, default)

    def summary(self) -> str:
        """
        Return a one-line summary of the evaluation results.

        Examples:
            >>> result.summary()
            'Classification: accuracy=0.95, f1_score=0.94'

            >>> result.summary()
            'Regression: r2_score=0.88, rmse=2.34'
        """
        if not self._metrics:
            return "No metrics available"

        # Identify the primary metric based on model type
        model_type = self._model_type.lower() if self._model_type else ""

        if "classification" in model_type:
            # Priority: accuracy, then F1, then ROC-AUC
            primary = self._metrics.get("accuracy") or self._metrics.get("f1_score")
            secondary = self._metrics.get("precision", 0), self._metrics.get("recall", 0)
            return (
                f"Classification: accuracy={self._metrics.get('accuracy', 0):.3f}, "
                f"precision={secondary[0]:.3f}, recall={secondary[1]:.3f}, "
                f"f1_score={self._metrics.get('f1_score', 0):.3f}"
            )
        elif "regression" in model_type:
            # Priority: R², then RMSE, then MAE
            return (
                f"Regression: r2_score={self._metrics.get('r2_score', 0):.3f}, "
                f"rmse={self._metrics.get('rmse', 0):.3f}, "
                f"mae={self._metrics.get('mae', 0):.3f}"
            )
        else:
            # Generic: show top 3 metrics
            items = list(self._metrics.items())[:3]
            parts = [f"{k}={v:.3f}" for k, v in items]
            return f"Evaluation: {', '.join(parts)}"

    # ── Terminal display ────────────────────────────────

    def __repr__(self):
        title = self._model_type.title() if self._model_type else "Evaluation"
        max_key = max(len(k) for k in self._metrics)
        col_w = max(max_key, 6)
        sep = "─" * (col_w + 16)

        lines = [f"{title} Results", sep]
        for key, value in self._metrics.items():
            lines.append(f"  {key:<{col_w}}  │  {value:.6f}")

        # ── Confusion matrix (text) ─────────────────────
        if self.confusion_matrix is not None:
            cm = np.asarray(self.confusion_matrix)
            labels = self.class_labels or list(range(cm.shape[0]))
            str_labels = [str(l) for l in labels]
            lw = max(len(s) for s in str_labels)
            cw = max(lw, max(len(str(v)) for v in cm.flat))

            lines.append("")
            lines.append("Confusion Matrix")
            lines.append(sep)

            # header row
            header = " " * (lw + 2) + "  ".join(f"{s:>{cw}}" for s in str_labels)
            lines.append(header)

            # data rows
            for i, row_label in enumerate(str_labels):
                row_vals = "  ".join(f"{cm[i, j]:>{cw}}" for j in range(cm.shape[1]))
                lines.append(f"{row_label:>{lw}}  {row_vals}")

        return "\n".join(lines)

    # ── Jupyter display ─────────────────────────────────

    # shared styles
    _TBL = (
        "border-collapse:collapse;font-size:14px;"
        "border:1px solid #c8ccd0;border-radius:8px;overflow:hidden"
    )
    _TH_TITLE = (
        "padding:10px 18px;text-align:left;color:#fff;"
        "font-weight:700;font-size:15px;background:#3b3f5c"
    )
    _TH_COL = (
        "padding:8px 18px;text-align:left;color:#1a1a2e;"
        "font-weight:600;background:#e8ecf1"
    )
    _TD_KEY = "padding:7px 18px;font-weight:500;color:#1a1a2e"
    _TD_VAL = "padding:7px 18px;font-family:monospace;color:#1a1a2e;text-align:right"
    _ROW_ALT = "background:#f6f8fb"

    def _repr_html_(self):
        title = self._model_type.title() if self._model_type else "Evaluation"

        items = list(self._metrics.items())
        n = len(items)
        mid = (n + 1) // 2  # left column gets the extra row if odd

        # ── Build 2-column metric rows ──────────────────
        metric_rows = ""
        for r in range(mid):
            bg = f' style="{self._ROW_ALT}"' if r % 2 == 1 else ""
            k1, v1 = items[r]
            cell_left = (
                f'<td style="{self._TD_KEY}">{k1}</td>'
                f'<td style="{self._TD_VAL}">{v1:.6f}</td>'
            )
            ri = r + mid
            if ri < n:
                k2, v2 = items[ri]
                cell_right = (
                    f'<td style="{self._TD_KEY};border-left:1px solid #d5dbe3">{k2}</td>'
                    f'<td style="{self._TD_VAL}">{v2:.6f}</td>'
                )
            else:
                cell_right = (
                    f'<td style="{self._TD_KEY};border-left:1px solid #d5dbe3"></td>'
                    f'<td style="{self._TD_VAL}"></td>'
                )
            metric_rows += f"<tr{bg}>{cell_left}{cell_right}</tr>"

        metrics_table = (
            f'<table style="{self._TBL}">'
            f"<thead>"
            f'<tr><th style="{self._TH_TITLE}" colspan="4">'
            f"{title} Results</th></tr>"
            f'<tr>'
            f'<th style="{self._TH_COL}">Metric</th>'
            f'<th style="{self._TH_COL};text-align:right">Value</th>'
            f'<th style="{self._TH_COL};border-left:1px solid #d5dbe3">Metric</th>'
            f'<th style="{self._TH_COL};text-align:right">Value</th>'
            f'</tr>'
            f"</thead>"
            f"<tbody>{metric_rows}</tbody>"
            f"</table>"
        )

        # ── Confusion matrix (HTML) ─────────────────────
        cm_table = ""
        if self.confusion_matrix is not None:
            cm = np.asarray(self.confusion_matrix)
            labels = self.class_labels or list(range(cm.shape[0]))
            cm_max = cm.max()

            header_cells = "".join(
                f'<th style="padding:8px 16px;text-align:center;'
                f'color:#1a1a2e;font-weight:600;min-width:56px">{l}</th>'
                for l in labels
            )

            cm_rows = ""
            for i, row_label in enumerate(labels):
                cells = ""
                for j in range(cm.shape[1]):
                    val = int(cm[i, j])
                    intensity = val / cm_max if cm_max > 0 else 0
                    if i == j:
                        bg = f"rgba(16,185,129,{intensity * 0.5:.2f})"
                    else:
                        bg = (
                            f"rgba(239,68,68,{intensity * 0.45:.2f})"
                            if val > 0 else "transparent"
                        )
                    cells += (
                        f'<td style="padding:8px 16px;text-align:center;'
                        f'font-family:monospace;font-weight:600;'
                        f'color:#1a1a2e;background:{bg}">'
                        f"{val}</td>"
                    )
                cm_rows += (
                    f"<tr>"
                    f'<th style="padding:8px 16px;text-align:right;'
                    f'font-weight:600;color:#1a1a2e">{row_label}</th>'
                    f"{cells}</tr>"
                )

            cm_table = (
                f'<table style="{self._TBL}">'
                f"<thead>"
                f'<tr><th style="{self._TH_TITLE}" '
                f'colspan="{len(labels) + 1}">Confusion Matrix</th></tr>'
                f'<tr style="background:#e8ecf1">'
                f'<th style="padding:8px 16px;text-align:right;'
                f'font-size:12px;color:#555">True ↓ &nbsp;Pred →</th>'
                f"{header_cells}</tr>"
                f"</thead>"
                f"<tbody>{cm_rows}</tbody>"
                f"</table>"
            )

        # ── Layout: side-by-side if confusion matrix exists ──
        if cm_table:
            html = (
                f'<div style="display:flex;gap:12px;align-items:flex-start;'
                f'margin:8px 0;flex-wrap:wrap">'
                f"{metrics_table}"
                f"{cm_table}"
                f"</div>"
            )
        else:
            html = f'<div style="margin:8px 0">{metrics_table}</div>'

        return html

    # ── Plotting ────────────────────────────────────────

    def plot(self):
        """
        Auto-detect and display up to 4 diagnostic plots
        in a single 2×2 figure.

        Classification: confusion matrix, ROC curve,
        precision-recall curve, per-class metric bars.

        Regression: actual vs predicted, residuals,
        residual distribution, residuals vs predicted.
        """
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            raise ImportError(
                "matplotlib is required for plotting. "
                "Install it with: pip install matplotlib"
            )

        if self._y_true is None or self._y_pred is None:
            raise RuntimeError(
                "Plot data not available. "
                "Use snaplab.evaluate() to generate plottable results."
            )

        # Detect plot type by capabilities, not by name
        has_cm = self.confusion_matrix is not None
        has_proba = self._y_proba is not None

        if has_cm:
            return self._plot_classification(plt)
        else:
            return self._plot_regression(plt)

    # ── Classification plots ────────────────────────────

    def _plot_classification(self, plt):
        from sklearn.metrics import (
            roc_curve, auc,
            precision_recall_curve, average_precision_score,
            precision_score, recall_score, f1_score,
        )

        ACCENT = "#3b3f5c"
        GREEN = "#10b981"
        RED = "#ef4444"
        GRAY = "#94a3b8"

        y_true = self._y_true
        y_pred = self._y_pred
        y_proba = self._y_proba
        cm = np.asarray(self.confusion_matrix)
        labels = self.class_labels or list(range(cm.shape[0]))
        is_binary = len(labels) == 2

        fig, axes = plt.subplots(2, 2, figsize=(11, 9))
        fig.suptitle("Classification Diagnostics", fontsize=16,
                     fontweight="bold", color=ACCENT, y=0.98)

        # ── 1. Confusion matrix heatmap ─────────────────
        ax = axes[0, 0]
        im = ax.imshow(cm, interpolation="nearest", cmap="Blues")
        ax.set_title("Confusion Matrix", fontweight="600", fontsize=12)
        ax.set_xlabel("Predicted", fontsize=10)
        ax.set_ylabel("Actual", fontsize=10)
        ax.set_xticks(range(len(labels)))
        ax.set_yticks(range(len(labels)))
        ax.set_xticklabels(labels)
        ax.set_yticklabels(labels)
        # annotate cells
        thresh = cm.max() / 2
        for i in range(cm.shape[0]):
            for j in range(cm.shape[1]):
                ax.text(j, i, str(cm[i, j]),
                        ha="center", va="center", fontsize=12, fontweight="600",
                        color="white" if cm[i, j] > thresh else ACCENT)

        # ── 2. ROC curve ────────────────────────────────
        ax = axes[0, 1]
        if y_proba is not None and is_binary:
            fpr, tpr, _ = roc_curve(y_true, y_proba[:, 1])
            roc_auc = auc(fpr, tpr)
            ax.plot(fpr, tpr, color=ACCENT, linewidth=2.5,
                    label=f"AUC = {roc_auc:.3f}")
            ax.fill_between(fpr, tpr, alpha=0.12, color=ACCENT)
            ax.plot([0, 1], [0, 1], color=GRAY, linestyle="--", linewidth=1)
            ax.set_title("ROC Curve", fontweight="600", fontsize=12)
            ax.set_xlabel("False Positive Rate", fontsize=10)
            ax.set_ylabel("True Positive Rate", fontsize=10)
            ax.legend(loc="lower right", fontsize=10)
        elif y_proba is not None:
            # multiclass: one curve per class
            for idx, label in enumerate(labels):
                y_bin = (y_true == label).astype(int)
                fpr, tpr, _ = roc_curve(y_bin, y_proba[:, idx])
                ax.plot(fpr, tpr, linewidth=1.8, label=f"Class {label}")
            ax.plot([0, 1], [0, 1], color=GRAY, linestyle="--", linewidth=1)
            ax.set_title("ROC Curves (per class)", fontweight="600", fontsize=12)
            ax.set_xlabel("False Positive Rate", fontsize=10)
            ax.set_ylabel("True Positive Rate", fontsize=10)
            ax.legend(fontsize=8, loc="lower right")
        else:
            ax.text(0.5, 0.5, "ROC curve\nnot available\n(no predict_proba)",
                    ha="center", va="center", fontsize=11, color=GRAY,
                    transform=ax.transAxes)
            ax.set_title("ROC Curve", fontweight="600", fontsize=12, color=GRAY)

        # ── 3. Precision-Recall curve ───────────────────
        ax = axes[1, 0]
        if y_proba is not None and is_binary:
            prec, rec, _ = precision_recall_curve(y_true, y_proba[:, 1])
            ap = average_precision_score(y_true, y_proba[:, 1])
            ax.plot(rec, prec, color=GREEN, linewidth=2.5,
                    label=f"AP = {ap:.3f}")
            ax.fill_between(rec, prec, alpha=0.12, color=GREEN)
            ax.set_title("Precision-Recall Curve", fontweight="600", fontsize=12)
            ax.set_xlabel("Recall", fontsize=10)
            ax.set_ylabel("Precision", fontsize=10)
            ax.legend(loc="lower left", fontsize=10)
        elif y_proba is not None:
            for idx, label in enumerate(labels):
                y_bin = (y_true == label).astype(int)
                prec, rec, _ = precision_recall_curve(y_bin, y_proba[:, idx])
                ax.plot(rec, prec, linewidth=1.8, label=f"Class {label}")
            ax.set_title("PR Curves (per class)", fontweight="600", fontsize=12)
            ax.set_xlabel("Recall", fontsize=10)
            ax.set_ylabel("Precision", fontsize=10)
            ax.legend(fontsize=8, loc="lower left")
        else:
            ax.text(0.5, 0.5, "PR curve\nnot available\n(no predict_proba)",
                    ha="center", va="center", fontsize=11, color=GRAY,
                    transform=ax.transAxes)
            ax.set_title("Precision-Recall", fontweight="600", fontsize=12,
                         color=GRAY)

        # ── 4. Per-class metric bars ────────────────────
        ax = axes[1, 1]
        per_prec = precision_score(y_true, y_pred, average=None,
                                   zero_division=0, labels=labels)
        per_rec = recall_score(y_true, y_pred, average=None,
                               zero_division=0, labels=labels)
        per_f1 = f1_score(y_true, y_pred, average=None,
                          zero_division=0, labels=labels)

        x = np.arange(len(labels))
        w = 0.25
        ax.bar(x - w, per_prec, w, label="Precision", color=ACCENT, alpha=0.85)
        ax.bar(x, per_rec, w, label="Recall", color=GREEN, alpha=0.85)
        ax.bar(x + w, per_f1, w, label="F1", color="#f59e0b", alpha=0.85)
        ax.set_xticks(x)
        ax.set_xticklabels([str(l) for l in labels])
        ax.set_ylim(0, 1.15)
        ax.set_title("Per-Class Metrics", fontweight="600", fontsize=12)
        ax.set_xlabel("Class", fontsize=10)
        ax.legend(fontsize=9, loc="upper right")

        fig.tight_layout(rect=[0, 0, 1, 0.95])
        plt.show()
        return fig

    # ── Regression plots ────────────────────────────────

    def _plot_regression(self, plt):
        ACCENT = "#3b3f5c"
        BLUE = "#3b82f6"
        CORAL = "#f97316"

        y_true = self._y_true
        y_pred = self._y_pred
        residuals = y_true - y_pred

        fig, axes = plt.subplots(2, 2, figsize=(11, 9))
        fig.suptitle("Regression Diagnostics", fontsize=16,
                     fontweight="bold", color=ACCENT, y=0.98)

        # ── 1. Actual vs Predicted ──────────────────────
        ax = axes[0, 0]
        ax.scatter(y_true, y_pred, alpha=0.5, s=28, color=BLUE, edgecolors="white",
                   linewidth=0.4)
        mn = min(y_true.min(), y_pred.min())
        mx = max(y_true.max(), y_pred.max())
        ax.plot([mn, mx], [mn, mx], color=ACCENT, linestyle="--", linewidth=1.5,
                label="Ideal")
        ax.set_title("Actual vs Predicted", fontweight="600", fontsize=12)
        ax.set_xlabel("Actual", fontsize=10)
        ax.set_ylabel("Predicted", fontsize=10)
        ax.legend(fontsize=9)

        # ── 2. Residuals vs Predicted ───────────────────
        ax = axes[0, 1]
        ax.scatter(y_pred, residuals, alpha=0.5, s=28, color=CORAL,
                   edgecolors="white", linewidth=0.4)
        ax.axhline(0, color=ACCENT, linestyle="--", linewidth=1.5)
        ax.set_title("Residuals vs Predicted", fontweight="600", fontsize=12)
        ax.set_xlabel("Predicted", fontsize=10)
        ax.set_ylabel("Residual", fontsize=10)

        # ── 3. Residual distribution ────────────────────
        ax = axes[1, 0]
        ax.hist(residuals, bins="auto", color=BLUE, alpha=0.7,
                edgecolor="white", linewidth=0.5)
        ax.axvline(0, color=ACCENT, linestyle="--", linewidth=1.5)
        ax.set_title("Residual Distribution", fontweight="600", fontsize=12)
        ax.set_xlabel("Residual", fontsize=10)
        ax.set_ylabel("Count", fontsize=10)

        # ── 4. Sorted residuals (scale-location) ───────
        ax = axes[1, 1]
        abs_res = np.abs(residuals)
        sorted_idx = np.argsort(y_pred)
        ax.scatter(y_pred[sorted_idx], np.sqrt(abs_res[sorted_idx]),
                   alpha=0.5, s=28, color=CORAL, edgecolors="white",
                   linewidth=0.4)
        ax.set_title("Scale-Location", fontweight="600", fontsize=12)
        ax.set_xlabel("Predicted", fontsize=10)
        ax.set_ylabel("√|Residual|", fontsize=10)

        fig.tight_layout(rect=[0, 0, 1, 0.95])
        plt.show()
        return fig
