use std::collections::{HashMap, HashSet};

use crate::{
    agent::Agent, game::Game, gaussian::Gaussian, player::Player, tuple_gt, tuple_max, Index, N_INF,
};

#[derive(Debug)]
pub(crate) struct Skill {
    pub(crate) forward: Gaussian,
    backward: Gaussian,
    likelihood: Gaussian,
    pub(crate) elapsed: u64,
    pub(crate) online: Gaussian,
}

impl Default for Skill {
    fn default() -> Self {
        Self {
            forward: N_INF,
            backward: N_INF,
            likelihood: N_INF,
            elapsed: 0,
            online: N_INF,
        }
    }
}

#[derive(Debug)]
struct Item {
    agent: Index,
    likelihood: Gaussian,
}

#[derive(Debug)]
struct Team {
    items: Vec<Item>,
    output: f64,
}

#[derive(Debug)]
pub(crate) struct Event {
    teams: Vec<Team>,
    evidence: f64,
    weights: Vec<Vec<f64>>,
}

impl Event {
    fn outputs(&self) -> Vec<f64> {
        self.teams
            .iter()
            .map(|team| team.output)
            .collect::<Vec<_>>()
    }
}

#[derive(Debug)]
pub struct Batch {
    pub(crate) events: Vec<Event>,
    pub(crate) skills: HashMap<Index, Skill>,
    pub(crate) time: u64,
    p_draw: f64,
}

impl Batch {
    pub(crate) fn new(
        composition: Vec<Vec<Vec<Index>>>,
        results: Vec<Vec<f64>>,
        weights: Vec<Vec<Vec<f64>>>,
        time: u64,
        p_draw: f64,
        agents: &mut HashMap<Index, Agent>,
    ) -> Self {
        assert!(
            results.is_empty() || results.len() == composition.len(),
            "TODO: Add a comment here"
        );
        assert!(
            weights.is_empty() || weights.len() == composition.len(),
            "TODO: Add a comment here"
        );

        let this_agent = composition
            .iter()
            .flatten()
            .flatten()
            .cloned()
            .collect::<HashSet<_>>();

        let elapsed = this_agent
            .iter()
            .map(|&idx| (idx, compute_elapsed(agents[&idx].last_time, time)))
            .collect::<HashMap<_, _>>();

        let skills = this_agent
            .iter()
            .map(|&idx| {
                (
                    idx,
                    Skill {
                        forward: agents[&idx].receive(elapsed[&idx]),
                        elapsed: elapsed[&idx],
                        ..Default::default()
                    },
                )
            })
            .collect::<HashMap<_, _>>();

        let events = (0..composition.len())
            .map(|e| {
                let teams = (0..composition[e].len())
                    .map(|t| {
                        let items = (0..composition[e][t].len())
                            .map(|a| Item {
                                agent: composition[e][t][a],
                                likelihood: N_INF,
                            })
                            .collect::<Vec<_>>();

                        Team {
                            items,
                            output: if results.is_empty() {
                                (composition[e].len() - (t + 1)) as f64
                            } else {
                                results[e][t]
                            },
                        }
                    })
                    .collect::<Vec<_>>();

                Event {
                    teams,
                    evidence: 0.0,
                    weights: if weights.is_empty() {
                        Vec::new()
                    } else {
                        weights[e].clone()
                    },
                }
            })
            .collect::<Vec<_>>();

        let mut this = Self {
            time,
            events,
            skills,
            p_draw,
        };

        this.iteration(0, agents);

        this
    }

