//! The joint must span slices, because Through Time reads each competitor at
//! their own last appearance.
//!
//! The exact posterior of a multi-slice scored history is still Gaussian: the
//! prior, the drift between appearances, and the scored likelihoods are all
//! Gaussian. So it can be written out by hand and compared against, which is
//! the check a single-slice fixture cannot make.

use smallvec::smallvec;
use trueskill_tt::{
    ConstantDrift, ConvergenceOptions, Event, History, Member, Outcome, Team, UnknownKeys,
};

const SIGMA0: f64 = 6.0;
const BETA: f64 = 1.0;
const SCORE_SIGMA: f64 = 2.0;
const GAMMA: f64 = 0.5;

type H = History;

fn history(gamma: f64) -> H {
    History::builder()
        .mu(0.0)
        .sigma(SIGMA0)
        .beta(BETA)
        .score_sigma(SCORE_SIGMA)
        .drift(ConstantDrift::new(gamma))
        .unknown_keys(UnknownKeys::Reject)
        .convergence(ConvergenceOptions {
            max_iter: 20_000,
            epsilon: 1e-13,
            alpha: 1.0,
        })
        .build()
}

fn duel(a: &'static str, b: &'static str, t: i64, sa: f64, sb: f64) -> Event<i64, &'static str> {
    Event {
        time: t,
        teams: smallvec![
            Team::with_members([Member::new(a)]),
            Team::with_members([Member::new(b)]),
        ],
        outcome: Outcome::scores([sa, sb]),
    }
}

fn inverse(mut a: Vec<Vec<f64>>) -> Vec<Vec<f64>> {
    let n = a.len();
    let mut inv: Vec<Vec<f64>> = (0..n)
        .map(|i| (0..n).map(|j| f64::from(u8::from(i == j))).collect())
        .collect();
    for col in 0..n {
        let mut piv = col;
        for r in col + 1..n {
            if a[r][col].abs() > a[piv][col].abs() {
                piv = r;
            }
        }
        a.swap(col, piv);
        inv.swap(col, piv);
        let d = a[col][col];
        for j in 0..n {
            a[col][j] /= d;
            inv[col][j] /= d;
        }
        for r in 0..n {
            if r == col {
                continue;
            }
            let f = a[r][col];
            for j in 0..n {
                a[r][j] -= f * a[col][j];
                inv[r][j] -= f * inv[col][j];
            }
        }
    }
    inv
}

/// Two competitors, two slices ten units apart, one duel in each.
///
/// The exact precision is written out explicitly here rather than obtained
/// from the crate, so this is an independent check rather than a restatement.
/// Variables are `[a0, b0, a1, b1]`.
#[test]
fn a_two_slice_joint_matches_the_exact_posterior() {
    let mut h = history(GAMMA);
    h.add_events(vec![
        duel("a", "b", 0, 5.0, 2.0),
        duel("a", "b", 10, 4.0, 3.0),
    ])
    .unwrap();
    let report = h.converge().unwrap();
    assert!(report.converged, "{:?}", report.final_step);

    let prior_prec = 1.0 / (SIGMA0 * SIGMA0);
    let drift_prec = 1.0 / (10.0 * GAMMA * GAMMA);
    let obs_prec = 1.0 / (SCORE_SIGMA * SCORE_SIGMA + 2.0 * BETA * BETA);

    let mut lambda = vec![vec![0.0; 4]; 4];
    // priors on the first appearances
    lambda[0][0] += prior_prec;
    lambda[1][1] += prior_prec;
    // drift a0-a1 and b0-b1
    for (p, q) in [(0usize, 2usize), (1, 3)] {
        lambda[p][p] += drift_prec;
        lambda[q][q] += drift_prec;
        lambda[p][q] -= drift_prec;
        lambda[q][p] -= drift_prec;
    }
    // one duel per slice: contrast (+1, -1) on that slice's variables
    for (p, q) in [(0usize, 1usize), (2, 3)] {
        lambda[p][p] += obs_prec;
        lambda[q][q] += obs_prec;
        lambda[p][q] -= obs_prec;
        lambda[q][p] -= obs_prec;
    }
    let cov = inverse(lambda);

    // The crate reads each competitor at their latest appearance: a1, b1.
    let exact_gap = (cov[2][2] + cov[3][3] - 2.0 * cov[2][3]).sqrt();
    let got = h
        .joint()
        .unwrap()
        .posterior_of(&[(&"a", 1.0), (&"b", -1.0)])
        .unwrap();
    assert!(
        (got.sigma() - exact_gap).abs() / exact_gap < 1e-9,
        "difference: got {} exact {exact_gap}",
        got.sigma()
    );

    let exact_single = cov[2][2].sqrt();
    let got_single = h.joint().unwrap().posterior_of(&[(&"a", 1.0)]).unwrap();
    assert!(
        (got_single.sigma() - exact_single).abs() / exact_single < 1e-9,
        "single node: got {} exact {exact_single}",
        got_single.sigma()
    );
}

