//! Synthetic test for the MF router.
//!
//! Each query has one row per stage. Each cheap stage has its own coverage:
//! the chance it resolves a query is a logistic function of one feature of
//! the embedding (Cache's centred at 0.5, Kv's at 0, Graph's at −0.5), and
//! Neural resolves everything. The router must learn to skip to the stage
//! that is cheapest to answer from in expectation, not to the one likeliest
//! to succeed, which is Neural.
//!
//! The coverage is soft, as real coverage is, and of the family the router
//! models, so its learned chances can be right. They have to be: with the
//! default costs a wrong skip pays a Neural call, a thousand times Graph,
//! while a right one saves at most a hundredth of one. A router sure of
//! itself on noise-free boundaries skips wrongly at their edges and spends
//! more than the plain walk.
//!
//! With the default costs the plain walk is itself within a fraction of a
//! percent of the best choices, so there is little for routing to win, and
//! the router lands within 1% of the best. Where Kv and Graph cost a real
//! share of a Neural call, the walk is 4–5% above the best, but at this
//! training budget MF's chances are off by 5–9% a stage, and the router lands
//! at about the walk's cost rather than capturing that. Training longer made
//! it beat the walk by 2% while its chances got worse, so that is a property
//! of the estimator, not of the rule; the test holds it to not losing.

use std::collections::HashMap;

use async_trait::async_trait;
use eoc_core::{Query, Stage};
use eoc_route_learned::evaluation::evaluate;
use eoc_route_learned::matrix_factorization::MfRouter;
use eoc_route_learned::router::{
    ALL_STAGES, LearnedRouter, RouterState, StagePrediction, cheapest_in_expectation,
    default_stage_joules, expected_joules,
};
use eoc_route_learned::training::{Example, TrainingConfig, train_mf};
use rand::{RngExt, SeedableRng};
use rand_distr::{Distribution, Normal};

const DIM: usize = 8;

/// Kv and Graph at a real share of a Neural call.
fn dear_cheap_stages() -> HashMap<Stage, u64> {
    let mut c = default_stage_joules();
    c.insert(Stage::Kv, 2_000_000);
    c.insert(Stage::Graph, 10_000_000);
    c
}

/// The true chance that each stage resolves the query with this embedding,
/// in `ALL_STAGES` order.
fn chances(emb: &[f32]) -> [f32; 4] {
    let s = |x: f32| 1.0 / (1.0 + (-4.0 * x).exp());
    [s(emb[0] - 0.5), s(emb[1]), s(emb[2] + 0.5), 1.0]
}

fn synthesise(n: usize, seed: u64, costs: &HashMap<Stage, u64>) -> Vec<Example> {
    let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
    let dist = Normal::new(0.0_f32, 1.0).expect("valid normal");
    let mut out = Vec::with_capacity(n * 4);
    for _ in 0..n {
        let emb: Vec<f32> = (0..DIM).map(|_| dist.sample(&mut rng)).collect();
        let p = chances(&emb);
        for (i, stage) in ALL_STAGES.into_iter().enumerate() {
            let success = rng.random_bool(f64::from(p[i]));
            out.push(Example::new(emb.clone(), stage, success, costs[&stage]));
        }
    }
    out
}

fn router(costs: &HashMap<Stage, u64>) -> MfRouter {
    let cfg = TrainingConfig {
        epochs: 30,
        learning_rate: 0.1,
        l2: 1e-5,
        ..TrainingConfig::default()
    };
    train_mf(&synthesise(1000, 17, costs), cfg)
}

/// The rule the routers used before: the stage likeliest to succeed.
struct LikeliestToSucceed<'a>(&'a MfRouter);

#[async_trait]
impl LearnedRouter for LikeliestToSucceed<'_> {
    async fn route(&self, q: &Query) -> StagePrediction {
        let probs = self.0.predict_probs(q.embedding.as_ref().unwrap()).unwrap();
        let best = (0..4).fold(0, |b, i| if probs[i] > probs[b] { i } else { b });
        StagePrediction::new(ALL_STAGES[best], probs[best])
    }
    async fn observe(&mut self, _q: &Query, _s: Stage, _ok: bool, _j: u64) {}
    fn export_state(&self) -> RouterState {
        unimplemented!()
    }
    fn import_state(_s: RouterState) -> eoc_route_learned::Result<Self> {
        unimplemented!()
    }
}

