//! Calibration of the crate's marginals against the EXACT posterior.
//!
//! A scored history is linear-Gaussian — `MarginFactor` encodes
//! `score_a - score_b ~ N(perf_a - perf_b, score_sigma^2)` — so the true joint
//! posterior has a closed form and the crate can be checked against ground
//! truth rather than against intuition. That is not possible for ranked
//! outcomes, whose truncation likelihood EP genuinely approximates.
//!
//! Two things are pinned here, and one is deliberately only recorded.
//!
//! **Pinned: on a tree the crate is exact**, means and variances both. Message
//! passing has no approximation to make when the factor graph has no cycles, so
//! any drift here would be a real defect.
//!
//! **Pinned: means are exact even with cycles.** This is the standard result
//! for Gaussian belief propagation (Weiss & Freeman 2001) and it is what makes
//! ratings trustworthy.
//!
//! **Recorded, not asserted: with cycles, marginal variances are too narrow.**
//! Measured on the round-robin fixture below, the crate reports sigma 1.430
//! where the exact posterior is 2.851 — a ratio of 0.502. That is the known
//! behaviour of loopy Gaussian BP, not a bug in this crate, and it is left
//! unasserted because fixing it is exactly what #46 proposes.
//!
//! Why that matters for a consumer, and why #46 cannot be implemented as "add
//! a covariance accessor": the exact correlation between two nodes here is
//! +0.857, so a consumer computing `sqrt(sa^2 + sb^2)` for a difference
//! overstates its width. But the too-narrow marginals partially cancel that,
//! leaving 1.327x rather than 2.646x. Adding true correlations to these
//! marginals without also correcting them would give 0.765 against a true
//! 1.524 — *overconfident*, which is the unsafe direction.

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

const N: usize = 5;
const MU0: f64 = 0.0;
const SIGMA0: f64 = 6.0;
const BETA: f64 = 1.0;
const SCORE_SIGMA: f64 = 2.0;

/// A STAR: every event touches c0, so the node-event graph is a tree and
/// Gaussian BP is exact. Any discrepancy here is not caused by loops.
fn tree_fixture() -> Vec<(usize, usize, f64)> {
    vec![(0, 1, 3.0), (0, 2, 5.0), (0, 3, 4.0), (0, 4, 6.0)]
}

/// (winner, loser, score_diff)
fn fixture() -> Vec<(usize, usize, f64)> {
    vec![
        (0, 1, 3.0),
        (0, 2, 5.0),
        (1, 2, 2.0),
        (3, 4, 1.0),
        (0, 3, 4.0),
        (1, 4, 2.5),
        (2, 3, 0.5),
        (0, 4, 6.0),
        (1, 3, 1.5),
        (2, 4, 3.0),
    ]
}

/// Invert a small symmetric positive-definite matrix by Gauss-Jordan.
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| if i == j { 1.0 } else { 0.0 }).collect())
        .collect();
    for col in 0..n {
        // partial pivot
        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
}

/// The exact posterior of a linear-Gaussian model:
/// precision = prior precision + sum of a_k a_k^T / v_k.
fn exact_for(obs: &[(usize, usize, f64)]) -> (Vec<f64>, Vec<Vec<f64>>) {
    let mut lambda = vec![vec![0.0; N]; N];
    let mut eta = [0.0; N];
    for (i, row) in lambda.iter_mut().enumerate() {
        row[i] = 1.0 / (SIGMA0 * SIGMA0);
        eta[i] = MU0 / (SIGMA0 * SIGMA0);
    }

    // Each 1v1 observation: d ~ N(x_a - x_b, score_sigma^2 + 2 beta^2)
    let v = SCORE_SIGMA * SCORE_SIGMA + 2.0 * BETA * BETA;
    for &(a, b, d) in obs {
        let mut vec_a = [0.0; N];
        vec_a[a] = 1.0;
        vec_a[b] = -1.0;
        for i in 0..N {
            for j in 0..N {
                lambda[i][j] += vec_a[i] * vec_a[j] / v;
            }
            eta[i] += vec_a[i] * d / v;
        }
    }

    let cov = inverse(lambda);
    let mean: Vec<f64> = (0..N)
        .map(|i| (0..N).map(|j| cov[i][j] * eta[j]).sum())
        .collect();
    (mean, cov)
}

