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,
    ):
        self._metrics = dict(metrics)
        self._model_type = model_type
        self.confusion_matrix = confusion_matrix
        self.class_labels = class_labels

    # ── 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)

    # ── 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>"

        html = (
            f'<div style="margin:8px 0">'
            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) ─────────────────────
        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>"
                )

            html += (
                f'<table style="{self._TBL};margin-top:10px">'
                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>"
            )

        html += "</div>"
        return html
