//! State import/export round-trip across all router families.

use eoc_route_learned::bandit::{LinUcbRouter, ThompsonSamplingRouter};
use eoc_route_learned::classifier::LogRegRouter;
use eoc_route_learned::matrix_factorization::MfRouter;
use eoc_route_learned::router::LearnedRouter;

#[test]
fn mf_roundtrip() {
    let router = MfRouter::new(8, 4, 11);
    let state = router.export_state();
    let restored = MfRouter::import_state(state).expect("import");
    assert_eq!(router.latent_dim(), restored.latent_dim());
    assert_eq!(router.embedding_dim(), restored.embedding_dim());
}

#[test]
fn logreg_roundtrip() {
    let router = LogRegRouter::new(8, 1e-4, 11);
    let state = router.export_state();
    let restored = LogRegRouter::import_state(state).expect("import");
    assert_eq!(router.embedding_dim(), restored.embedding_dim());
}

#[test]
fn linucb_roundtrip() {
    let router = LinUcbRouter::new(8, 1.0, 1.0);
    let state = router.export_state();
    let restored = LinUcbRouter::import_state(state).expect("import");
    assert_eq!(router.dim, restored.dim);
    assert_eq!(router.alpha, restored.alpha);
}

#[test]
fn thompson_roundtrip() {
    let router = ThompsonSamplingRouter::new(7);
    let state = router.export_state();
    let restored = ThompsonSamplingRouter::import_state(state).expect("import");
    assert_eq!(router.seed, restored.seed);
}

#[test]
fn rejects_wrong_algorithm_tag() {
    let r = MfRouter::new(4, 4, 0);
    let mut s = r.export_state();
    s.algorithm = "wrong".into();
    assert!(MfRouter::import_state(s).is_err());
}

/// The bandits learn each stage's cost. A learned cost survives export, and
/// a state saved before they learned costs imports with the defaults rather
/// than failing.
#[test]
fn bandit_stage_costs_roundtrip_and_older_states_import() {
    use eoc_core::Stage;
    use eoc_route_learned::router::default_stage_joules;

    fn without_joules(state: &mut eoc_route_learned::router::RouterState) {
        let mut v: serde_json::Value = serde_json::from_slice(&state.weights_blob).unwrap();
        v.as_object_mut()
            .unwrap()
            .remove("joules")
            .expect("the field is exported");
        state.weights_blob = serde_json::to_vec(&v).unwrap();
    }

    let mut lin = LinUcbRouter::new(4, 1.0, 1.0);
    lin.observe_cost(Stage::Neural, 10_000);
    let restored = LinUcbRouter::import_state(lin.export_state()).expect("import");
    assert_eq!(restored.joules, lin.joules);
    let mut old = lin.export_state();
    without_joules(&mut old);
    assert_eq!(
        LinUcbRouter::import_state(old).expect("import").joules,
        default_stage_joules()
    );

    let mut ts = ThompsonSamplingRouter::new(3);
    ts.observe_cost(Stage::Kv, 7);
    let restored = ThompsonSamplingRouter::import_state(ts.export_state()).expect("import");
    assert_eq!(restored.joules, ts.joules);
    let mut old = ts.export_state();
    without_joules(&mut old);
    assert_eq!(
        ThompsonSamplingRouter::import_state(old)
            .expect("import")
            .joules,
        default_stage_joules()
    );
}