use std::collections::HashMap;

use crate::{
    agent::{self, Agent},
    batch::{self, Batch},
    gaussian::Gaussian,
    player::Player,
    sort_time, tuple_gt, tuple_max, Index, BETA, GAMMA, MU, P_DRAW, SIGMA,
};

#[derive(Clone)]
pub struct HistoryBuilder {
    time: bool,
    mu: f64,
    sigma: f64,
    beta: f64,
    gamma: f64,
    p_draw: f64,
    online: bool,
}

impl HistoryBuilder {
    pub fn time(mut self, time: bool) -> Self {
        self.time = time;
        self
    }

    pub fn mu(mut self, mu: f64) -> Self {
        self.mu = mu;
        self
    }

    pub fn sigma(mut self, sigma: f64) -> Self {
        self.sigma = sigma;
        self
    }

    pub fn beta(mut self, beta: f64) -> Self {
        self.beta = beta;
        self
    }

    pub fn gamma(mut self, gamma: f64) -> Self {
        self.gamma = gamma;
        self
    }

    pub fn p_draw(mut self, p_draw: f64) -> Self {
        self.p_draw = p_draw;
        self
    }

    pub fn online(mut self, online: bool) -> Self {
        self.online = online;
        self
    }

    pub fn build(self) -> History {
        History {
            size: 0,
            batches: Vec::new(),
            agents: HashMap::new(),
            time: self.time,
            mu: self.mu,
            sigma: self.sigma,
            beta: self.beta,
            gamma: self.gamma,
            p_draw: self.p_draw,
            online: self.online,
        }
    }
}

impl Default for HistoryBuilder {
    fn default() -> Self {
        Self {
            time: true,
            mu: MU,
            sigma: SIGMA,
            beta: BETA,
            gamma: GAMMA,
            p_draw: P_DRAW,
            online: false,
        }
    }
}

pub struct History {
    size: usize,
    pub(crate) batches: Vec<Batch>,
    agents: HashMap<Index, Agent>,
    time: bool,
    mu: f64,
    sigma: f64,
    beta: f64,
    gamma: f64,
    p_draw: f64,
    online: bool,
}

impl Default for History {
    fn default() -> Self {
        Self {
            size: 0,
            batches: Vec::new(),
            agents: HashMap::new(),
            time: true,
            mu: MU,
            sigma: SIGMA,
            beta: BETA,
            gamma: GAMMA,
            p_draw: P_DRAW,
            online: false,
        }
    }
}

impl History {
    pub fn builder() -> HistoryBuilder {
        HistoryBuilder::default()
    }

    fn iteration(&mut self) -> (f64, f64) {
        let mut step = (0.0, 0.0);

        agent::clean(self.agents.values_mut(), false);

        for j in (0..self.batches.len() - 1).rev() {
            for agent in self.batches[j + 1].skills.keys() {
                self.agents.get_mut(agent).unwrap().message =
                    self.batches[j + 1].backward_prior_out(agent, &self.agents);
            }

            let old = self.batches[j].posteriors();

            self.batches[j].new_backward_info(&self.agents);

            let new = self.batches[j].posteriors();

            step = old
                .iter()
                .fold(step, |step, (a, old)| tuple_max(step, old.delta(new[a])));
        }

        agent::clean(self.agents.values_mut(), false);

        for j in 1..self.batches.len() {
            for agent in self.batches[j - 1].skills.keys() {
                self.agents.get_mut(agent).unwrap().message =
                    self.batches[j - 1].forward_prior_out(agent);
            }

            let old = self.batches[j].posteriors();

            self.batches[j].new_forward_info(&self.agents);

            let new = self.batches[j].posteriors();

            step = old
                .iter()
                .fold(step, |step, (a, old)| tuple_max(step, old.delta(new[a])));
        }

        if self.batches.len() == 1 {
            let old = self.batches[0].posteriors();

            self.batches[0].iteration(0, &self.agents);

            let new = self.batches[0].posteriors();

            step = old
                .iter()
                .fold(step, |step, (a, old)| tuple_max(step, old.delta(new[a])));
        }

        step
    }