fn key(i: usize) -> &'static str {
    ["c0", "c1", "c2", "c3", "c4"][i]
}

/// Returns (worst mean error, worst sd ratio).
fn fitted(
    obs: &[(usize, usize, f64)],
) -> History<i64, ConstantDrift, trueskill_tt::NullObserver, &'static str> {
    let mut h: History<i64, _, _, &'static str> = History::builder()
        .mu(MU0)
        .sigma(SIGMA0)
        .beta(BETA)
        .score_sigma(SCORE_SIGMA)
        .drift(ConstantDrift::new(0.0))
        .convergence(ConvergenceOptions {
            max_iter: 20_000,
            epsilon: 1e-13,
            alpha: 1.0,
        })
        .build();

    let events: Vec<Event<i64, &'static str>> = obs
        .iter()
        .copied()
        .map(|(a, b, d)| Event {
            time: 1,
            teams: smallvec![
                Team::with_members([Member::new(key(a))]),
                Team::with_members([Member::new(key(b))]),
            ],
            outcome: Outcome::scores([d, 0.0]),
        })
        .collect();
    h.add_events(events).unwrap();
    let report = h.converge().unwrap();
    assert!(
        report.converged,
        "fixture must converge: {:?}",
        report.final_step
    );

    h
}

/// Returns (worst mean error, worst sd ratio gap).
fn run(name: &str, obs: Vec<(usize, usize, f64)>) -> (f64, f64) {
    println!("\n########## {name} ##########");
    let h = fitted(&obs);
    let (mean, cov) = exact_for(&obs);

    println!("\n== marginals: crate vs the exact linear-Gaussian posterior ==");
    println!(
        "{:>4} {:>12} {:>12} {:>12} {:>12} {:>8}",
        "node", "crate mu", "exact mu", "crate sd", "exact sd", "sd ratio"
    );
    for i in 0..N {
        let g = h.current_skill(&key(i)).unwrap();
        let exact_sd = cov[i][i].sqrt();
        println!(
            "{:>4} {:>12.6} {:>12.6} {:>12.6} {:>12.6} {:>8.3}",
            key(i),
            g.mu(),
            mean[i],
            g.sigma(),
            exact_sd,
            g.sigma() / exact_sd
        );
    }

    let mut worst_mean = 0.0f64;
    let mut worst_ratio_gap = 0.0f64;
    for i in 0..N {
        let g = h.current_skill(&key(i)).unwrap();
        worst_mean = worst_mean.max((g.mu() - mean[i]).abs());
        worst_ratio_gap = worst_ratio_gap.max((g.sigma() / cov[i][i].sqrt() - 1.0).abs());
    }

    println!("\n== what a consumer actually computes for a DIFFERENCE ==");
    println!(
        "{:>8} {:>12} {:>14} {:>14} {:>12}",
        "pair", "exact", "naive(exact)", "naive(crate)", "crate err"
    );
    for i in 0..N {
        for j in i + 1..N {
            if i != 0 && j != 1 {
                continue;
            }
            let gi = h.current_skill(&key(i)).unwrap();
            let gj = h.current_skill(&key(j)).unwrap();
            let exact_sd = (cov[i][i] + cov[j][j] - 2.0 * cov[i][j]).sqrt();
            let naive_exact = (cov[i][i] + cov[j][j]).sqrt();
            let naive_crate = (gi.sigma().powi(2) + gj.sigma().powi(2)).sqrt();
            let corr = cov[i][j] / (cov[i][i].sqrt() * cov[j][j].sqrt());
            println!(
                "{:>8} {:>12.6} {:>14.6} {:>14.6} {:>11.3}x   (corr {corr:.4})",
                format!("{}-{}", key(i), key(j)),
                exact_sd,
                naive_exact,
                naive_crate,
                naive_crate / exact_sd
            );
        }
    }

    (worst_mean, worst_ratio_gap)
}

/// With no cycles there is nothing for message passing to approximate.
#[test]
fn on_a_tree_the_marginals_are_exact() {
    let (mean_err, sd_gap) = run("TREE (star: no loops, BP is exact)", tree_fixture());
    assert!(
        mean_err < 1e-9,
        "tree means should be exact, worst error {mean_err}"
    );
    assert!(
        sd_gap < 1e-9,
        "tree sigmas should be exact, worst ratio gap {sd_gap}"
    );
}