    pub(crate) fn add_events(
        &mut self,
        composition: Vec<Vec<Vec<Index>>>,
        results: Vec<Vec<f64>>,
        weights: Vec<Vec<Vec<f64>>>,
        agents: &mut HashMap<Index, Agent>,
    ) {
        let this_agent = composition
            .iter()
            .flatten()
            .flatten()
            .cloned()
            .collect::<HashSet<_>>();

        for idx in this_agent {
            let elapsed = compute_elapsed(agents[&idx].last_time, self.time);

            if let Some(skill) = self.skills.get_mut(&idx) {
                skill.elapsed = elapsed;
                skill.forward = agents[&idx].receive(elapsed);
            } else {
                self.skills.insert(
                    idx,
                    Skill {
                        forward: agents[&idx].receive(elapsed),
                        elapsed,
                        ..Default::default()
                    },
                );
            }
        }

        let from = self.events.len();

        for e in 0..composition.len() {
            let teams = (0..composition[e].len())
                .map(|t| {
                    let items = (0..composition[e][t].len())
                        .map(|a| Item {
                            agent: composition[e][t][a],
                            likelihood: N_INF,
                        })
                        .collect::<Vec<_>>();

                    Team {
                        items,
                        output: if results.is_empty() {
                            (composition[e].len() - (t + 1)) as f64
                        } else {
                            results[e][t]
                        },
                    }
                })
                .collect::<Vec<_>>();

            let event = Event {
                teams,
                evidence: 0.0,
                weights: if weights.is_empty() {
                    Vec::new()
                } else {
                    weights[e].clone()
                },
            };

            self.events.push(event);
        }

        self.iteration(from, agents);
    }

    pub fn posterior(&self, agent: Index) -> Gaussian {
        let skill = &self.skills[&agent];

        skill.likelihood * skill.backward * skill.forward
    }

    pub(crate) fn posteriors(&self) -> HashMap<Index, Gaussian> {
        self.skills
            .keys()
            .map(|&idx| (idx, self.posterior(idx)))
            .collect::<HashMap<_, _>>()
    }

    fn within_prior(
        &self,
        item: &Item,
        online: bool,
        forward: bool,
        agents: &mut HashMap<Index, Agent>,
    ) -> Player {
        let r = &agents[&item.agent].player;

        if online {
            Player::new(self.skills[&item.agent].online, r.beta, r.gamma)
        } else if forward {
            Player::new(self.skills[&item.agent].forward, r.beta, r.gamma)
        } else {
            let wp = self.posterior(item.agent) / item.likelihood;

            Player::new(wp, r.beta, r.gamma)
        }
    }

