"""Metrics + fairness + plots for binary cyberbullying detection.

Public API:
- ``compute_metrics(y_true, y_pred, y_proba)`` - full classification metrics
  (acc, P/R/F1 macro+weighted, ROC-AUC, MCC, Cohen's κ, balanced acc, CM).
- ``compute_fairness(y_true, y_pred)`` - Jayanti & Rohman (2026) label-as-group
  framework: accuracy gap, equal-opportunity gap, demographic-parity gap,
  precision gap.
- ``format_metrics_report(metrics, fairness)`` - pretty-print to string.

Plots (basic):
- ``plot_confusion_matrix``, ``plot_roc_curve``.

Plots (Jayanti & Rohman 2026 paper-figure replicas - usable for `indobert.ipynb`
paper repro AND `ibt_hybrid_eval.ipynb` proposed-model eval):
- ``plot_confusion_matrix_dual``       - Fig 3: raw + normalized side-by-side
- ``plot_per_class_acc_and_errors``    - Fig 4: per-class acc bars + FPR/FNR bars
- ``plot_confidence_distribution``     - Fig 5: overlay histogram of P(CB) by true class
- ``length_quartile_analysis``         - returns dict for the two length plots
- ``plot_accuracy_by_length``          - Fig 6: bar of accuracy per Q1–Q4
- ``plot_correct_vs_incorrect_by_length`` - Fig 7: stacked bar per Q1–Q4
- ``plot_paper_repro_suite``           - convenience: all 5 figures in one call

Label convention (Kaggle canonical, used throughout):
    0 = cyberbullying, 1 = non-cyberbullying
The paper figures invert the visual labels; the numerical metrics are
permutation-invariant (acc, F1m, ROC-AUC) or trivially flipped (per-class).
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Sequence

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    cohen_kappa_score,
    confusion_matrix,
    f1_score,
    matthews_corrcoef,
    precision_score,
    recall_score,
    roc_auc_score,
    roc_curve,
)

logger = logging.getLogger(__name__)

LABEL_NAMES = ("cyberbullying", "non-cyberbullying")  # index matches label value


def compute_metrics(
    y_true: Sequence[int],
    y_pred: Sequence[int],
    y_proba: np.ndarray,
) -> dict:
    """Full classification-report-grade metrics for binary {0,1} labels.

    ``y_proba`` must have shape ``(n, 2)``; column 1 (positive class) is used for ROC-AUC.

    Output schema:

    - aggregates: ``accuracy``, ``balanced_accuracy``, ``cohens_kappa``,
      ``precision_macro``, ``recall_macro``, ``f1_macro``, ``f1_weighted``,
      ``roc_auc``, ``mcc``
    - per-class arrays (index 0 = cyberbullying, index 1 = non-cyberbullying):
      ``precision_per_class``, ``recall_per_class``, ``f1_per_class``,
      ``support_per_class``
    - ``confusion_matrix``: ``{"tn": ..., "fp": ..., "fn": ..., "tp": ...}``
      where positive class = label 1 (non-cyberbullying)
    """
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)

    p_per = precision_score(y_true, y_pred, labels=[0, 1], average=None, zero_division=0)
    r_per = recall_score(y_true, y_pred, labels=[0, 1], average=None, zero_division=0)
    f1_per = f1_score(y_true, y_pred, labels=[0, 1], average=None, zero_division=0)
    support = np.bincount(y_true, minlength=2)

    cm = confusion_matrix(y_true, y_pred, labels=[0, 1])
    # rows = true label, cols = predicted; pos class = 1
    tn, fp, fn, tp = int(cm[0, 0]), int(cm[0, 1]), int(cm[1, 0]), int(cm[1, 1])

    return {
        "accuracy": float(accuracy_score(y_true, y_pred)),
        "balanced_accuracy": float(balanced_accuracy_score(y_true, y_pred)),
        "cohens_kappa": float(cohen_kappa_score(y_true, y_pred)),
        "precision_macro": float(precision_score(y_true, y_pred, average="macro", zero_division=0)),
        "recall_macro": float(recall_score(y_true, y_pred, average="macro", zero_division=0)),
        "f1_macro": float(f1_score(y_true, y_pred, average="macro", zero_division=0)),
        "f1_weighted": float(f1_score(y_true, y_pred, average="weighted", zero_division=0)),
        "roc_auc": float(roc_auc_score(y_true, y_proba[:, 1])),
        "mcc": float(matthews_corrcoef(y_true, y_pred)),
        "precision_per_class": [float(x) for x in p_per],
        "recall_per_class": [float(x) for x in r_per],
        "f1_per_class": [float(x) for x in f1_per],
        "support_per_class": [int(x) for x in support],
        "confusion_matrix": {"tn": tn, "fp": fp, "fn": fn, "tp": tp},
    }


def compute_fairness(y_true: Sequence[int], y_pred: Sequence[int]) -> dict:
    """Compute fairness gaps using the Jayanti & Rohman (2026) convention.

    The paper treats the true label itself as the protected group (no external
    demographic attribute is available). Under that convention:

    - ``accuracy_gap`` = ``|recall(0) − recall(1)|`` = ``|TNR − TPR|``
    - ``equal_opportunity_gap`` = ``|recall(0) − recall(1)|`` - the paper's
      formula collapses to the same value as ``accuracy_gap`` when the group is
      the label. Reported separately for direct numerical comparison with the paper.
    - ``demographic_parity_gap`` = ``|PPR(0) − PPR(1)|`` = ``|FPR − TPR|``,
      where ``PPR(k) = P(Ŷ=1 | Y=k)``.
    - ``precision_gap`` (bonus, not in paper) = ``|precision(0) − precision(1)|``
      = ``|NPV − PPV|`` - captures class-asymmetric error behavior that the
      paper's formula conflates.
    """
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)

    recall_0 = recall_score(y_true, y_pred, pos_label=0, zero_division=0)
    recall_1 = recall_score(y_true, y_pred, pos_label=1, zero_division=0)
    precision_0 = precision_score(y_true, y_pred, pos_label=0, zero_division=0)
    precision_1 = precision_score(y_true, y_pred, pos_label=1, zero_division=0)

    # PPR(k) = P(Ŷ=1 | Y=k)
    mask_y0 = (y_true == 0)
    mask_y1 = (y_true == 1)
    ppr_0 = (y_pred[mask_y0] == 1).mean() if mask_y0.any() else 0.0
    ppr_1 = (y_pred[mask_y1] == 1).mean() if mask_y1.any() else 0.0

    accuracy_gap = abs(recall_0 - recall_1)
    return {
        "accuracy_gap": float(accuracy_gap),
        "equal_opportunity_gap": float(accuracy_gap),  # paper-literal: collapses to accuracy_gap
        "demographic_parity_gap": float(abs(ppr_0 - ppr_1)),
        "precision_gap": float(abs(precision_0 - precision_1)),
        "per_class_accuracy": {"0": float(recall_0), "1": float(recall_1)},
    }


def format_metrics_report(metrics: dict, fairness: dict) -> str:
    """Pretty-printed single-string report for logging."""
    lines = ["=== Metrics ==="]
    for k, v in metrics.items():
        if isinstance(v, float):
            lines.append(f"  {k:22s} {v:.4f}")
        elif isinstance(v, list):
            joined = ", ".join(
                f"{x:.4f}" if isinstance(x, float) else str(x) for x in v
            )
            lines.append(f"  {k:22s} [{joined}]")
        elif isinstance(v, dict):
            joined = ", ".join(f"{kk}={vv}" for kk, vv in v.items())
            lines.append(f"  {k:22s} {{{joined}}}")
        else:
            lines.append(f"  {k:22s} {v}")
    lines.append("=== Fairness ===")
    for k, v in fairness.items():
        if k == "per_class_accuracy":
            lines.append(f"  {k:24s} {{0: {v['0']:.4f}, 1: {v['1']:.4f}}}")
        else:
            lines.append(f"  {k:24s} {v:.4f}")
    return "\n".join(lines)


def plot_confusion_matrix(
    y_true: Sequence[int],
    y_pred: Sequence[int],
    labels: Sequence[str] = LABEL_NAMES,
    normalize: bool = False,
    save_path: Path | None = None,
    ax=None,
    title: str | None = None,
) -> None:
    """Plot a confusion matrix heatmap. Axis 0 = true label, axis 1 = predicted label."""
    cm = confusion_matrix(y_true, y_pred, labels=[0, 1])
    if normalize:
        row_sums = cm.sum(axis=1, keepdims=True)
        cm_norm = np.divide(cm, row_sums, out=np.zeros_like(cm, dtype=float), where=row_sums != 0)
        data, fmt = cm_norm, ".2f"
    else:
        data, fmt = cm, "d"

    own_fig = ax is None
    if own_fig:
        _, ax = plt.subplots(figsize=(5.5, 4.5))
    sns.heatmap(
        data, annot=True, fmt=fmt, cmap="Blues",
        xticklabels=labels, yticklabels=labels, ax=ax, cbar=False,
    )
    ax.set_xlabel("Predicted")
    ax.set_ylabel("Actual")
    ax.set_title(title or ("Confusion matrix (normalized)" if normalize else "Confusion matrix"))

    if save_path is not None:
        plt.tight_layout()
        plt.savefig(str(save_path), dpi=150, bbox_inches="tight")


def plot_confusion_matrix_dual(
    y_true: Sequence[int],
    y_pred: Sequence[int],
    labels: Sequence[str] = LABEL_NAMES,
    save_path: Path | None = None,
    title_prefix: str = "",
) -> tuple:
    """Paper Fig 3 replica - raw counts + normalized percentages side-by-side.

    Returns ``(fig, (ax_raw, ax_norm))`` so the caller can further customize.
    """
    fig, (ax_raw, ax_norm) = plt.subplots(1, 2, figsize=(12, 4.5))
    plot_confusion_matrix(
        y_true, y_pred, labels=labels, normalize=False, ax=ax_raw,
        title=f"{title_prefix}Confusion Matrix".strip(),
    )
    plot_confusion_matrix(
        y_true, y_pred, labels=labels, normalize=True, ax=ax_norm,
        title=f"{title_prefix}Normalized Confusion Matrix".strip(),
    )
    plt.tight_layout()
    if save_path is not None:
        fig.savefig(str(save_path), dpi=150, bbox_inches="tight")
    return fig, (ax_raw, ax_norm)


def plot_per_class_acc_and_errors(
    y_true: Sequence[int],
    y_pred: Sequence[int],
    labels: Sequence[str] = LABEL_NAMES,
    save_path: Path | None = None,
    title_prefix: str = "",
) -> tuple:
    """Paper Fig 4 replica - per-class accuracy bars + FPR/FNR bars side-by-side.

    FPR (False Positive Rate) = error rate on non-cyberbullying samples =
        ``1 − recall(label=1)``.
    FNR (False Negative Rate) = error rate on cyberbullying samples =
        ``1 − recall(label=0)``.
    The semantics match the paper regardless of label convention.
    """
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)

    # Per-class accuracy = recall per class.
    recall_0 = recall_score(y_true, y_pred, pos_label=0, zero_division=0)
    recall_1 = recall_score(y_true, y_pred, pos_label=1, zero_division=0)
    cb_acc, ncb_acc = float(recall_0), float(recall_1)
    fpr, fnr = 1.0 - ncb_acc, 1.0 - cb_acc
    avg_acc = (cb_acc + ncb_acc) / 2.0

    fig, (ax_acc, ax_err) = plt.subplots(1, 2, figsize=(12, 4.5))

    bars = ax_acc.bar(
        labels, [cb_acc, ncb_acc],
        color=["#e07b91", "#7cc9a3"],  # CB pink-red, non-CB green
        edgecolor="#444", linewidth=0.8,
    )
    ax_acc.axhline(avg_acc, linestyle="--", color="#5b6cff", linewidth=1.2, label="Average")
    for b, v in zip(bars, [cb_acc, ncb_acc]):
        ax_acc.text(b.get_x() + b.get_width() / 2, v + 0.015, f"{v:.4f}",
                    ha="center", fontsize=10, fontweight="bold")
    ax_acc.set_ylim(0, 1.0)
    ax_acc.set_ylabel("Accuracy")
    ax_acc.set_title(f"{title_prefix}Per-Class Accuracy".strip())
    ax_acc.legend(loc="upper right")

    bars2 = ax_err.bar(
        ["False Positive\nRate", "False Negative\nRate"],
        [fpr, fnr],
        color=["#f5a960", "#e07b91"],
        edgecolor="#444", linewidth=0.8,
    )
    for b, v in zip(bars2, [fpr, fnr]):
        ax_err.text(b.get_x() + b.get_width() / 2, v + 0.008, f"{v:.4f}",
                    ha="center", fontsize=10, fontweight="bold")
    ax_err.set_ylim(0, max(0.5, max(fpr, fnr) * 1.25))
    ax_err.set_ylabel("Rate")
    ax_err.set_title(f"{title_prefix}Error Rates".strip())

    plt.tight_layout()
    if save_path is not None:
        fig.savefig(str(save_path), dpi=150, bbox_inches="tight")
    return fig, (ax_acc, ax_err)


def plot_confidence_distribution(
    y_true: Sequence[int],
    y_proba: np.ndarray,
    cb_label: int = 0,
    bins: int = 30,
    save_path: Path | None = None,
    ax=None,
    title: str | None = None,
) -> None:
    """Paper Fig 5 replica - overlay histogram of P(cyberbullying), split by true class.

    X-axis is ``y_proba[:, cb_label]`` (predicted probability of the
    cyberbullying class - label=0 in our Kaggle-canonical convention).
    Two overlaid histograms colored by true label; vertical dashed line at 0.5
    marks the default decision threshold.
    """
    y_true = np.asarray(y_true)
    p_cb = np.asarray(y_proba)[:, cb_label]

    p_when_cb = p_cb[y_true == cb_label]
    p_when_ncb = p_cb[y_true != cb_label]

    own_fig = ax is None
    if own_fig:
        _, ax = plt.subplots(figsize=(7.5, 4.5))

    ax.hist(p_when_ncb, bins=bins, range=(0, 1), color="#7cc9a3",
            edgecolor="#2a7a4a", alpha=0.85, label="Non-Cyberbullying (true)")
    ax.hist(p_when_cb, bins=bins, range=(0, 1), color="#e07b91",
            edgecolor="#a23a52", alpha=0.7, label="Cyberbullying (true)")
    ax.axvline(0.5, linestyle="--", color="#444", linewidth=1.1)
    ax.set_xlabel("Predicted Probability (Cyberbullying)")
    ax.set_ylabel("Frequency")
    ax.set_title(title or "Prediction Confidence Distribution")
    ax.legend(loc="upper center")

    if save_path is not None and own_fig:
        plt.tight_layout()
        plt.savefig(str(save_path), dpi=150, bbox_inches="tight")


def length_quartile_analysis(
    texts: Sequence[str],
    y_true: Sequence[int],
    y_pred: Sequence[int],
    by: str = "char",
) -> dict:
    """Bin test samples into 4 length quartiles, return per-quartile counts + accuracy.

    Parameters
    ----------
    by : ``"char"`` (default) or ``"word"`` - length metric for quartiling.

    Returns
    -------
    dict with:
      - ``quartiles``: list of 4 labels [``"Q1"``, ``"Q2"``, ``"Q3"``, ``"Q4"``]
      - ``length_bounds``: list of 4 ``(lo, hi)`` tuples per quartile
      - ``n_total`` / ``n_correct`` / ``accuracy``: list of length 4 each
    """
    if by not in ("char", "word"):
        raise ValueError(f"by must be 'char' or 'word', got {by!r}")
    lengths = np.array(
        [len(t) if by == "char" else len(str(t).split()) for t in texts]
    )
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    correct = (y_true == y_pred).astype(int)

    # 4 equal-frequency bins on lengths. ``rank(method='first')`` breaks ties
    # deterministically so reruns produce identical quartile assignments.
    import pandas as pd  # local import - avoid hard dep at module import time
    ranks = pd.Series(lengths).rank(method="first")
    bin_labels = ["Q1", "Q2", "Q3", "Q4"]
    qcat = pd.qcut(ranks, q=4, labels=bin_labels)

    out = {
        "by": by,
        "quartiles": bin_labels,
        "length_bounds": [],
        "n_total": [],
        "n_correct": [],
        "accuracy": [],
    }
    for q in bin_labels:
        mask = (qcat == q).to_numpy()
        n = int(mask.sum())
        nc = int(correct[mask].sum())
        acc = float(nc / n) if n > 0 else 0.0
        lo = int(lengths[mask].min()) if n > 0 else 0
        hi = int(lengths[mask].max()) if n > 0 else 0
        out["length_bounds"].append((lo, hi))
        out["n_total"].append(n)
        out["n_correct"].append(nc)
        out["accuracy"].append(acc)
    return out


def plot_accuracy_by_length(
    quartile_data: dict,
    save_path: Path | None = None,
    ax=None,
    title: str | None = None,
) -> None:
    """Paper Fig 6 replica - bar of accuracy per length quartile."""
    own_fig = ax is None
    if own_fig:
        _, ax = plt.subplots(figsize=(7.5, 4.5))

    bars = ax.bar(
        quartile_data["quartiles"], quartile_data["accuracy"],
        color="#8ec7e8", edgecolor="#2a6fa3", linewidth=0.8,
    )
    for b, v in zip(bars, quartile_data["accuracy"]):
        ax.text(b.get_x() + b.get_width() / 2, v + 0.012, f"{v:.4f}",
                ha="center", fontsize=10, fontweight="bold")
    ax.set_ylim(0, 1.0)
    ax.set_xlabel("Comment Length Quartile")
    ax.set_ylabel("Accuracy")
    ax.set_title(title or "Model Accuracy by Comment Length")

    if save_path is not None and own_fig:
        plt.tight_layout()
        plt.savefig(str(save_path), dpi=150, bbox_inches="tight")


def plot_correct_vs_incorrect_by_length(
    quartile_data: dict,
    save_path: Path | None = None,
    ax=None,
    title: str | None = None,
) -> None:
    """Paper Fig 7 replica - stacked bar of correct vs incorrect per length quartile."""
    own_fig = ax is None
    if own_fig:
        _, ax = plt.subplots(figsize=(7.5, 4.5))

    n_correct = np.array(quartile_data["n_correct"])
    n_total = np.array(quartile_data["n_total"])
    n_incorrect = n_total - n_correct

    x = quartile_data["quartiles"]
    ax.bar(x, n_incorrect, color="#e07b91", edgecolor="#a23a52", label="Incorrect")
    ax.bar(x, n_correct, bottom=n_incorrect, color="#7cc9a3", edgecolor="#2a7a4a",
           label="Correct")
    ax.set_xlabel("Comment Length Quartile")
    ax.set_ylabel("Count")
    ax.set_title(title or "Correct vs Incorrect Predictions by Length")
    ax.legend(loc="upper right")

    if save_path is not None and own_fig:
        plt.tight_layout()
        plt.savefig(str(save_path), dpi=150, bbox_inches="tight")


def plot_paper_repro_suite(
    texts: Sequence[str],
    y_true: Sequence[int],
    y_pred: Sequence[int],
    y_proba: np.ndarray,
    out_dir: Path,
    model_tag: str,
    length_by: str = "char",
    labels: Sequence[str] = LABEL_NAMES,
) -> dict:
    """Generate all 5 Jayanti & Rohman (2026) figure replicas in one call.

    Saves five PNGs into ``out_dir`` with the naming pattern
    ``<model_tag>_<figure>.png``. Returns the ``length_quartile_analysis``
    dict so the caller can log/log-table the per-quartile accuracy.
    """
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    plot_confusion_matrix_dual(
        y_true, y_pred, labels=labels,
        save_path=out_dir / f"{model_tag}_cm_dual.png",
        title_prefix=f"{model_tag} - ",
    )
    plt.close("all")

    plot_per_class_acc_and_errors(
        y_true, y_pred, labels=labels,
        save_path=out_dir / f"{model_tag}_per_class_acc_errors.png",
        title_prefix=f"{model_tag} - ",
    )
    plt.close("all")

    plot_confidence_distribution(
        y_true, y_proba,
        save_path=out_dir / f"{model_tag}_confidence_distribution.png",
        title=f"{model_tag} - Prediction Confidence Distribution",
    )
    plt.close("all")

    quartile_data = length_quartile_analysis(texts, y_true, y_pred, by=length_by)

    plot_accuracy_by_length(
        quartile_data,
        save_path=out_dir / f"{model_tag}_accuracy_by_length.png",
        title=f"{model_tag} - Model Accuracy by Comment Length",
    )
    plt.close("all")

    plot_correct_vs_incorrect_by_length(
        quartile_data,
        save_path=out_dir / f"{model_tag}_correct_vs_incorrect_by_length.png",
        title=f"{model_tag} - Correct vs Incorrect Predictions by Length",
    )
    plt.close("all")

    return quartile_data


def plot_roc_curve(
    y_true: Sequence[int],
    y_proba: np.ndarray,
    save_path: Path | None = None,
    ax=None,
    title: str | None = None,
    label: str | None = None,
) -> None:
    """Plot ROC for the positive class (column 1 of ``y_proba``)."""
    fpr, tpr, _ = roc_curve(y_true, y_proba[:, 1])
    auc = roc_auc_score(y_true, y_proba[:, 1])

    own_fig = ax is None
    if own_fig:
        _, ax = plt.subplots(figsize=(5.5, 4.5))
    curve_label = f"{label} (AUC={auc:.3f})" if label else f"AUC={auc:.3f}"
    ax.plot(fpr, tpr, label=curve_label, linewidth=2)
    ax.plot([0, 1], [0, 1], color="gray", linestyle="--", linewidth=1, label="chance")
    ax.set_xlabel("False positive rate")
    ax.set_ylabel("True positive rate")
    ax.set_title(title or "ROC curve")
    ax.legend(loc="lower right")

    if save_path is not None:
        plt.tight_layout()
        plt.savefig(str(save_path), dpi=150, bbox_inches="tight")