/// With cycles the means stay exact — the property ratings depend on — while
/// the variances do not. The variance gap is measured and reported rather than
/// asserted; see the module docs.
#[test]
fn with_cycles_the_means_stay_exact_but_the_variances_shrink() {
    let (mean_err, sd_gap) = run("LOOPY (round robin)", fixture());
    assert!(
        mean_err < 1e-9,
        "loopy means must still be exact, worst error {mean_err}"
    );
    assert!(
        sd_gap > 0.1,
        "the loopy variance gap is the premise of #46; if it has closed, that \
         issue and these docs need revisiting (worst ratio gap {sd_gap})"
    );
}

/// The point of #46: `posterior_of` must reproduce the exact joint, including
/// the correlation that marginals cannot express.
#[test]
fn posterior_of_matches_the_exact_joint() {
    for (name, obs) in [("tree", tree_fixture()), ("loopy", fixture())] {
        let h = fitted(&obs);
        let (_, cov) = exact_for(&obs);

        println!("\n== posterior_of vs exact ({name}) ==");
        println!(
            "{:>12} {:>14} {:>14} {:>10}",
            "functional", "posterior_of", "exact", "ratio"
        );

        for (i, j) in [(0usize, 1usize), (0, 2), (1, 3), (2, 4)] {
            let got = h
                .posterior_of(&[(&key(i), 1.0), (&key(j), -1.0)])
                .expect("scored slice should have a joint");
            let exact_sd = (cov[i][i] + cov[j][j] - 2.0 * cov[i][j]).sqrt();
            println!(
                "{:>12} {:>14.6} {:>14.6} {:>10.4}",
                format!("{}-{}", key(i), key(j)),
                got.sigma(),
                exact_sd,
                got.sigma() / exact_sd
            );
            assert!(
                (got.sigma() - exact_sd).abs() / exact_sd < 1e-9,
                "{name} {}-{}: posterior_of gave {} where the exact joint is {exact_sd}",
                key(i),
                key(j),
                got.sigma()
            );
        }

        // A single competitor: this is where the loopy marginal was 2x narrow.
        for (i, row) in cov.iter().enumerate() {
            let got = h.posterior_of(&[(&key(i), 1.0)]).unwrap();
            let exact_sd = row[i].sqrt();
            assert!(
                (got.sigma() - exact_sd).abs() / exact_sd < 1e-9,
                "{name} {}: posterior_of gave {} where exact is {exact_sd}",
                key(i),
                got.sigma()
            );
        }
        println!("  single-competitor marginals also exact");
    }
}

/// Cost of the dense solve as the slice grows. Recorded, not asserted.
#[test]
#[ignore = "timing probe, run explicitly"]
fn cost_scaling() {
    use std::time::Instant;
    for n in [50usize, 100, 200, 400, 800] {
        let names: Vec<String> = (0..n).map(|i| format!("c{i}")).collect();
        let mut h: History<i64, _, _, String> = History::builder_with_key()
            .score_sigma(2.0)
            .drift(ConstantDrift::new(0.0))
            .convergence(ConvergenceOptions {
                max_iter: 200,
                epsilon: 1e-8,
                alpha: 1.0,
            })
            .build();
        let mut seed = 5u64;
        let mut rnd = move || {
            seed ^= seed << 13;
            seed ^= seed >> 7;
            seed ^= seed << 17;
            seed
        };
        let events: Vec<Event<i64, String>> = (0..n * 4)
            .map(|_| {
                let a = (rnd() as usize) % n;
                let mut b = (rnd() as usize) % n;
                if b == a {
                    b = (b + 1) % n;
                }
                Event {
                    time: 1,
                    teams: smallvec![
                        Team::with_members([Member::new(names[a].clone())]),
                        Team::with_members([Member::new(names[b].clone())]),
                    ],
                    outcome: Outcome::scores([1.0, 0.0]),
                }
            })
            .collect();
        h.add_events(events).unwrap();
        let _ = h.converge().unwrap();

        let t = Instant::now();
        let g = h
            .posterior_of(&[(&names[0], 1.0), (&names[1], -1.0)])
            .unwrap();
        println!("  n={n:>4}: {:>10.2?}  sigma {:.6}", t.elapsed(), g.sigma());
    }
}