"""Inference math matching notebook 7's rank-averaging and calibration."""
import numpy as np
from scipy.stats import rankdata


def normalize_scores(values):
    values = np.clip(np.asarray(values, dtype=float), 1e-12, None)
    return values / values.sum(axis=-1, keepdims=True)


def calibrate(values, temperature):
    temperature = float(temperature)
    if not np.isfinite(temperature) or temperature <= 0:
        raise ValueError('Calibration temperature must be positive')
    logits = np.log(normalize_scores(values)) / temperature
    logits -= logits.max(axis=-1, keepdims=True)
    return normalize_scores(np.exp(logits))


def rank_average(probabilities):
    stack = np.asarray(probabilities, dtype=float)
    if stack.ndim != 3 or not np.isfinite(stack).all():
        raise ValueError('Expected finite model-by-phrase-by-class probabilities')
    ranks = np.stack([np.apply_along_axis(rankdata, 1, values) for values in stack])
    return normalize_scores(np.mean(ranks, axis=0))