/// Expected µJ over held-out queries, priced at their true chances: the best
/// choices, the router's, the plain walk's and the likeliest-stage rule's;
/// and the share of queries where the router's choice is the best one.
struct Priced {
    best: f64,
    routed: f64,
    walk: f64,
    likeliest: f64,
    agree: f64,
}

fn priced(costs: &HashMap<Stage, u64>) -> Priced {
    let router = router(costs);
    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
    let dist = Normal::new(0.0_f32, 1.0).expect("valid normal");
    let mut out = Priced {
        best: 0.0,
        routed: 0.0,
        walk: 0.0,
        likeliest: 0.0,
        agree: 0.0,
    };
    let mut fit = [0.0_f64; 4];
    let n = 2000;
    for _ in 0..n {
        let emb: Vec<f32> = (0..DIM).map(|_| dist.sample(&mut rng)).collect();
        let p = chances(&emb);
        let (choice, _, learned) = router.predict(&emb).expect("predict");
        let ideal = cheapest_in_expectation(&p, costs);
        let top = (0..4).fold(0, |b, i| if learned[i] > learned[b] { i } else { b });
        for i in 0..4 {
            fit[i] += f64::from((learned[i] - p[i]).abs()) / f64::from(n);
        }
        out.agree += f64::from(u8::from(choice == ideal)) / f64::from(n);
        out.routed += expected_joules(choice, &p, costs);
        out.walk += expected_joules(Stage::Cache, &p, costs);
        out.best += expected_joules(ideal, &p, costs);
        out.likeliest += expected_joules(ALL_STAGES[top], &p, costs);
    }
    println!(
        "expected µJ: best {:.0}, routed {:.0}, walk {:.0}, likeliest {:.0}; agrees with best {:.3}",
        out.best, out.routed, out.walk, out.likeliest, out.agree
    );
    println!("mean |learned − true| chance per stage: {fit:.3?}");
    out
}

#[test]
fn with_the_default_costs_the_router_lands_within_a_percent_of_the_best() {
    let r = priced(&default_stage_joules());
    assert!(
        r.agree > 0.9,
        "the router's choice is the best for only {} of queries",
        r.agree
    );
    assert!(
        r.routed <= r.best * 1.01,
        "routed {:.0} µJ is over 1% above the best {:.0}",
        r.routed,
        r.best
    );
    assert!(
        r.routed * 2.0 < r.likeliest,
        "routed {:.0} µJ is not far below the likeliest-stage rule's {:.0}",
        r.routed,
        r.likeliest
    );
}

#[test]
fn where_cheap_stages_are_dear_the_router_does_not_lose_to_the_walk() {
    let r = priced(&dear_cheap_stages());
    assert!(
        r.routed <= r.walk * 1.01,
        "routed {:.0} µJ is over 1% above the walk's {:.0}",
        r.routed,
        r.walk
    );
    assert!(
        r.routed * 2.0 < r.likeliest,
        "routed {:.0} µJ is not far below the likeliest-stage rule's {:.0}",
        r.routed,
        r.likeliest
    );
}

/// The same, replayed over sampled outcomes by `evaluate`.
#[tokio::test]
async fn mf_router_spends_far_less_than_the_likeliest_stage_rule() {
    let costs = default_stage_joules();
    let router = router(&costs);
    let test = synthesise(500, 42, &costs);
    let routed = evaluate(&router, &test).await;
    let likeliest = evaluate(&LikeliestToSucceed(&router), &test).await;
    println!("routed:     {routed:?}");
    println!("likeliest:  {likeliest:?}");
    assert_eq!((routed.n, routed.unscored), (500, 0));
    assert!(
        routed.total_microjoules * 2 < likeliest.total_microjoules,
        "routed {} µJ is not far below the likeliest-stage rule's {}",
        routed.total_microjoules,
        likeliest.total_microjoules
    );
}