/// The case that motivated this: competitors read at *different* slices, with
/// the last slice holding only one of them. Under the old latest-slice joint
/// this was `UnknownKey`.
#[test]
fn competitors_last_seen_in_different_slices_are_comparable() {
    let mut h = history(GAMMA);
    h.add_events(vec![
        duel("a", "b", 0, 5.0, 2.0),
        duel("a", "c", 10, 4.0, 3.0),
        // the final slice holds one duel that does not involve b at all
        duel("a", "c", 20, 6.0, 1.0),
    ])
    .unwrap();
    let _ = h.converge().unwrap();

    // b last appeared at time 0; a and c at time 20. All three must resolve.
    for (x, y) in [("a", "b"), ("b", "c"), ("a", "c")] {
        let g = h
            .joint()
            .unwrap()
            .posterior_of(&[(&x, 1.0), (&y, -1.0)])
            .unwrap_or_else(|e| panic!("{x} - {y} should resolve across slices: {e}"));
        assert!(g.sigma() > 0.0 && g.sigma().is_finite());
    }
}

/// The mean must agree with what message passing reports, which is exact even
/// with cycles. Only the second moment needs the joint.
#[test]
fn means_agree_with_the_marginals() {
    let mut h = history(GAMMA);
    h.add_events(vec![
        duel("a", "b", 0, 5.0, 2.0),
        duel("b", "c", 5, 3.0, 1.0),
        duel("a", "c", 10, 4.0, 2.0),
    ])
    .unwrap();
    let _ = h.converge().unwrap();

    for k in ["a", "b", "c"] {
        let marginal = h.current_skill(&k).unwrap().mu();
        let joint = h.joint().unwrap().posterior_of(&[(&k, 1.0)]).unwrap().mu();
        assert!(
            (marginal - joint).abs() < 1e-9,
            "{k}: marginal {marginal}, joint {joint}"
        );
    }
}

/// With zero drift a competitor has one latent skill however many slices it
/// appears in, so spreading the same events over time must not change the
/// answer. This exercises the appearance-merging path.
#[test]
fn zero_drift_makes_slice_layout_irrelevant() {
    let spread = {
        let mut h = history(0.0);
        h.add_events(vec![
            duel("a", "b", 0, 5.0, 2.0),
            duel("a", "b", 10, 4.0, 3.0),
            duel("a", "b", 20, 6.0, 1.0),
        ])
        .unwrap();
        let _ = h.converge().unwrap();
        h.joint()
            .unwrap()
            .posterior_of(&[(&"a", 1.0), (&"b", -1.0)])
            .unwrap()
    };
    let together = {
        let mut h = history(0.0);
        h.add_events(vec![
            duel("a", "b", 0, 5.0, 2.0),
            duel("a", "b", 0, 4.0, 3.0),
            duel("a", "b", 0, 6.0, 1.0),
        ])
        .unwrap();
        let _ = h.converge().unwrap();
        h.joint()
            .unwrap()
            .posterior_of(&[(&"a", 1.0), (&"b", -1.0)])
            .unwrap()
    };

    assert!(
        (spread.sigma() - together.sigma()).abs() < 1e-9,
        "zero drift: spread {} vs together {}",
        spread.sigma(),
        together.sigma()
    );
}