    pub fn convergence(
        &mut self,
        iterations: usize,
        epsilon: f64,
        verbose: bool,
    ) -> ((f64, f64), usize) {
        let mut step = (f64::INFINITY, f64::INFINITY);
        let mut i = 0;

        while tuple_gt(step, epsilon) && i < iterations {
            if verbose {
                print!("Iteration = {}", i);
            }

            step = self.iteration();

            i += 1;

            if verbose {
                println!(", step = {:?}", step);
            }
        }

        if verbose {
            println!("End");
        }

        (step, i)
    }

    pub fn learning_curves(&self) -> HashMap<Index, Vec<(i64, Gaussian)>> {
        let mut data: HashMap<Index, Vec<(i64, Gaussian)>> = HashMap::new();

        for b in &self.batches {
            for (agent, skill) in b.skills.iter() {
                let point = (b.time, skill.posterior());

                if let Some(entry) = data.get_mut(agent) {
                    entry.push(point);
                } else {
                    data.insert(*agent, vec![point]);
                }
            }
        }

        data
    }

    pub fn log_evidence(&mut self, forward: bool, targets: &[Index]) -> f64 {
        self.batches
            .iter()
            .map(|batch| batch.log_evidence(self.online, targets, forward, &self.agents))
            .sum()
    }

    pub fn add_events(
        &mut self,
        composition: Vec<Vec<Vec<Index>>>,
        results: Vec<Vec<f64>>,
        times: Vec<i64>,
        weights: Vec<Vec<Vec<f64>>>,
    ) {
        self.add_events_with_prior(composition, results, times, weights, HashMap::new())
    }

    pub fn add_events_with_prior(
        &mut self,
        composition: Vec<Vec<Vec<Index>>>,
        results: Vec<Vec<f64>>,
        times: Vec<i64>,
        weights: Vec<Vec<Vec<f64>>>,
        mut priors: HashMap<Index, Player>,
    ) {
        assert!(times.is_empty() || self.time, "length(times)>0 but !h.time");
        assert!(
            !times.is_empty() || !self.time,
            "length(times)==0 but h.time"
        );
        assert!(
            results.is_empty() || results.len() == composition.len(),
            "(length(results) > 0) & (length(composition) != length(results))"
        );
        assert!(
            times.is_empty() || times.len() == composition.len(),
            "length(times) > 0) & (length(composition) != length(times))"
        );
        assert!(
            weights.is_empty() || weights.len() == composition.len(),
            "(length(weights) > 0) & (length(composition) != length(weights))"
        );

        agent::clean(self.agents.values_mut(), true);

        let mut this_agent = Vec::with_capacity(1024);

        for agent in composition.iter().flatten().flatten() {
            if this_agent.contains(agent) {
                continue;
            }

            this_agent.push(*agent);

            if !self.agents.contains_key(agent) {
                self.agents.insert(
                    *agent,
                    Agent {
                        player: priors.remove(agent).unwrap_or_else(|| {
                            Player::new(
                                Gaussian::from_ms(self.mu, self.sigma),
                                self.beta,
                                self.gamma,
                            )
                        }),
                        ..Default::default()
                    },
                );
            }
        }

        let n = composition.len();
        let o = if self.time {
            sort_time(&times, false)
        } else {
            (0..composition.len()).collect::<Vec<_>>()
        };

        let mut i = 0;
        let mut k = 0;

        while i < n {
            let mut j = i + 1;
            let t = if self.time { times[o[i]] } else { i as i64 + 1 };

            while self.time && j < n && times[o[j]] == t {
                j += 1;
            }

            while (!self.time && (self.size > k))
                || (self.time && self.batches.len() > k && self.batches[k].time < t)
            {
                let batch = &mut self.batches[k];

                if k > 0 {
                    batch.new_forward_info(&self.agents);
                }

                // TODO: Is it faster to iterate over agents in batch instead?
                for agent_idx in &this_agent {
                    if let Some(skill) = batch.skills.get_mut(agent_idx) {
                        skill.elapsed =
                            batch::compute_elapsed(self.agents[agent_idx].last_time, batch.time);

                        let agent = self.agents.get_mut(agent_idx).unwrap();

                        agent.last_time = if self.time { batch.time } else { i64::MAX };
                        agent.message = batch.forward_prior_out(agent_idx);
                    }
                }

                k += 1;
            }

            let composition = (i..j)
                .map(|e| composition[o[e]].clone())
                .collect::<Vec<_>>();

            let results = if results.is_empty() {
                Vec::new()
            } else {
                (i..j).map(|e| results[o[e]].clone()).collect::<Vec<_>>()
            };

            let weights = if weights.is_empty() {
                Vec::new()
            } else {
                (i..j).map(|e| weights[o[e]].clone()).collect::<Vec<_>>()
            };

            if self.time && self.batches.len() > k && self.batches[k].time == t {
                let batch = &mut self.batches[k];
                batch.add_events(composition, results, weights, &self.agents);

                for agent_idx in batch.skills.keys() {
                    let agent = self.agents.get_mut(agent_idx).unwrap();

                    agent.last_time = if self.time { t } else { i64::MAX };
                    agent.message = batch.forward_prior_out(agent_idx);
                }
            } else {
                let mut batch: Batch = Batch::new(t, self.p_draw);
                batch.add_events(composition, results, weights, &self.agents);

                self.batches.insert(k, batch);

                let batch = &self.batches[k];

                for agent_idx in batch.skills.keys() {
                    let agent = self.agents.get_mut(agent_idx).unwrap();

                    agent.last_time = if self.time { t } else { i64::MAX };
                    agent.message = batch.forward_prior_out(agent_idx);
                }

                k += 1;
            }

            i = j;
        }

        while self.time && self.batches.len() > k {
            let batch = &mut self.batches[k];

            batch.new_forward_info(&self.agents);

            // TODO: Is it faster to iterate over agents in batch instead?
            for agent_idx in &this_agent {
                if let Some(skill) = batch.skills.get_mut(agent_idx) {
                    skill.elapsed =
                        batch::compute_elapsed(self.agents[agent_idx].last_time, batch.time);

                    let agent = self.agents.get_mut(agent_idx).unwrap();

                    agent.last_time = if self.time { batch.time } else { i64::MAX };
                    agent.message = batch.forward_prior_out(agent_idx);
                }
            }

            k += 1;
        }

        self.size += n;
    }
}

