# llmindex.io — Bradley-Terry fit (MM algorithm, ties as half-wins)
# Author:  Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# License: Proprietary — © Simon-Pierre Boucher, all rights reserved

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass
class BTResult:
    log_strength: np.ndarray  # (M,) log-strengths, mean 0
    log_strength_se: np.ndarray  # (M,) SEs from the Fisher information
    iterations: int
    converged: bool


def fit_bradley_terry(
    wins: np.ndarray,
    max_iterations: int = 1000,
    tolerance: float = 1e-10,
    damping: float = 0.1,
) -> BTResult:
    """Fit Bradley-Terry strengths from a wins matrix.

    wins[i, j] = (possibly fractional) number of wins of i over j.
    Ties should be pre-encoded as 0.5 win to each side.
    `damping` adds a tiny uniform prior so isolated models stay finite.
    """
    W = np.asarray(wins, dtype=float)
    if W.ndim != 2 or W.shape[0] != W.shape[1]:
        raise ValueError("wins must be a square matrix")
    M = W.shape[0]
    # Regularization: everyone gets `damping` phantom wins vs everyone else.
    W = W + damping * (np.ones((M, M)) - np.eye(M))
    N = W + W.T  # total comparisons between each pair
    p = np.ones(M)
    converged = False
    it = 0
    for it in range(1, max_iterations + 1):
        denom = (N / (p[:, None] + p[None, :] + 1e-300)).sum(1) - np.diag(
            N / (p[:, None] + p[None, :] + 1e-300)
        )
        new_p = W.sum(1) / np.maximum(denom, 1e-300)
        new_p = new_p / np.exp(np.log(new_p + 1e-300).mean())  # geometric-mean normalize
        if np.max(np.abs(new_p - p)) < tolerance:
            p = new_p
            converged = True
            break
        p = new_p
    log_strength = np.log(p + 1e-300)
    log_strength -= log_strength.mean()
    # Fisher information of log-strengths: I_ii = sum_j N_ij * p_ij * (1 - p_ij)
    # where p_ij = p_i / (p_i + p_j); SE_i = 1 / sqrt(I_ii).
    P = p[:, None] / (p[:, None] + p[None, :] + 1e-300)
    info = (N * P * (1 - P)).sum(1) - np.diag(N * P * (1 - P))
    se = 1.0 / np.sqrt(np.maximum(info, 1e-9))
    return BTResult(log_strength=log_strength, log_strength_se=se, iterations=it, converged=converged)