//! Synthetic bandit tests: LinUCB and Thompson sampling.
//!
//! Each stage resolves a query with a known probability. Each round the
//! router picks a stage, the pick is run the way `LearnedCascade` dispatches
//! it (Kv or Graph falling through to Neural on a miss, Neural alone, or for
//! Cache the full walk), and every stage that ran reports its outcome and
//! cost. The router must learn the dispatch that spends least, not the stage
//! likeliest to succeed: before, both bandits converged on Neural.

use eoc_core::Stage;
use eoc_route_learned::bandit::{LinUcbRouter, ThompsonSamplingRouter};
use eoc_route_learned::router::ALL_STAGES;
use rand::{RngExt, SeedableRng};
use rand_distr::{Distribution, Normal};
use std::collections::HashMap;

const DIM: usize = 4;

fn cost(stage: Stage) -> u64 {
    match stage {
        Stage::Cache => 1_000,
        Stage::Kv => 50_000,
        Stage::Graph => 500_000,
        Stage::Neural => 50_000_000,
    }
}

/// The stages a dispatch to `stage` runs, each with whether it resolved.
fn dispatch(stage: Stage, resolves: &HashMap<Stage, bool>) -> Vec<(Stage, bool)> {
    let order: &[Stage] = match stage {
        Stage::Cache => &ALL_STAGES,
        Stage::Kv => &[Stage::Kv, Stage::Neural],
        Stage::Graph => &[Stage::Graph, Stage::Neural],
        Stage::Neural => &[Stage::Neural],
    };
    let mut ran = Vec::new();
    for s in order {
        ran.push((*s, resolves[s]));
        if resolves[s] {
            break;
        }
    }
    ran
}

/// Pick counts over the last `tail` of `rounds`, and the joules spent there.
struct Run {
    tail_picks: HashMap<Stage, usize>,
    tail_microjoules: u64,
}

fn run(
    rounds: usize,
    tail: usize,
    chance: [f64; 4],
    mut pick: impl FnMut(&[f32]) -> Stage,
    mut learn: impl FnMut(&[f32], Stage, bool, u64),
) -> Run {
    let mut rng = rand::rngs::StdRng::seed_from_u64(0xBEEF);
    let noise = Normal::new(0.0_f32, 1.0).expect("valid normal");
    let mut out = Run {
        tail_picks: HashMap::new(),
        tail_microjoules: 0,
    };
    for round in 0..rounds {
        // A constant feature first, so LinUCB can learn each stage's base rate.
        let mut ctx = vec![1.0_f32];
        ctx.extend((1..DIM).map(|_| 0.1 * noise.sample(&mut rng)));
        let resolves: HashMap<Stage, bool> = ALL_STAGES
            .iter()
            .zip(chance)
            .map(|(s, p)| (*s, rng.random_bool(p)))
            .collect();
        let stage = pick(&ctx);
        for (s, ok) in dispatch(stage, &resolves) {
            learn(&ctx, s, ok, cost(s));
            if round >= rounds - tail {
                out.tail_microjoules += cost(s);
            }
        }
        if round >= rounds - tail {
            *out.tail_picks.entry(stage).or_insert(0) += 1;
        }
    }
    out
}

fn linucb(rounds: usize, tail: usize, chance: [f64; 4]) -> Run {
    let router = std::cell::RefCell::new(LinUcbRouter::new(DIM, 0.1, 1.0));
    run(
        rounds,
        tail,
        chance,
        |ctx| router.borrow().pick(ctx).expect("pick").0,
        |ctx, s, ok, uj| {
            let mut r = router.borrow_mut();
            r.update(ctx, s, if ok { 1.0 } else { 0.0 })
                .expect("update");
            r.observe_cost(s, uj);
        },
    )
}

fn thompson(rounds: usize, tail: usize, chance: [f64; 4]) -> Run {
    let router = std::cell::RefCell::new(ThompsonSamplingRouter::new(7));
    run(
        rounds,
        tail,
        chance,
        |_| router.borrow().pick().0,
        |_, s, ok, uj| {
            let mut r = router.borrow_mut();
            r.update(s, ok);
            r.observe_cost(s, uj);
        },
    )
}

fn share(run: &Run, stage: Stage, tail: usize) -> f32 {
    *run.tail_picks.get(&stage).unwrap_or(&0) as f32 / tail as f32
}

/// Graph resolves 95% of queries and nothing cheaper resolves any: skip to
/// Graph, at about 3 million µJ a query, not Neural at 50 million.
#[test]
fn bandits_skip_to_the_stage_that_resolves() {
    let chance = [0.0, 0.0, 0.95, 1.0];
    for (name, r) in [
        ("linucb", linucb(3000, 500, chance)),
        ("thompson", thompson(3000, 500, chance)),
    ] {
        println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules);
        assert!(
            share(&r, Stage::Graph, 500) > 0.8,
            "{name} picks Graph only {}",
            share(&r, Stage::Graph, 500)
        );
        assert!(
            share(&r, Stage::Neural, 500) < 0.05,
            "{name} still skips to Neural"
        );
        assert!(
            r.tail_microjoules / 500 < 5_000_000,
            "{name} spends {} µJ a query",
            r.tail_microjoules / 500
        );
    }
}

/// Nothing short of Neural resolves anything, so going straight to Neural,
/// at 50 million µJ, is cheapest; the full walk costs 50.551 million. The
/// bandits come within 1% of Neural. They still probe Kv now and then: a
/// probe costs a thousandth of a Neural call and would save the whole call
/// if it hit, so it is worth trying until the bound on Kv's chance falls
/// below 0.1%.
#[test]
fn bandits_come_within_a_percent_of_neural_when_nothing_cheaper_resolves() {
    let chance = [0.0, 0.0, 0.0, 1.0];
    for (name, r) in [
        ("linucb", linucb(4000, 500, chance)),
        ("thompson", thompson(4000, 500, chance)),
    ] {
        println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules);
        let per_query = r.tail_microjoules / 500;
        assert!(
            per_query < 50_500_000,
            "{name} spends {per_query} µJ a query"
        );
    }
}

/// The old test's arms. Neural is likeliest to succeed (0.8), and the old
/// bandits learned to pick it every time. But the cheap stages resolve often
/// enough that walking them first costs least, so the walk is the pick.
#[test]
fn bandits_walk_when_cheap_stages_often_resolve() {
    let chance = [0.1, 0.2, 0.3, 0.8];
    for (name, r) in [
        ("linucb", linucb(1000, 500, chance)),
        ("thompson", thompson(1000, 500, chance)),
    ] {
        println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules);
        assert!(
            share(&r, Stage::Cache, 500) > 0.8,
            "{name} walks only {}",
            share(&r, Stage::Cache, 500)
        );
        assert!(
            share(&r, Stage::Neural, 500) < 0.05,
            "{name} skips to Neural"
        );
    }
}