#[cfg(test)]
mod tests {
    use approx::assert_ulps_eq;

    use crate::{Game, Gaussian, IndexMap, Player, EPSILON, ITERATIONS, P_DRAW};

    use super::*;

    #[test]
    fn test_init() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");

        let composition = vec![
            vec![vec![a], vec![b]],
            vec![vec![a], vec![c]],
            vec![vec![b], vec![c]],
        ];
        let results = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]];

        let mut priors = HashMap::new();

        for agent in [a, b, c] {
            priors.insert(
                agent,
                Player::new(
                    Gaussian::from_ms(25.0, 25.0 / 3.0),
                    25.0 / 6.0,
                    0.15 * 25.0 / 3.0,
                ),
            );
        }

        let mut h = History::default();

        h.add_events_with_prior(composition, results, vec![1, 2, 3], vec![], priors);

        let p0 = h.batches[0].posteriors();

        assert_ulps_eq!(
            p0[&a],
            Gaussian::from_ms(29.205220, 7.194481),
            epsilon = 1e-6
        );

        let observed = h.batches[1].skills[&a].forward.sigma;
        let gamma: f64 = 0.15 * 25.0 / 3.0;
        let expected = (gamma.powi(2) + h.batches[0].skills[&a].posterior().sigma.powi(2)).sqrt();

        assert_ulps_eq!(observed, expected, epsilon = 0.000001);

        let observed = h.batches[1].skills[&a].posterior();

        let w = [vec![1.0], vec![1.0]];
        let p = Game::new(
            h.batches[1].events[0].within_priors(false, false, &h.batches[1].skills, &h.agents),
            &[0.0, 1.0],
            &w,
            P_DRAW,
        )
        .posteriors();
        let expected = p[0][0];

        assert_ulps_eq!(observed, expected, epsilon = 1e-6);
    }

    #[test]
    fn test_one_batch() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");

        let composition = vec![
            vec![vec![a], vec![b]],
            vec![vec![b], vec![c]],
            vec![vec![c], vec![a]],
        ];
        let results = vec![vec![1.0, 0.0], vec![1.0, 0.0], vec![1.0, 0.0]];
        let times = vec![1, 1, 1];

        let mut priors = HashMap::new();

        for agent in [a, b, c] {
            priors.insert(
                agent,
                Player::new(
                    Gaussian::from_ms(25.0, 25.0 / 3.0),
                    25.0 / 6.0,
                    0.15 * 25.0 / 3.0,
                ),
            );
        }

        let mut h1 = History::default();

        h1.add_events_with_prior(composition, results, times, vec![], priors);

        assert_ulps_eq!(
            h1.batches[0].skills[&a].posterior(),
            Gaussian::from_ms(22.904409, 6.010330),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h1.batches[0].skills[&c].posterior(),
            Gaussian::from_ms(25.110318, 5.866311),
            epsilon = 1e-6
        );

        h1.convergence(ITERATIONS, EPSILON, false);

        assert_ulps_eq!(
            h1.batches[0].skills[&a].posterior(),
            Gaussian::from_ms(25.000000, 5.419212),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h1.batches[0].skills[&c].posterior(),
            Gaussian::from_ms(25.000000, 5.419212),
            epsilon = 1e-6
        );

        let composition = vec![
            vec![vec![a], vec![b]],
            vec![vec![b], vec![c]],
            vec![vec![c], vec![a]],
        ];
        let results = vec![vec![1.0, 0.0], vec![1.0, 0.0], vec![1.0, 0.0]];
        let times = vec![1, 2, 3];

        let mut priors = HashMap::new();

        for agent in [a, b, c] {
            priors.insert(
                agent,
                Player::new(
                    Gaussian::from_ms(25.0, 25.0 / 3.0),
                    25.0 / 6.0,
                    25.0 / 300.0,
                ),
            );
        }

        let mut h2 = History::default();

        h2.add_events_with_prior(composition, results, times, vec![], priors);

        assert_ulps_eq!(
            h2.batches[2].skills[&a].posterior(),
            Gaussian::from_ms(22.903522, 6.011017),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h2.batches[2].skills[&c].posterior(),
            Gaussian::from_ms(25.110702, 5.866811),
            epsilon = 1e-6
        );

        h2.convergence(ITERATIONS, EPSILON, false);

        assert_ulps_eq!(
            h2.batches[2].skills[&a].posterior(),
            Gaussian::from_ms(24.998668, 5.420053),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h2.batches[2].skills[&c].posterior(),
            Gaussian::from_ms(25.000532, 5.419827),
            epsilon = 1e-6
        );
    }

    #[test]
    fn test_learning_curves() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");

        let composition = vec![
            vec![vec![a], vec![b]],
            vec![vec![b], vec![c]],
            vec![vec![c], vec![a]],
        ];
        let results = vec![vec![1.0, 0.0], vec![1.0, 0.0], vec![1.0, 0.0]];
        let times = vec![5, 6, 7];

        let mut priors = HashMap::new();

        for agent in [a, b, c] {
            priors.insert(
                agent,
                Player::new(
                    Gaussian::from_ms(25.0, 25.0 / 3.0),
                    25.0 / 6.0,
                    25.0 / 300.0,
                ),
            );
        }

        let mut h = History::default();

        h.add_events_with_prior(composition, results, times, vec![], priors);
        h.convergence(ITERATIONS, EPSILON, false);

        let lc = h.learning_curves();

        let aj_e = lc[&a].len();
        let cj_e = lc[&c].len();

        assert_eq!(lc[&a][0].0, 5);
        assert_eq!(lc[&a][aj_e - 1].0, 7);

        assert_ulps_eq!(
            lc[&a][aj_e - 1].1,
            Gaussian::from_ms(24.998668, 5.420053),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            lc[&c][cj_e - 1].1,
            Gaussian::from_ms(25.000532, 5.419827),
            epsilon = 1e-6
        );
    }

    #[test]
    fn test_env_ttt() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");

        let composition = vec![
            vec![vec![a], vec![b]],
            vec![vec![a], vec![c]],
            vec![vec![b], vec![c]],
        ];
        let results = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]];

        let mut h = History::builder()
            .mu(25.0)
            .sigma(25.0 / 3.0)
            .beta(25.0 / 6.0)
            .gamma(25.0 / 300.0)
            .time(false)
            .build();

        h.add_events(composition, results, vec![], vec![]);

        h.convergence(ITERATIONS, EPSILON, false);

        assert_eq!(h.batches[2].skills[&b].elapsed, 1);
        assert_eq!(h.batches[2].skills[&c].elapsed, 1);

        assert_ulps_eq!(
            h.batches[0].skills[&a].posterior(),
            Gaussian::from_ms(25.000267, 5.419381),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[0].skills[&b].posterior(),
            Gaussian::from_ms(24.999465, 5.419425),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[2].skills[&b].posterior(),
            Gaussian::from_ms(25.000532, 5.419696),
            epsilon = 1e-6
        );
    }

    #[test]
    fn test_teams() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");
        let d = index_map.get_or_create("d");
        let e = index_map.get_or_create("e");
        let f = index_map.get_or_create("f");

        let composition = vec![
            vec![vec![a, b], vec![c, d]],
            vec![vec![e, f], vec![b, c]],
            vec![vec![a, d], vec![e, f]],
        ];
        let results = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]];

        let mut h = History::builder()
            .mu(0.0)
            .sigma(6.0)
            .beta(1.0)
            .gamma(0.0)
            .time(false)
            .build();

        h.add_events(composition, results, vec![], vec![]);

        let trueskill_log_evidence = h.log_evidence(false, &[]);
        let trueskill_log_evidence_online = h.log_evidence(true, &[]);

        assert_ulps_eq!(
            trueskill_log_evidence,
            trueskill_log_evidence_online,
            epsilon = 1e-6
        );

        assert_ulps_eq!(
            h.batches[0].skills[&b].posterior().mu,
            -1.0 * h.batches[0].skills[&c].posterior().mu,
            epsilon = 1e-6
        );

        let evidence_second_event = h.log_evidence(false, &[b]).exp() * 2.0;
        assert_ulps_eq!(0.5, evidence_second_event, epsilon = 1e-6);

        let evidence_third_event = h.log_evidence(false, &[a]).exp() * 2.0;
        assert_ulps_eq!(0.669885, evidence_third_event, epsilon = 1e-6);

        h.convergence(ITERATIONS, EPSILON, false);

        let loocv_hat = h.log_evidence(false, &[]).exp();
        let p_d_m_hat = h.log_evidence(true, &[]).exp();

        assert_ulps_eq!(loocv_hat, 0.241027, epsilon = 1e-6);
        assert_ulps_eq!(p_d_m_hat, 0.172432, epsilon = 1e-6);

        assert_ulps_eq!(
            h.batches[0].skills[&a].posterior(),
            h.batches[0].skills[&b].posterior(),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[0].skills[&c].posterior(),
            h.batches[0].skills[&d].posterior(),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[1].skills[&e].posterior(),
            h.batches[1].skills[&f].posterior(),
            epsilon = 1e-6
        );

        assert_ulps_eq!(
            h.batches[0].skills[&a].posterior(),
            Gaussian::from_ms(4.084902, 5.106919),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[0].skills[&c].posterior(),
            Gaussian::from_ms(-0.533029, 5.106919),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[2].skills[&e].posterior(),
            Gaussian::from_ms(-3.551872, 5.154569),
            epsilon = 1e-6
        );
    }

    #[test]
    fn test_add_events() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");

        let composition = vec![
            vec![vec![a], vec![b]],
            vec![vec![a], vec![c]],
            vec![vec![b], vec![c]],
        ];
        let results = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]];

        let mut h = History::builder()
            .mu(0.0)
            .sigma(2.0)
            .beta(1.0)
            .gamma(0.0)
            .time(false)
            .build();

        h.add_events(composition.clone(), results.clone(), vec![], vec![]);

        h.convergence(ITERATIONS, EPSILON, false);

        assert_eq!(h.batches[2].skills[&b].elapsed, 1);
        assert_eq!(h.batches[2].skills[&c].elapsed, 1);

        assert_ulps_eq!(
            h.batches[0].skills[&a].posterior(),
            Gaussian::from_ms(0.000000, 1.300610),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[0].skills[&b].posterior(),
            Gaussian::from_ms(0.000000, 1.300610),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[2].skills[&b].posterior(),
            Gaussian::from_ms(0.000000, 1.300610),
            epsilon = 1e-6
        );

        h.add_events(composition, results, vec![], vec![]);

        assert_eq!(h.batches.len(), 6);

        assert_eq!(
            h.batches
                .iter()
                .map(|b| b.get_composition())
                .collect::<Vec<_>>(),
            vec![
                vec![vec![vec![a], vec![b]]],
                vec![vec![vec![a], vec![c]]],
                vec![vec![vec![b], vec![c]]],
                vec![vec![vec![a], vec![b]]],
                vec![vec![vec![a], vec![c]]],
                vec![vec![vec![b], vec![c]]]
            ]
        );

        h.convergence(ITERATIONS, EPSILON, false);

        assert_ulps_eq!(
            h.batches[0].skills[&a].posterior(),
            Gaussian::from_ms(0.000000, 0.931236),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[3].skills[&a].posterior(),
            Gaussian::from_ms(0.000000, 0.931236),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[3].skills[&b].posterior(),
            Gaussian::from_ms(0.000000, 0.931236),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[5].skills[&b].posterior(),
            Gaussian::from_ms(0.000000, 0.931236),
            epsilon = 1e-6
        );
    }

    #[test]
    fn test_only_add_events() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");

        let composition = vec![
            vec![vec![a], vec![b]],
            vec![vec![a], vec![c]],
            vec![vec![b], vec![c]],
        ];
        let results = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]];

        let mut h = History::builder()
            .mu(0.0)
            .sigma(2.0)
            .beta(1.0)
            .gamma(0.0)
            .time(false)
            .build();

        h.add_events(composition.clone(), results.clone(), vec![], vec![]);

        h.convergence(ITERATIONS, EPSILON, false);

        assert_eq!(h.batches[2].skills[&b].elapsed, 1);
        assert_eq!(h.batches[2].skills[&c].elapsed, 1);

        assert_ulps_eq!(
            h.batches[0].skills[&a].posterior(),
            Gaussian::from_ms(0.000000, 1.300610),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[0].skills[&b].posterior(),
            Gaussian::from_ms(0.000000, 1.300610),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[2].skills[&b].posterior(),
            Gaussian::from_ms(0.000000, 1.300610),
            epsilon = 1e-6
        );

        h.add_events(composition, results, vec![], vec![]);

        assert_eq!(h.batches.len(), 6);

        assert_eq!(
            h.batches
                .iter()
                .map(|b| b.get_composition())
                .collect::<Vec<_>>(),
            vec![
                vec![vec![vec![a], vec![b]]],
                vec![vec![vec![a], vec![c]]],
                vec![vec![vec![b], vec![c]]],
                vec![vec![vec![a], vec![b]]],
                vec![vec![vec![a], vec![c]]],
                vec![vec![vec![b], vec![c]]]
            ]
        );

        h.convergence(ITERATIONS, EPSILON, false);

        assert_ulps_eq!(
            h.batches[0].skills[&a].posterior(),
            Gaussian::from_ms(0.000000, 0.931236),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[3].skills[&a].posterior(),
            Gaussian::from_ms(0.000000, 0.931236),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[3].skills[&b].posterior(),
            Gaussian::from_ms(0.000000, 0.931236),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            h.batches[5].skills[&b].posterior(),
            Gaussian::from_ms(0.000000, 0.931236),
            epsilon = 1e-6
        );
    }

    #[test]
    fn test_log_evidence() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");

        let composition = vec![vec![vec![a], vec![b]], vec![vec![b], vec![a]]];

        let mut h = History::builder().time(false).build();

        h.add_events(composition.clone(), vec![], vec![], vec![]);

        let p_d_m_2 = h.log_evidence(false, &[]).exp() * 2.0;

        assert_ulps_eq!(p_d_m_2, 0.17650911, epsilon = 1e-6);
        assert_ulps_eq!(
            p_d_m_2,
            h.log_evidence(true, &[]).exp() * 2.0,
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            p_d_m_2,
            h.log_evidence(true, &[a]).exp() * 2.0,
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            p_d_m_2,
            h.log_evidence(false, &[a]).exp() * 2.0,
            epsilon = 1e-6
        );

        h.convergence(11, EPSILON, false);

        let loocv_approx_2 = h.log_evidence(false, &[]).exp().sqrt();

        assert_ulps_eq!(loocv_approx_2, 0.001976774, epsilon = 0.000001);

        let p_d_m_approx_2 = h.log_evidence(true, &[]).exp() * 2.0;

        assert!(loocv_approx_2 - p_d_m_approx_2 < 1e-4);

        assert_ulps_eq!(
            loocv_approx_2,
            h.log_evidence(true, &[b]).exp() * 2.0,
            epsilon = 1e-4
        );

        let mut h = History::builder().time(false).build();

        h.add_events(composition, vec![], vec![], vec![]);

        assert_ulps_eq!(
            ((0.5f64 * 0.1765).ln() / 2.0).exp(),
            (h.log_evidence(false, &[]) / 2.0).exp(),
            epsilon = 1e-4
        );
    }

    #[test]
    fn test_add_events_with_time() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");

        let composition = vec![
            vec![vec![a], vec![b]],
            vec![vec![a], vec![c]],
            vec![vec![b], vec![c]],
        ];
        let results = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]];

        let mut h = History::builder()
            .mu(0.0)
            .sigma(2.0)
            .beta(1.0)
            .gamma(0.0)
            .build();

        h.add_events(
            composition.clone(),
            results.clone(),
            vec![0, 10, 20],
            vec![],
        );

        h.convergence(ITERATIONS, EPSILON, false);

        h.add_events(composition, results, vec![15, 10, 0], vec![]);

        assert_eq!(h.batches.len(), 4);

        assert_eq!(
            h.batches
                .iter()
                .map(|batch| batch.events.len())
                .collect::<Vec<_>>(),
            vec![2, 2, 1, 1]
        );

        assert_eq!(
            h.batches
                .iter()
                .map(|b| b.get_composition())
                .collect::<Vec<_>>(),
            vec![
                vec![vec![vec![a], vec![b]], vec![vec![b], vec![c]]],
                vec![vec![vec![a], vec![c]], vec![vec![a], vec![c]]],
                vec![vec![vec![a], vec![b]]],
                vec![vec![vec![b], vec![c]]]
            ]
        );

        assert_eq!(
            h.batches
                .iter()
                .map(|b| b.get_results())
                .collect::<Vec<_>>(),
            vec![
                vec![vec![1.0, 0.0], vec![1.0, 0.0]],
                vec![vec![0.0, 1.0], vec![0.0, 1.0]],
                vec![vec![1.0, 0.0]],
                vec![vec![1.0, 0.0]]
            ]
        );

        let end = h.batches.len() - 1;

        assert_eq!(h.batches[0].skills[&c].elapsed, 0);
        assert_eq!(h.batches[end].skills[&c].elapsed, 10);

        assert_eq!(h.batches[0].skills[&a].elapsed, 0);
        assert_eq!(h.batches[2].skills[&a].elapsed, 5);

        assert_eq!(h.batches[0].skills[&b].elapsed, 0);
        assert_eq!(h.batches[end].skills[&b].elapsed, 5);

        h.convergence(ITERATIONS, EPSILON, false);

        assert_ulps_eq!(
            h.batches[0].skills[&b].posterior(),
            h.batches[end].skills[&b].posterior(),
            epsilon = 1e-6
        );

        assert_ulps_eq!(
            h.batches[0].skills[&c].posterior(),
            h.batches[end].skills[&c].posterior(),
            epsilon = 1e-6
        );

        assert_ulps_eq!(
            h.batches[0].skills[&c].posterior(),
            h.batches[0].skills[&b].posterior(),
            epsilon = 1e-6
        );

        // ---------------------------------------

        let composition = vec![
            vec![vec![a], vec![b]],
            vec![vec![c], vec![a]],
            vec![vec![b], vec![c]],
        ];

        let mut h = History::builder()
            .mu(0.0)
            .sigma(2.0)
            .beta(1.0)
            .gamma(0.0)
            .build();

        h.add_events(composition.clone(), vec![], vec![0, 10, 20], vec![]);

        h.convergence(ITERATIONS, EPSILON, false);

        h.add_events(composition, vec![], vec![15, 10, 0], vec![]);

        assert_eq!(h.batches.len(), 4);

        assert_eq!(
            h.batches
                .iter()
                .map(|batch| batch.events.len())
                .collect::<Vec<_>>(),
            vec![2, 2, 1, 1]
        );

        assert_eq!(
            h.batches
                .iter()
                .map(|b| b.get_composition())
                .collect::<Vec<_>>(),
            vec![
                vec![vec![vec![a], vec![b]], vec![vec![b], vec![c]]],
                vec![vec![vec![c], vec![a]], vec![vec![c], vec![a]]],
                vec![vec![vec![a], vec![b]]],
                vec![vec![vec![b], vec![c]]]
            ]
        );

        assert_eq!(
            h.batches
                .iter()
                .map(|b| b.get_results())
                .collect::<Vec<_>>(),
            vec![
                vec![vec![1.0, 0.0], vec![1.0, 0.0]],
                vec![vec![1.0, 0.0], vec![1.0, 0.0]],
                vec![vec![1.0, 0.0]],
                vec![vec![1.0, 0.0]]
            ]
        );

        let end = h.batches.len() - 1;

        assert_eq!(h.batches[0].skills[&c].elapsed, 0);
        assert_eq!(h.batches[end].skills[&c].elapsed, 10);

        assert_eq!(h.batches[0].skills[&a].elapsed, 0);
        assert_eq!(h.batches[2].skills[&a].elapsed, 5);

        assert_eq!(h.batches[0].skills[&b].elapsed, 0);
        assert_eq!(h.batches[end].skills[&b].elapsed, 5);

        h.convergence(ITERATIONS, EPSILON, false);

        assert_ulps_eq!(
            h.batches[0].skills[&b].posterior(),
            h.batches[end].skills[&b].posterior(),
            epsilon = 1e-6
        );

        assert_ulps_eq!(
            h.batches[0].skills[&c].posterior(),
            h.batches[end].skills[&c].posterior(),
            epsilon = 1e-6
        );

        assert_ulps_eq!(
            h.batches[0].skills[&c].posterior(),
            h.batches[0].skills[&b].posterior(),
            epsilon = 1e-6
        );
    }

    #[test]
    fn test_1vs1_weighted() {
        let mut index_map = IndexMap::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");

        let composition = vec![vec![vec![a], vec![b]], vec![vec![b], vec![a]]];
        let weights = vec![vec![vec![5.0], vec![4.0]], vec![vec![5.0], vec![4.0]]];

        let mut h = History::builder()
            .mu(2.0)
            .sigma(6.0)
            .beta(1.0)
            .gamma(0.0)
            .time(false)
            .build();

        h.add_events(composition, vec![], vec![], weights);

        let lc = h.learning_curves();

        assert_ulps_eq!(
            lc[&a][0].1,
            Gaussian::from_ms(5.537659, 4.758722),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            lc[&b][0].1,
            Gaussian::from_ms(-0.830127, 5.239568),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            lc[&a][1].1,
            Gaussian::from_ms(1.792277, 4.099566),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            lc[&b][1].1,
            Gaussian::from_ms(4.845533, 3.747616),
            epsilon = 1e-6
        );

        h.convergence(ITERATIONS, EPSILON, false);

        let lc = h.learning_curves();

        assert_ulps_eq!(lc[&a][0].1, lc[&a][0].1, epsilon = 1e-6);
        assert_ulps_eq!(lc[&b][0].1, lc[&a][0].1, epsilon = 1e-6);
        assert_ulps_eq!(lc[&a][1].1, lc[&a][0].1, epsilon = 1e-6);
        assert_ulps_eq!(lc[&b][1].1, lc[&a][0].1, epsilon = 1e-6);
    }
}