AUC Report

class AUCReport:
    """Pre-computed linked AUC explorer that renders to a self-contained HTML page.

    Parameters
    ----------
    df : pl.DataFrame
        Data with one column per variable (continuous score) and one per
        outcome (binary 0/1 label), plus an optional ID column.
    variables : list[str]
        Column names of the continuous score variables.
    outcomes : list[str]
        Column names of the binary outcome variables.
    id_col : str or None
        Optional column with sample / patient identifiers.  When provided,
        hovering over points in the ROC and distribution charts reveals the ID.
    kind : {"box", "strip"}
        Distribution chart style: grouped boxplot or jittered strip plot.
    specificity_levels : list[float] or None
        Target specificity levels to annotate on each ROC curve.  Levels
        unreachable by a given curve are silently skipped.
        Defaults to ``[0.8, 0.9]``.
    chart_width, chart_height : int
        Pixel dimensions of each chart panel (ROC and distribution).
        The scatter panel is square with side ``chart_width + 20``.

    Examples
    --------
    One-liner for Quarto / Jupyter::

        from plotutils.auc import AUCReport
        AUCReport(df, variables, outcomes, id_col="patient_id").to_html("report.html")
    """

    _DEFAULT_SPEC_LEVELS: list[float] = [0.8, 0.9, 0.99]

    def __init__(
        self,
        df: pl.DataFrame,
        variables: list[str],
        outcomes: list[str],
        id_col: str | None = None,
        kind: Literal["box", "strip"] = "box",
        specificity_levels: list[float] | None = None,
        chart_width: int = 320,
        chart_height: int = 290,
        auto_reverse: bool = True,
        reference_outcome: str | None = None,
    ) -> None:
        self._df = df
        self._variables = variables
        self._outcomes = outcomes
        self._id_col = id_col
        self._kind = kind
        self._target_spec = (
            specificity_levels
            if specificity_levels is not None
            else self._DEFAULT_SPEC_LEVELS
        )
        self._w = chart_width
        self._h = chart_height
        self._plot_dist = (
            plot_bivariate_boxes if kind == "box" else plot_bivariate_strip
        )

        # Phase 1: reverse anti-correlated outcomes (before any AUC computation).
        self._reversed_outcomes: set[str] = set()
        if auto_reverse and len(outcomes) > 1:
            ref = reference_outcome if reference_outcome is not None else outcomes[0]
            for out in outcomes:
                if out == ref:
                    continue
                # Pearson correlation on non-null overlap with the reference.
                pair = self._df.select([ref, out]).drop_nulls()
                if len(pair) < 2:
                    continue
                corr = pair.select(
                    pl.corr(pl.col(ref).cast(pl.Float64), pl.col(out).cast(pl.Float64))
                ).item()
                if corr is not None and corr < -0.3:
                    self._reversed_outcomes.add(out)

            if self._reversed_outcomes:
                self._df = self._df.with_columns(
                    [pl.col(out) * -1 + 1 for out in self._reversed_outcomes]
                )

        # Phase 2: compute AUC on (possibly outcome-aligned) data.
        self._auc_df = self._build_auc_df()

        # Phase 3: reverse variables where AUC < 0.5 for ALL outcomes.
        self._reversed: set[str] = set()
        if auto_reverse:
            auc_cols = [f"auc_{out}" for out in outcomes]
            for var in variables:
                row = self._auc_df.filter(pl.col("variable") == var)
                if row.is_empty():
                    continue
                aucs = [row[c].item() for c in auc_cols]
                non_null = [a for a in aucs if a is not None]
                if non_null and all(a < 0.5 for a in non_null):
                    self._reversed.add(var)

            if self._reversed:
                self._df = self._df.with_columns(
                    [pl.col(var) * -1 for var in self._reversed]
                )
                self._auc_df = self._build_auc_df()

        # Build display names: add (+)/(-) prefixes when any variable is reversed.
        if self._reversed:
            self._display_names = {
                var: f"(-) {var}" if var in self._reversed else f"(+) {var}"
                for var in variables
            }
        else:
            self._display_names = {var: var for var in variables}

        # Build outcome display names: add (+)/(-) prefixes when any outcome is reversed.
        if self._reversed_outcomes:
            self._outcome_display_names = {
                out: f"(-) {out}" if out in self._reversed_outcomes else f"(+) {out}"
                for out in outcomes
            }
        else:
            self._outcome_display_names = {out: out for out in outcomes}

        # Add display_name column to _auc_df.
        self._auc_df = self._auc_df.with_columns(
            pl.col("variable").replace_strict(self._display_names).alias("display_name")
        )

        # Pre-compute all chart specs (expensive, done once).
        self._roc_specs = self._build_roc_specs()
        self._dist_specs = self._build_dist_specs()

    # ── Internal builders ──────────────────────────────────────────────────

    def _build_auc_df(self) -> pl.DataFrame:
        n_total = len(self._df)
        # Outcome-level missingness is the same for all variables.
        outcome_missing: dict[str, int] = {
            out: int(self._df[out].is_null().sum()) for out in self._outcomes
        }
        rows = []
        for var in self._variables:
            n_score_missing = int(self._df[var].is_null().sum())
            row: dict = {
                "variable": var,
                "n_total": n_total,
                "n_score_missing": n_score_missing,
            }
            for out in self._outcomes:
                # Drop rows where score OR outcome is null before ROC computation.
                df_vo = (
                    self._df.select([var, out])
                    .rename({var: "score", out: "label"})
                    .drop_nulls()
                )
                n_valid = len(df_vo)
                row[f"n_valid_{out}"] = n_valid
                row[f"n_outcome_missing_{out}"] = outcome_missing[out]

                try:
                    roc = _compute_roc(df_vo, "score", "label")
                except InsufficientClassesError:
                    row[f"auc_{out}"] = None
                    for sl in self._target_spec:
                        row[f"pauc_{out}_q{int(sl * 100)}"] = None
                    continue
                row[f"auc_{out}"] = round(_compute_auc(roc), 4)
                # pAUC columns: one per specificity level, None if unreachable.
                max_spec = float(
                    roc.filter(pl.col("threshold").is_not_null())["specificity"].max()  # type: ignore[arg-type]
                )
                for sl in self._target_spec:
                    col = f"pauc_{out}_q{int(sl * 100)}"
                    if sl <= max_spec:
                        row[col] = round(
                            _compute_pauc(
                                roc, sl, 1.0, focus="specificity", mcclish=True
                            ),
                            4,
                        )
                    else:
                        row[col] = None
            rows.append(row)
        return pl.DataFrame(rows)

    def _build_roc_specs(self) -> dict[str, dict]:
        specs: dict[str, dict] = {}
        for var in self._variables:
            for out in self._outcomes:
                select_cols = [self._id_col, var, out] if self._id_col else [var, out]
                df_vo = (
                    self._df.select(select_cols)
                    .rename({var: "score", out: "label"})
                    .drop_nulls(subset=["score", "label"])
                )
                try:
                    roc_check = _compute_roc(
                        df_vo.select(["score", "label"]), "score", "label"
                    )
                except InsufficientClassesError:
                    continue
                max_spec = float(
                    roc_check.filter(pl.col("threshold").is_not_null())[
                        "specificity"
                    ].max()  # type: ignore[arg-type]
                )
                achievable = [s for s in self._target_spec if s <= max_spec] or None
                chart = plot_roc_curve(
                    df_vo,
                    score_col="score",
                    label_col="label",
                    id_col=self._id_col,
                    specificity_levels=achievable,
                    title=self._display_names[var],
                    width=self._w,
                    height=self._h,
                )
                specs[f"{var}|{out}"] = json.loads(chart.to_json())
        return specs

    def _build_dist_specs(self) -> dict[str, dict]:
        specs: dict[str, dict] = {}
        for var in self._variables:
            for out_x in self._outcomes:
                for out_y in self._outcomes:
                    if out_x == out_y:
                        continue

                    select_cols = [var, out_x, out_y]
                    if self._id_col:
                        select_cols.insert(0, self._id_col)

                    rename_map = {var: "score", out_x: "label_x", out_y: "label_y"}
                    df_full = (
                        self._df.select(select_cols)
                        .rename(rename_map)
                        .with_columns(
                            pl.col("label_x").cast(pl.Utf8).fill_null("Missing"),
                            pl.col("label_y").cast(pl.Utf8).fill_null("Missing"),
                        )
                    )

                    # Patients with null score → separate layer.
                    df_valid = df_full.filter(pl.col("score").is_not_null())
                    df_missing = df_full.filter(pl.col("score").is_null())

                    d: dict = {
                        "score": df_valid["score"].to_list(),
                        "label_x": df_valid["label_x"].to_list(),
                        "label_y": df_valid["label_y"].to_list(),
                    }
                    if self._id_col:
                        d[self._id_col] = df_valid[self._id_col].to_list()

                    missing_score_df: pl.DataFrame | None = None
                    if len(df_missing) > 0:
                        miss_d: dict = {
                            "label_x": df_missing["label_x"].to_list(),
                            "label_y": df_missing["label_y"].to_list(),
                        }
                        if self._id_col:
                            miss_d[self._id_col] = df_missing[self._id_col].to_list()
                        missing_score_df = pl.DataFrame(miss_d)

                    chart = self._plot_dist(
                        pl.DataFrame(d),
                        score_col="score",
                        label_x_col="label_x",
                        label_y_col="label_y",
                        id_col=self._id_col,
                        title=self._display_names[var],
                        x_title=self._outcome_display_names[out_x].replace("_", " "),
                        color_title=self._outcome_display_names[out_y].replace("_", " "),
                        y_title="score",
                        width=self._w,
                        height=self._h,
                        missing_score_df=missing_score_df,
                    )
                    specs[f"{var}|{out_x}|{out_y}"] = json.loads(chart.to_json())
        return specs

    @staticmethod
    def _auc_lookup(signal: str, outcomes: list[str], field: str = "auc_{out}") -> str:
        """Nested Vega ternary to pick the right datum field by signal value.

        Parameters
        ----------
        signal : str
            Name of the Vega signal (e.g. ``"outcome_x"``).
        outcomes : list[str]
            All possible outcome values.
        field : str
            Field name template with ``{out}`` placeholder (e.g.
            ``"auc_{out}"`` or ``"n_valid_{out}"``).
        """

        def _f(out: str) -> str:
            return field.format(out=out)

        expr = f"datum.{_f(outcomes[-1])}"
        for out in reversed(outcomes[:-1]):
            expr = f"{signal} == '{out}' ? datum.{_f(out)} : {expr}"
        return expr

    def _build_scatter(self) -> alt.Chart:
        outcomes = self._outcomes
        outcome_labels = [self._outcome_display_names[o].replace("_", " ") for o in outcomes]
        scatter_size = self._w + 20

        outcome_x = alt.param(
            name="outcome_x",
            value=outcomes[0],
            bind=alt.binding_select(
                options=outcomes, labels=outcome_labels, name="X-axis: "
            ),
        )
        outcome_y = alt.param(
            name="outcome_y",
            value=outcomes[1],
            bind=alt.binding_select(
                options=outcomes, labels=outcome_labels, name="Y-axis: "
            ),
        )
        click_sel = alt.selection_point(
            fields=["variable"], name="var_sel", on="click", empty=False
        )

        return (
            alt.Chart(self._auc_df)
            .transform_calculate(
                x_auc=self._auc_lookup("outcome_x", outcomes),
                y_auc=self._auc_lookup("outcome_y", outcomes),
                n_valid_x=self._auc_lookup("outcome_x", outcomes, "n_valid_{out}"),
                n_valid_y=self._auc_lookup("outcome_y", outcomes, "n_valid_{out}"),
                n_omiss_x=self._auc_lookup(
                    "outcome_x", outcomes, "n_outcome_missing_{out}"
                ),
                n_omiss_y=self._auc_lookup(
                    "outcome_y", outcomes, "n_outcome_missing_{out}"
                ),
            )
            .transform_calculate(
                _sz_x=(
                    '"X: " + datum.n_valid_x + "/" + datum.n_total'
                    ' + " (" + datum.n_omiss_x + " outcomes missing)"'
                ),
                _sz_y=(
                    '"Y: " + datum.n_valid_y + "/" + datum.n_total'
                    ' + " (" + datum.n_omiss_y + " outcomes missing)"'
                ),
                _sz_score='"Score missing: " + datum.n_score_missing',
            )
            .mark_point(filled=True)
            .encode(
                x=alt.X(
                    "x_auc:Q",
                    title="AUC (X outcome)",
                    scale=alt.Scale(domain=[0.4, 1.05]),
                ),
                y=alt.Y(
                    "y_auc:Q",
                    title="AUC (Y outcome)",
                    scale=alt.Scale(domain=[0.4, 1.05]),
                ),
                color=alt.condition(
                    click_sel,
                    alt.Color("display_name:N", legend=None),
                    alt.value("lightgray"),
                ),
                size=alt.condition(click_sel, alt.value(180), alt.value(80)),
                tooltip=[
                    alt.Tooltip("display_name:N", title="Variable"),
                    alt.Tooltip("x_auc:Q", format=".3f", title="AUC (X)"),
                    alt.Tooltip("y_auc:Q", format=".3f", title="AUC (Y)"),
                    alt.Tooltip("_sz_x:N", title="Sample size"),
                    alt.Tooltip("_sz_y:N", title=""),
                    alt.Tooltip("_sz_score:N", title=""),
                ],
            )
            .add_params(click_sel, outcome_x, outcome_y)
            .properties(
                title=alt.Title(
                    "AUC – AUC scatter",
                    subtitle="Click a variable to reveal its charts →",
                ),
                width=scatter_size,
                height=scatter_size,
            )
            .configure_axis(
                gridColor="#444", gridWidth=0.75, gridDash=[3, 3], gridOpacity=0.4
            )
            .configure_view(strokeWidth=0)
        )

    # ── Public API ─────────────────────────────────────────────────────────────

    def to_html(self, path: str | None = None) -> str:
        """Render the linked AUC explorer as a self-contained HTML page.

        Parameters
        ----------
        path : str or None
            If provided, write the HTML to this file path in addition to
            returning it.

        Returns
        -------
        str
            Full HTML page as a string (suitable for ``IPython.display.HTML``
            or Quarto ``{python}`` cells with ``output: asis``).
        """
        dist_kind_label = "Boxplot" if self._kind == "box" else "Strip plot"
        roc_specs_js = json.dumps(self._roc_specs)
        dist_specs_js = json.dumps(self._dist_specs)
        scatter_js = json.dumps(json.loads(self._build_scatter().to_json()))
        out_display_js = json.dumps(self._outcome_display_names)

        html = f"""<!DOCTYPE html>
<html>
<head>
  <meta charset="UTF-8">
  <title>Linked AUC Explorer</title>
  <script src="https://cdn.jsdelivr.net/npm/vega@6"></script>
  <script src="https://cdn.jsdelivr.net/npm/vega-lite@6.1.0"></script>
  <script src="https://cdn.jsdelivr.net/npm/vega-embed@7"></script>
  <style>
    body {{
      font-family: sans-serif;
      margin: 0;
      padding: 20px 28px;
      background: #fafafa;
      color: #222;
    }}
    h1 {{ font-size: 18px; margin: 0 0 4px; }}
    p.hint {{ color: #777; font-size: 13px; margin: 0 0 18px; }}

    #grid {{
      display: grid;
      grid-template-columns: 1fr 1fr;
      grid-template-rows: auto auto;
      gap: 20px;
      align-items: start;
    }}

    .cell-label {{
      font-size: 12px;
      font-weight: bold;
      color: #666;
      margin-bottom: 4px;
      min-height: 16px;
    }}

    .slot {{
      border: 2px dashed #ddd;
      border-radius: 6px;
      display: flex;
      align-items: center;
      justify-content: center;
      color: #bbb;
      font-size: 13px;
      box-sizing: border-box;
      min-height: 320px;
    }}

    /* Prevent Vega canvases from overflowing their grid cell */
    .vega-embed {{
      width: 100%;
    }}
    .vega-embed canvas, .vega-embed svg {{
      max-width: 100%;
      height: auto;
    }}
  </style>
</head>
<body>
  <h1>Linked AUC Explorer</h1>
  <p class="hint">
    Select outcomes for the scatter axes. Click a variable point to populate the other three charts.
  </p>

  <div id="grid">

    <!-- top-left: scatter -->
    <div>
      <div class="cell-label">AUC scatter</div>
      <div id="scatter"></div>
    </div>

    <!-- top-right: ROC for outcome X -->
    <div>
      <div class="cell-label" id="lbl-roc0">ROC — outcome X</div>
      <div class="slot" id="roc0"><span>← click a variable</span></div>
    </div>

    <!-- bottom-left: distribution chart -->
    <div>
      <div class="cell-label" id="lbl-dist">{dist_kind_label} — score distribution</div>
      <div class="slot" id="dist"><span>← click a variable</span></div>
    </div>

    <!-- bottom-right: ROC for outcome Y -->
    <div>
      <div class="cell-label" id="lbl-roc1">ROC — outcome Y</div>
      <div class="slot" id="roc1"><span>← click a variable</span></div>
    </div>

  </div>

<script>
const ROC  = {roc_specs_js};
const DIST = {dist_specs_js};
const SCATTER = {scatter_js};
const OUT_DISPLAY = {out_display_js};

let activeVar = null;
let activeDisplayName = null;

function embed(id, spec) {{
  vegaEmbed('#' + id, spec, {{actions: false}});
}}

function update(view) {{
  if (!activeVar) return;
  const ox = view.signal('outcome_x');
  const oy = view.signal('outcome_y');

  // ROC charts
  const oxLabel = (OUT_DISPLAY[ox] || ox).replace(/_/g, ' ');
  const oyLabel = (OUT_DISPLAY[oy] || oy).replace(/_/g, ' ');
  document.getElementById('lbl-roc0').textContent = 'ROC — ' + oxLabel + ' (' + activeDisplayName + ')';
  document.getElementById('lbl-roc1').textContent = 'ROC — ' + oyLabel + ' (' + activeDisplayName + ')';
  const roc0 = ROC[activeVar + '|' + ox];
  const roc1 = ROC[activeVar + '|' + oy];
  if (roc0) embed('roc0', roc0);
  if (roc1) embed('roc1', roc1);

  // Distribution chart
  document.getElementById('lbl-dist').textContent =
    '{dist_kind_label} — ' + activeDisplayName + ' score × ' + oxLabel + ' × ' + oyLabel;
  const dist = DIST[activeVar + '|' + ox + '|' + oy];
  if (dist) embed('dist', dist);
}}

vegaEmbed('#scatter', SCATTER, {{actions: false}}).then(function(r) {{
  const view = r.view;

  view.addEventListener('click', function(event, item) {{
    if (item && item.datum && item.datum.variable) {{
      activeVar = item.datum.variable;
      activeDisplayName = item.datum.display_name || activeVar;
      update(view);
    }}
  }});

  view.addSignalListener('outcome_x', function() {{ update(view); }});
  view.addSignalListener('outcome_y', function() {{ update(view); }});
}});
</script>
</body>
</html>
"""
        if path is not None:
            with open(path, "w", encoding="utf-8") as f:
                f.write(html)
        return html

    def _repr_html_(self) -> str:
        """Rich HTML display for Jupyter / Quarto notebooks."""
        return self.to_html()