//! Active learning: which comparison teaches you the most.
//!
//! [`quality`](crate::quality) answers "is this matchup *fair*". That is a
//! different question from "is this matchup *informative*", and the two
//! coincide only for two evenly matched competitors. When each observation
//! costs something — a human click, a scheduled fixture — the question worth
//! asking is the second one.
//!
//! The quantity here is expected information gain: the outcome-weighted
//! divergence between what you believe now and what you would believe after
//! seeing the result.
//!
//! ```text
//! EIG(matchup) = SUM  P(outcome) * KL( posterior_after(outcome) || prior )
//!              outcome
//! ```
//!
//! It is the mutual information between the observed outcome and the skills,
//! which is worth remembering because it pins the scale: information gain
//! cannot exceed the entropy of the thing you are about to observe. A contest
//! with `k` distinguishable outcomes can teach you at most `ln k` nats,
//! whatever the ratings. That ceiling is the sharpest available test of an
//! implementation — see [`expected_information_gain`].

use crate::{
    GameOptions, Gaussian, InferenceError, Outcome, Rating, drift::Drift, predict, time::Time,
};

/// Outcomes below this probability contribute nothing measurable and are not
/// worth an inference pass.
///
/// The contribution of an outcome is `P * KL`, and `KL` is bounded in practice
/// by tens of nats, so a probability this small moves the total by less than
/// the quadrature error already present in `P` itself.
const NEGLIGIBLE: f64 = 1e-12;

/// `KL(q || p)` for two univariate Gaussians, in nats.
///
/// Both arguments are proper posteriors from inference, so the degenerate
/// cases guarded here (zero or infinite variance) indicate that inference has
/// broken down rather than anything a caller did.
fn kl_divergence(q: Gaussian, p: Gaussian) -> f64 {
    let (var_q, var_p) = (q.sigma().powi(2), p.sigma().powi(2));

    if !(var_q.is_finite() && var_p.is_finite()) || var_q <= 0.0 || var_p <= 0.0 {
        return 0.0;
    }

    let mean_gap = q.mu() - p.mu();

    // Algebraically `0.5 * (ln(var_p/var_q) + (var_q + gap^2)/var_p - 1)`, but
    // written so that neither term can go negative.
    //
    // The direct form cancels against its `- 1.0` for two near-identical
    // distributions and returns a *negative* divergence — measured, 762 082 of
    // 3 000 000 near-identical pairs, worst `-5.55e-17`, which is exactly one
    // ULP of the 1.0. It also loses the answer entirely where it is small:
    // at `var_q/var_p - 1 = 1e-9` the direct form gives `0.0` where the true
    // value is `2.5e-19`.
    //
    // With `u = var_q/var_p - 1` the variance part is `0.5 * (u - ln(1+u))`,
    // which is non-negative for every `u > -1`, and the mean part is a square
    // over a positive variance. Non-negativity is then structural rather than
    // incidental.
    let u = var_q / var_p - 1.0;
    0.5 * u_minus_ln1p(u) + mean_gap * mean_gap / (2.0 * var_p)
}

/// `u - ln(1 + u)`, without the cancellation that spelling invites.
///
/// Both terms are approximately `u` for small `u`, so the subtraction loses
/// everything just where the result matters. The Taylor series
/// `u^2/2 - u^3/3 + u^4/4 - ...` is exact in that regime and manifestly
/// non-negative, since `u^2/2` dominates.
fn u_minus_ln1p(u: f64) -> f64 {
    if u.abs() < 1e-4 {
        let u2 = u * u;
        u2 * (0.5 - u / 3.0 + u2 / 4.0)
    } else {
        u - libm::log1p(u)
    }
}