    pub(crate) fn within_priors(
        &self,
        event: usize,
        online: bool,
        forward: bool,
        agents: &mut HashMap<Index, Agent>,
    ) -> Vec<Vec<Player>> {
        self.events[event]
            .teams
            .iter()
            .map(|team| {
                team.items
                    .iter()
                    .map(|item| self.within_prior(item, online, forward, agents))
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    }

    pub(crate) fn iteration(&mut self, from: usize, agents: &mut HashMap<Index, Agent>) {
        for e in from..self.events.len() {
            let teams = self.within_priors(e, false, false, agents);
            let result = self.events[e].outputs();

            let g = Game::new(teams, result, self.events[e].weights.clone(), self.p_draw);

            for (t, team) in self.events[e].teams.iter_mut().enumerate() {
                for (i, item) in team.items.iter_mut().enumerate() {
                    self.skills.get_mut(&item.agent).unwrap().likelihood =
                        (self.skills[&item.agent].likelihood / item.likelihood)
                            * g.likelihoods[t][i];

                    item.likelihood = g.likelihoods[t][i];
                }
            }

            self.events[e].evidence = g.evidence;
        }
    }

    pub(crate) fn convergence(&mut self, agents: &mut HashMap<Index, Agent>) -> usize {
        let epsilon = 1e-6;
        let iterations = 20;

        let mut step = (f64::INFINITY, f64::INFINITY);
        let mut i = 0;

        while tuple_gt(step, epsilon) && i < iterations {
            let old = self.posteriors();

            self.iteration(0, agents);

            let new = self.posteriors();

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

            i += 1;
        }

        i
    }

    pub(crate) fn forward_prior_out(&self, agent: &Index) -> Gaussian {
        let skill = &self.skills[agent];

        skill.forward * skill.likelihood
    }

    pub(crate) fn backward_prior_out(
        &self,
        agent: &Index,
        agents: &mut HashMap<Index, Agent>,
    ) -> Gaussian {
        let skill = &self.skills[agent];
        let n = skill.likelihood * skill.backward;

        n.forget(agents[agent].player.gamma, skill.elapsed)
    }

    pub(crate) fn new_backward_info(&mut self, agents: &mut HashMap<Index, Agent>) {
        for (agent, skill) in self.skills.iter_mut() {
            skill.backward = agents[agent].message;
        }

        self.iteration(0, agents);
    }

    pub(crate) fn new_forward_info(&mut self, agents: &mut HashMap<Index, Agent>) {
        for (agent, skill) in self.skills.iter_mut() {
            skill.forward = agents[agent].receive(skill.elapsed);
        }

        self.iteration(0, agents);
    }

    pub(crate) fn log_evidence(
        &self,
        online: bool,
        targets: &[Index],
        forward: bool,
        agents: &mut HashMap<Index, Agent>,
    ) -> f64 {
        if targets.is_empty() {
            if online || forward {
                self.events
                    .iter()
                    .enumerate()
                    .map(|(e, event)| {
                        Game::new(
                            self.within_priors(e, online, forward, agents),
                            event.outputs(),
                            event.weights.clone(),
                            self.p_draw,
                        )
                        .evidence
                        .ln()
                    })
                    .sum()
            } else {
                self.events.iter().map(|event| event.evidence.ln()).sum()
            }
        } else {
            if online || forward {
                self.events
                    .iter()
                    .enumerate()
                    .filter(|(_, event)| {
                        event
                            .teams
                            .iter()
                            .flat_map(|team| &team.items)
                            .any(|item| targets.contains(&item.agent))
                    })
                    .map(|(e, event)| {
                        Game::new(
                            self.within_priors(e, online, forward, agents),
                            event.outputs(),
                            event.weights.clone(),
                            self.p_draw,
                        )
                        .evidence
                        .ln()
                    })
                    .sum()
            } else {
                self.events
                    .iter()
                    .filter(|event| {
                        event
                            .teams
                            .iter()
                            .flat_map(|team| &team.items)
                            .any(|item| targets.contains(&item.agent))
                    })
                    .map(|event| event.evidence.ln())
                    .sum()
            }
        }
    }

    pub(crate) fn get_composition(&self) -> Vec<Vec<Vec<Index>>> {
        self.events
            .iter()
            .map(|event| {
                event
                    .teams
                    .iter()
                    .map(|team| team.items.iter().map(|item| item.agent).collect::<Vec<_>>())
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    }

    pub(crate) fn get_results(&self) -> Vec<Vec<f64>> {
        self.events
            .iter()
            .map(|event| {
                event
                    .teams
                    .iter()
                    .map(|team| team.output)
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    }
}

pub(crate) fn compute_elapsed(last_time: u64, actual_time: u64) -> u64 {
    if last_time == u64::MIN {
        0
    } else if last_time == u64::MAX {
        1
    } else {
        actual_time - last_time
    }
}

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

    use crate::{agent::Agent, player::Player, IndexMap};

    use super::*;

    #[test]
    fn test_one_event_each() {
        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 mut agents = HashMap::new();

        for agent in [a, b, c, d, e, f] {
            agents.insert(
                agent,
                Agent {
                    player: Player::new(Gaussian::new(25.0, 25.0 / 3.0), 25.0 / 6.0, 25.0 / 300.0),
                    ..Default::default()
                },
            );
        }

        let mut batch = Batch::new(
            vec![
                vec![vec![a], vec![b]],
                vec![vec![c], vec![d]],
                vec![vec![e], vec![f]],
            ],
            vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]],
            vec![],
            0,
            0.0,
            &mut agents,
        );

        let post = batch.posteriors();

        assert_ulps_eq!(post[&a], Gaussian::new(29.205220, 7.194481), epsilon = 1e-6);
        assert_ulps_eq!(post[&b], Gaussian::new(20.794779, 7.194481), epsilon = 1e-6);
        assert_ulps_eq!(post[&c], Gaussian::new(20.794779, 7.194481), epsilon = 1e-6);
        assert_ulps_eq!(post[&d], Gaussian::new(29.205220, 7.194481), epsilon = 1e-6);
        assert_ulps_eq!(post[&e], Gaussian::new(29.205220, 7.194481), epsilon = 1e-6);
        assert_ulps_eq!(post[&f], Gaussian::new(20.794779, 7.194481), epsilon = 1e-6);

        assert_eq!(batch.convergence(&mut agents), 1);
    }

    #[test]
    fn test_same_strength() {
        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 mut agents = HashMap::new();

        for agent in [a, b, c, d, e, f] {
            agents.insert(
                agent,
                Agent {
                    player: Player::new(Gaussian::new(25.0, 25.0 / 3.0), 25.0 / 6.0, 25.0 / 300.0),
                    ..Default::default()
                },
            );
        }

        let mut batch = Batch::new(
            vec![
                vec![vec![a], vec![b]],
                vec![vec![a], vec![c]],
                vec![vec![b], vec![c]],
            ],
            vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]],
            vec![],
            0,
            0.0,
            &mut agents,
        );

        let post = batch.posteriors();

        assert_ulps_eq!(post[&a], Gaussian::new(24.960978, 6.298544), epsilon = 1e-6);
        assert_ulps_eq!(post[&b], Gaussian::new(27.095590, 6.010330), epsilon = 1e-6);
        assert_ulps_eq!(post[&c], Gaussian::new(24.889681, 5.866311), epsilon = 1e-6);

        assert!(batch.convergence(&mut agents) > 1);

        let post = batch.posteriors();

        assert_ulps_eq!(post[&a], Gaussian::new(25.000000, 5.419212), epsilon = 1e-6);
        assert_ulps_eq!(post[&b], Gaussian::new(25.000000, 5.419212), epsilon = 1e-6);
        assert_ulps_eq!(post[&c], Gaussian::new(25.000000, 5.419212), 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 d = index_map.get_or_create("d");
        let e = index_map.get_or_create("e");
        let f = index_map.get_or_create("f");

        let mut agents = HashMap::new();

        for agent in [a, b, c, d, e, f] {
            agents.insert(
                agent,
                Agent {
                    player: Player::new(Gaussian::new(25.0, 25.0 / 3.0), 25.0 / 6.0, 25.0 / 300.0),
                    ..Default::default()
                },
            );
        }

        let mut batch = Batch::new(
            vec![
                vec![vec![a], vec![b]],
                vec![vec![a], vec![c]],
                vec![vec![b], vec![c]],
            ],
            vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]],
            vec![],
            0,
            0.0,
            &mut agents,
        );

        batch.convergence(&mut agents);

        let post = batch.posteriors();

        assert_ulps_eq!(post[&a], Gaussian::new(25.000000, 5.419212), epsilon = 1e-6);
        assert_ulps_eq!(post[&b], Gaussian::new(25.000000, 5.419212), epsilon = 1e-6);
        assert_ulps_eq!(post[&c], Gaussian::new(25.000000, 5.419212), epsilon = 1e-6);

        batch.add_events(
            vec![
                vec![vec![a], vec![b]],
                vec![vec![a], vec![c]],
                vec![vec![b], vec![c]],
            ],
            vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]],
            vec![],
            &mut agents,
        );

        assert_eq!(batch.events.len(), 6);

        batch.convergence(&mut agents);

        let post = batch.posteriors();

        assert_ulps_eq!(post[&a], Gaussian::new(25.000003, 3.880150), epsilon = 1e-6);
        assert_ulps_eq!(post[&b], Gaussian::new(25.000003, 3.880150), epsilon = 1e-6);
        assert_ulps_eq!(post[&c], Gaussian::new(25.000003, 3.880150), epsilon = 1e-6);
    }
}