/// More drift means less is carried forward from old evidence, so a comparison
/// against a competitor last seen long ago must widen.
#[test]
fn drift_widens_a_comparison_across_time() {
    let mut previous = 0.0;
    for gamma in [0.0f64, 0.1, 0.5, 2.0] {
        let mut h = history(gamma);
        h.add_events(vec![
            duel("a", "b", 0, 5.0, 2.0),
            duel("a", "c", 100, 4.0, 3.0),
        ])
        .unwrap();
        let _ = h.converge().unwrap();

        // b was last seen at time 0; a at time 100.
        let g = h
            .joint()
            .unwrap()
            .posterior_of(&[(&"a", 1.0), (&"b", -1.0)])
            .unwrap();
        assert!(
            g.sigma() > previous,
            "gamma={gamma}: sigma {} did not exceed {previous}",
            g.sigma()
        );
        previous = g.sigma();
    }
}

/// `posterior_of_at` pins the reading to a moment, where `posterior_of` takes
/// each competitor wherever they were last seen.
#[test]
fn posterior_of_at_reads_as_of_a_time() {
    let mut h = history(GAMMA);
    h.add_events(vec![
        duel("a", "b", 0, 5.0, 2.0),
        duel("a", "b", 10, 4.0, 3.0),
        duel("a", "b", 20, 6.0, 1.0),
    ])
    .unwrap();
    let _ = h.converge().unwrap();

    let early = h
        .joint()
        .unwrap()
        .posterior_of_at(0, &[(&"a", 1.0), (&"b", -1.0)])
        .unwrap();
    let late = h
        .joint()
        .unwrap()
        .posterior_of_at(20, &[(&"a", 1.0), (&"b", -1.0)])
        .unwrap();
    let latest = h
        .joint()
        .unwrap()
        .posterior_of(&[(&"a", 1.0), (&"b", -1.0)])
        .unwrap();

    // Asking as of the final slice is the same as asking for the latest.
    assert!((late.mu() - latest.mu()).abs() < 1e-9);
    assert!((late.sigma() - latest.sigma()).abs() < 1e-9);

    // Reading at time 0 is a different quantity, and the smoothed estimate
    // there is informed by everything that came after.
    assert!(
        (early.mu() - late.mu()).abs() > 1e-6,
        "as-of-0 and as-of-20 should differ: {} vs {}",
        early.mu(),
        late.mu()
    );

    // A time before any event has nothing to read.
    assert!(
        h.joint()
            .unwrap()
            .posterior_of_at(-1, &[(&"a", 1.0)])
            .is_err()
    );
}

/// Times between slices resolve to the latest appearance at or before them.
#[test]
fn a_time_between_slices_reads_the_previous_appearance() {
    let mut h = history(GAMMA);
    h.add_events(vec![
        duel("a", "b", 0, 5.0, 2.0),
        duel("a", "b", 100, 4.0, 3.0),
    ])
    .unwrap();
    let _ = h.converge().unwrap();

    let at_zero = h
        .joint()
        .unwrap()
        .posterior_of_at(0, &[(&"a", 1.0)])
        .unwrap();
    let between = h
        .joint()
        .unwrap()
        .posterior_of_at(50, &[(&"a", 1.0)])
        .unwrap();
    assert!((at_zero.mu() - between.mu()).abs() < 1e-12);
    assert!((at_zero.sigma() - between.sigma()).abs() < 1e-12);
}