"""
Correct aggregate group shares for classifier error.

If you classify a corpus and report the share of each group, those shares are
biased by classifier error, and the bias is largest exactly where accuracy is
lowest. With a confusion matrix M where

    M[i, j] = P(predicted = i | true = j)

observed and true shares are related by  p_obs = M @ p_true, so

    p_true ~= pinv(M) @ p_obs

This is the standard quantification / prior-correction estimator. Its bias goes
to zero as M approaches the identity, and it is base-rate invariant, so a
confusion matrix estimated on a stratified validation sample is valid here.

Two cautions:

  * Estimate M with the SAME pipeline that produced the counts you are
    correcting - same detector, same crop margin, same checkpoints. A confusion
    matrix from a different preprocessing is the wrong matrix.
  * When a class has very low recall the correction is poorly conditioned and
    can return a near-zero or clipped estimate. Report those classes as bounds.
    On our data the unweighted correction sent Hispanic to 0.0, while the
    inverse-probability-weighted version recovered 2.6%. Prefer the weighted
    matrix and say which you used.

Licence: MIT.
"""

from __future__ import annotations

import numpy as np

__all__ = [
    "confusion_matrix_from_labels",
    "correct_proportions",
    "correct_counts",
]


def confusion_matrix_from_labels(y_true, y_pred, labels, weights=None) -> np.ndarray:
    """
    Column-stochastic confusion matrix M[i, j] = P(pred = labels[i] | true = labels[j]).

    weights : optional per-sample weights. Pass inverse selection probabilities
              when the validation sample was stratified on predicted labels,
              otherwise M is conditioned on the predictions rather than the
              population.
    """
    labels = list(labels)
    idx = {l: i for i, l in enumerate(labels)}
    n = len(labels)
    if weights is None:
        weights = np.ones(len(y_true), dtype=float)
    weights = np.asarray(weights, dtype=float)

    M = np.zeros((n, n), dtype=float)
    for t, p, w in zip(y_true, y_pred, weights):
        if t in idx and p in idx:
            M[idx[p], idx[t]] += w

    for j in range(n):
        col = M[:, j].sum()
        if col > 0:
            M[:, j] /= col
        else:
            M[j, j] = 1.0          # unobserved class: assume no confusion
    return M


def correct_proportions(p_obs, M, clip_negative: bool = True) -> np.ndarray:
    """
    Recover true class proportions from observed ones.

    p_obs : observed proportions, same order as the confusion matrix labels
    M     : column-stochastic confusion matrix
    """
    p_obs = np.asarray(p_obs, dtype=float)
    q = np.linalg.pinv(M) @ p_obs
    if clip_negative:
        q = np.clip(q, 0.0, None)
    total = q.sum()
    return q / total if total > 0 else p_obs


def correct_counts(counts, M) -> np.ndarray:
    """Same correction expressed in counts rather than proportions."""
    counts = np.asarray(counts, dtype=float)
    total = counts.sum()
    if total <= 0:
        return counts
    return correct_proportions(counts / total, M) * total


if __name__ == "__main__":
    # Worked example: race shares of real advertisements, using the published
    # real-domain IPW matrix from benchmark/confusion_race_real_ipw.csv
    labels = ["white", "black", "asian", "hispanic"]
    p_obs  = [0.6494, 0.1889, 0.1552, 0.0065]
    census = [0.6090, 0.1329, 0.0697, 0.1884]

    M = np.array([
        [0.9603, 0.1138, 0.2073, 0.1376],
        [0.0057, 0.8383, 0.0106, 0.2646],
        [0.0327, 0.0448, 0.7804, 0.4092],
        [0.0013, 0.0031, 0.0017, 0.1886],
    ])

    p_cor = correct_proportions(p_obs, M)
    print(f"{'class':10s} {'observed':>9s} {'corrected':>10s} {'census':>8s} "
          f"{'dev obs':>8s} {'dev cor':>8s}")
    for l, o, c, r in zip(labels, p_obs, p_cor, census):
        print(f"{l:10s} {o*100:8.2f}% {c*100:9.2f}% {r*100:7.2f}% "
              f"{(o-r)*100:+7.2f} {(c-r)*100:+7.2f}")
    print("\nWhite over-representation: +4.04 pp observed -> +0.67 pp corrected.")