/// Expected information gain of a hypothetical matchup, in nats.
///
/// Enumerates the outcomes this matchup could have, runs inference for each to
/// get the belief it would produce, and weights the resulting divergence by
/// that outcome's probability. A higher value means the result would teach you
/// more.
///
/// # Interpreting the value
///
/// Nats. The upper bound is the entropy of the outcome variable: at most
/// `ln 2 ≈ 0.693` for a two-way result, `ln 3 ≈ 1.099` once draws are
/// possible, `ln k` for `k` outcomes. A value near the ceiling means the
/// result is close to a coin flip *and* would move the posteriors a long way;
/// a value near zero means you already know what will happen, or that the
/// result would barely change your beliefs if you saw it.
///
/// This is not a monotone transform of [`quality`](crate::quality). A lopsided
/// matchup between two uncertain competitors scores well on quality-times-
/// variance heuristics and poorly here, because the near-certain outcome
/// carries almost no information.
///
/// # Cost
///
/// One full inference pass per possible outcome, so this is far more expensive
/// than `quality()` — which is one closed-form evaluation. The outcome count
/// grows quickly with team count (3 outcomes for two teams that can draw, 13
/// for three, 75 for four), and scoring every candidate pairing among `n`
/// competitors is `O(n² × outcomes)` inference passes.
///
/// For a selector over many candidates, shortlist with the cheap
/// [`quality`](crate::quality) or
/// [`predict_win_probabilities`](crate::History::predict_win_probabilities)
/// first and score only the shortlist here. The expected-variance-reduction
/// proxy sometimes suggested as a cheaper alternative is *not* cheaper: it
/// needs the same hypothetical posteriors, so it shares the dominant cost.
///
/// # Errors
///
/// - `NotEnoughTeams` if fewer than two teams are supplied.
/// - `EmptyTeam` if any team has no members.
/// - `TooManyTeams` if the outcome space is too large to enumerate; see
///   [`MAX_PREDICTED_TEAMS`](crate::MAX_PREDICTED_TEAMS).
/// - `InvalidParameter` for a `p_draw` outside `[0.0, 1.0)`.
/// - `GridTooCoarse` when the performance sigmas are too far apart to
///   integrate on one grid. This comes from `outcome_distribution`, which runs
///   before any inference — so it is not covered by "anything `Game::ranked`
///   returns" below.
/// - Anything [`Game::ranked`](crate::Game::ranked) returns for a hypothetical
///   outcome.
pub fn expected_information_gain<T: Time, D: Drift<T>>(
    teams: &[&[Rating<T, D>]],
    options: &GameOptions,
) -> Result<f64, InferenceError> {
    if teams.len() < 2 {
        return Err(InferenceError::NotEnoughTeams { got: teams.len() });
    }
    if teams.len() > crate::MAX_PREDICTED_TEAMS {
        return Err(InferenceError::TooManyTeams {
            got: teams.len(),
            max: crate::MAX_PREDICTED_TEAMS,
        });
    }
    if !(0.0..1.0).contains(&options.p_draw) {
        return Err(InferenceError::InvalidParameter {
            parameter: crate::Parameter::PDraw,
            value: options.p_draw,
        });
    }
    for (idx, team) in teams.iter().enumerate() {
        if team.is_empty() {
            return Err(InferenceError::EmptyTeam { team: idx });
        }
    }

    // Prediction runs on performances: skill inflated by each member's beta.
    let performances: Vec<Gaussian> = teams
        .iter()
        .map(|team| {
            team.iter()
                .fold(crate::N00, |acc, rating| acc.convolve(rating.performance()))
        })
        .collect();

    // Draw margins per pair, derived from the teams' betas exactly as
    // inference derives them, so the outcomes weighted here are the outcomes
    // that would actually be fitted.
    let beta_sq: Vec<f64> = teams
        .iter()
        .map(|team| team.iter().map(|r| r.beta().powi(2)).sum())
        .collect();
    let p_draw = options.p_draw;
    let margins = predict::Margins::new(teams.len(), |i, j| {
        if p_draw == 0.0 {
            0.0
        } else {
            crate::compute_margin(p_draw, (beta_sq[i] + beta_sq[j]).sqrt())
        }
    });

    let mut gain = 0.0;

    for (ranks, probability) in predict::outcome_distribution(&performances, &margins)? {
        if probability <= NEGLIGIBLE {
            continue;
        }

        let game = crate::Game::ranked(teams, Outcome::ranking(ranks), options)?;
        let posteriors = game.posteriors();

        // Beliefs factorise across competitors, so the joint divergence is the
        // sum of the per-competitor ones.
        let divergence: f64 = teams
            .iter()
            .zip(&posteriors)
            .flat_map(|(team, posterior)| team.iter().zip(posterior))
            .map(|(rating, &after)| kl_divergence(after, rating.prior()))
            .sum();

        gain += probability * divergence;
    }

    Ok(gain)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{BETA, ConstantDrift, GAMMA};

    type R = Rating<i64, ConstantDrift>;

    fn rating(mu: f64, sigma: f64) -> R {
        R::new(
            Gaussian::from_ms(mu, sigma),
            BETA,
            ConstantDrift::new(GAMMA),
        )
    }

    fn options(p_draw: f64) -> GameOptions {
        GameOptions {
            p_draw,
            ..GameOptions::default()
        }
    }

    fn eig(teams: &[&[R]], p_draw: f64) -> f64 {
        expected_information_gain(teams, &options(p_draw)).unwrap()
    }

    /// The analytic ceiling. Information gain is the mutual information between
    /// the outcome and the skills, so it cannot exceed the entropy of the
    /// outcome variable — whatever the ratings. This is the check a subtly
    /// wrong implementation fails while still returning plausible numbers: an
    /// early prototype of this returned 4.77 nats from a sign error and passed
    /// every monotonicity test.
    #[test]
    fn never_exceeds_the_entropy_of_the_outcome() {
        let ceiling_two = std::f64::consts::LN_2;

        for (a, b) in [
            (rating(0.0, 6.0), rating(0.0, 6.0)),
            (rating(0.0, 0.5), rating(0.0, 0.5)),
            (rating(12.0, 6.0), rating(-12.0, 6.0)),
            (rating(40.0, 1.0), rating(-40.0, 1.0)),
            (rating(3.0, 6.0), rating(-2.0, 0.1)),
            (rating(0.0, 25.0), rating(0.0, 25.0)),
        ] {
            let g = eig(&[&[a], &[b]], 0.0);
            assert!(
                g >= 0.0 && g <= ceiling_two,
                "EIG {g} outside [0, ln 2] for mu=({}, {}) sigma=({}, {})",
                a.prior().mu(),
                b.prior().mu(),
                a.prior().sigma(),
                b.prior().sigma()
            );
        }
    }

    /// With draws enabled there are three outcomes, so the ceiling rises to
    /// `ln 3` — and the two-outcome bound no longer applies.
    #[test]
    fn the_ceiling_follows_the_outcome_count() {
        let ceiling_three = 3.0f64.ln();
        for sigma in [0.5, 3.0, 6.0, 25.0] {
            let g = eig(&[&[rating(0.0, sigma)], &[rating(0.0, sigma)]], 0.25);
            assert!(
                g >= 0.0 && g <= ceiling_three,
                "EIG {g} outside [0, ln 3] at sigma {sigma}"
            );
        }
    }

    /// An even matchup between uncertain competitors is the informative one.
    /// A hopelessly lopsided matchup teaches you almost nothing, because you
    /// already know how it ends.
    #[test]
    fn an_even_matchup_beats_a_lopsided_one() {
        let even = eig(&[&[rating(0.0, 6.0)], &[rating(0.0, 6.0)]], 0.0);
        let lopsided = eig(&[&[rating(12.0, 6.0)], &[rating(-12.0, 6.0)]], 0.0);
        assert!(
            even > lopsided,
            "even {even} should beat lopsided {lopsided}"
        );
    }

    /// Certainty is the thing information gain is measuring the absence of:
    /// the less you know, the more there is to learn.
    #[test]
    fn gain_falls_as_certainty_rises() {
        let mut previous = f64::INFINITY;
        for sigma in [12.0, 6.0, 3.0, 1.0, 0.5, 0.1] {
            let g = eig(&[&[rating(0.0, sigma)], &[rating(0.0, sigma)]], 0.0);
            assert!(
                g < previous,
                "sigma {sigma}: {g} did not fall below {previous}"
            );
            previous = g;
        }
        assert!(previous >= 0.0);
    }

    /// The heuristic this replaces is `quality * sigma_a^2 * sigma_b^2`. It is
    /// not a monotone transform of information gain — it ranks a lopsided
    /// matchup above a confident even one, and EIG ranks them the other way.
    /// Pinning the disagreement down is what stops a future "simplification"
    /// from quietly reverting to the heuristic.
    #[test]
    fn disagrees_with_the_quality_times_variance_heuristic() {
        let heuristic = |a: &R, b: &R| {
            crate::quality(&[&[a.prior()], &[b.prior()]], BETA)
                * a.prior().sigma().powi(2)
                * b.prior().sigma().powi(2)
        };

        let (confident_a, confident_b) = (rating(0.0, 0.5), rating(0.0, 0.5));
        let (lopsided_a, lopsided_b) = (rating(12.0, 6.0), rating(-12.0, 6.0));

        assert!(
            heuristic(&lopsided_a, &lopsided_b) > heuristic(&confident_a, &confident_b),
            "the heuristic should prefer the lopsided matchup"
        );
        assert!(
            eig(&[&[confident_a], &[confident_b]], 0.0) > eig(&[&[lopsided_a], &[lopsided_b]], 0.0),
            "information gain should prefer the even matchup"
        );
    }

    #[test]
    fn supports_more_than_two_teams() {
        let teams: Vec<Vec<R>> = vec![
            vec![rating(0.0, 6.0)],
            vec![rating(0.0, 6.0)],
            vec![rating(0.0, 6.0)],
        ];
        let refs: Vec<&[R]> = teams.iter().map(Vec::as_slice).collect();
        let g = expected_information_gain(&refs, &options(0.0)).unwrap();
        // Six distinguishable orderings with no draws.
        assert!(
            g > 0.0 && g <= 6.0f64.ln(),
            "three-team EIG {g} out of range"
        );
    }

    #[test]
    fn multi_member_teams_are_supported() {
        let a = [rating(0.0, 6.0), rating(1.0, 4.0)];
        let b = [rating(0.0, 6.0)];
        let g = expected_information_gain(&[&a, &b], &options(0.0)).unwrap();
        assert!(g > 0.0 && g <= std::f64::consts::LN_2, "{g}");
    }

    #[test]
    fn degenerate_shapes_are_errors() {
        let a = [rating(0.0, 6.0)];
        assert!(matches!(
            expected_information_gain(&[&a], &options(0.0)),
            Err(InferenceError::NotEnoughTeams { got: 1 })
        ));
        let empty: [R; 0] = [];
        assert!(matches!(
            expected_information_gain(&[&a, &empty], &options(0.0)),
            Err(InferenceError::EmptyTeam { team: 1 })
        ));
        assert!(matches!(
            expected_information_gain(&[&a, &a], &options(1.5)),
            Err(InferenceError::InvalidParameter {
                parameter: crate::Parameter::PDraw,
                ..
            })
        ));
    }

    #[test]
    fn kl_divergence_is_zero_for_identical_beliefs() {
        let g = Gaussian::from_ms(3.0, 2.0);
        assert!(kl_divergence(g, g).abs() < 1e-15);
    }

    #[test]
    fn kl_divergence_is_non_negative_and_grows_with_separation() {
        let prior = Gaussian::from_ms(0.0, 3.0);
        let mut previous = 0.0;
        for mu in [0.0, 0.5, 1.0, 2.0, 4.0] {
            let d = kl_divergence(Gaussian::from_ms(mu, 3.0), prior);
            assert!(d >= 0.0, "negative divergence at mu {mu}: {d}");
            assert!(d >= previous, "not increasing at mu {mu}");
            previous = d;
        }
    }
}