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

use crate::{
    agent::Agent, game::Game, gaussian::Gaussian, player::Player, tuple_gt, tuple_max, 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: String,
    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<String, Skill>,
    pub(crate) time: u64,
    p_draw: f64,
}

impl Batch {
    pub(crate) fn new(
        composition: Vec<Vec<Vec<&str>>>,
        results: Vec<Vec<f64>>,
        weights: Vec<Vec<Vec<f64>>>,
        time: u64,
        p_draw: f64,
        agents: &mut HashMap<String, 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(|&a| (a, compute_elapsed(agents[a].last_time, time)))
            .collect::<HashMap<_, _>>();

        let skills = this_agent
            .iter()
            .map(|&a| {
                (
                    a.to_string(),
                    Skill {
                        forward: agents[a].receive(elapsed[a]),
                        elapsed: elapsed[a],
                        ..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].to_string(),
                                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<&str>>>,
        results: Vec<Vec<f64>>,
        weights: Vec<Vec<Vec<f64>>>,
        agents: &mut HashMap<String, Agent>,
    ) {
        let this_agent = composition
            .iter()
            .flatten()
            .flatten()
            .cloned()
            .collect::<HashSet<_>>();

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

            if let Some(skill) = self.skills.get_mut(a) {
                skill.elapsed = elapsed;
                skill.forward = agents[a].receive(elapsed);
            } else {
                self.skills.insert(
                    a.to_string(),
                    Skill {
                        forward: agents[a].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].to_string(),
                            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: &str) -> Gaussian {
        let skill = &self.skills[agent];

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

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

    fn within_prior(
        &self,
        item: &Item,
        online: bool,
        forward: bool,
        agents: &mut HashMap<String, 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<String, 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<String, 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<String, 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: &str) -> Gaussian {
        let skill = &self.skills[agent];

        skill.forward * skill.likelihood
    }

    pub(crate) fn backward_prior_out(
        &self,
        agent: &str,
        agents: &mut HashMap<String, 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<String, 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<String, Agent>) {
        for (agent, skill) in self.skills.iter_mut() {
            skill.forward = agents[agent].receive(skill.elapsed);
        }

        self.iteration(0, agents);
    }

    pub(crate) fn log_evidence2(
        &self,
        online: bool,
        agents2: &Vec<&str>,
        forward: bool,
        agents: &mut HashMap<String, Agent>,
    ) -> f64 {
        if agents2.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| agents2.contains(&item.agent.as_str()))
                    })
                    .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| agents2.contains(&item.agent.as_str()))
                    })
                    .map(|event| event.evidence.ln())
                    .sum()
            }
        }
    }

    pub(crate) fn get_composition(&self) -> Vec<Vec<Vec<&str>>> {
        self.events
            .iter()
            .map(|event| {
                event
                    .teams
                    .iter()
                    .map(|team| {
                        team.items
                            .iter()
                            .map(|item| item.agent.as_str())
                            .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};

    use super::*;

    #[test]
    fn test_one_event_each() {
        let mut agents = HashMap::new();

        for agent in ["a", "b", "c", "d", "e", "f"] {
            agents.insert(
                agent.to_string(),
                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"].mu, 29.205220743876975, epsilon = 0.000001);
        assert_ulps_eq!(post["a"].sigma, 7.194481422570443, epsilon = 0.000001);

        assert_ulps_eq!(post["b"].mu, 20.79477925612302, epsilon = 0.000001);
        assert_ulps_eq!(post["b"].sigma, 7.194481422570443, epsilon = 0.000001);

        assert_ulps_eq!(post["c"].mu, 20.79477925612302, epsilon = 0.000001);
        assert_ulps_eq!(post["c"].sigma, 7.194481422570443, epsilon = 0.000001);

        assert_ulps_eq!(post["d"].mu, 29.205220743876975, epsilon = 0.000001);
        assert_ulps_eq!(post["d"].sigma, 7.194481422570443, epsilon = 0.000001);

        assert_ulps_eq!(post["e"].mu, 29.205220743876975, epsilon = 0.000001);
        assert_ulps_eq!(post["e"].sigma, 7.194481422570443, epsilon = 0.000001);

        assert_ulps_eq!(post["f"].mu, 20.79477925612302, epsilon = 0.000001);
        assert_ulps_eq!(post["f"].sigma, 7.194481422570443, epsilon = 0.000001);

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

    #[test]
    fn test_same_strength() {
        let mut agents = HashMap::new();

        for agent in ["a", "b", "c", "d", "e", "f"] {
            agents.insert(
                agent.to_string(),
                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"].mu, 24.96097857478182, epsilon = 0.000001);
        assert_ulps_eq!(post["a"].sigma, 6.298544763358269, epsilon = 0.000001);

        assert_ulps_eq!(post["b"].mu, 27.095590669107086, epsilon = 0.000001);
        assert_ulps_eq!(post["b"].sigma, 6.010330439043099, epsilon = 0.000001);

        assert_ulps_eq!(post["c"].mu, 24.88968178743119, epsilon = 0.000001);
        assert_ulps_eq!(post["c"].sigma, 5.866311348102562, epsilon = 0.000001);

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

        let post = batch.posteriors();

        assert_ulps_eq!(post["a"].mu, 25.000000, epsilon = 0.000001);
        assert_ulps_eq!(post["a"].sigma, 5.4192120, epsilon = 0.000001);

        assert_ulps_eq!(post["b"].mu, 25.000000, epsilon = 0.000001);
        assert_ulps_eq!(post["b"].sigma, 5.4192120, epsilon = 0.000001);

        assert_ulps_eq!(post["c"].mu, 25.000000, epsilon = 0.000001);
        assert_ulps_eq!(post["c"].sigma, 5.4192120, epsilon = 0.000001);
    }

    #[test]
    fn test_add_events() {
        let mut agents = HashMap::new();

        for agent in ["a", "b", "c", "d", "e", "f"] {
            agents.insert(
                agent.to_string(),
                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.4192120),
            epsilon = 0.000001
        );

        assert_ulps_eq!(
            post["b"],
            Gaussian::new(25.000000, 5.4192120),
            epsilon = 0.000001
        );

        assert_ulps_eq!(
            post["c"],
            Gaussian::new(25.000000, 5.4192120),
            epsilon = 0.000001
        );

        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.00000315330858, 3.880150268080797),
            epsilon = 0.000001
        );
        assert_ulps_eq!(
            post["b"],
            Gaussian::new(25.00000315330858, 3.880150268080797),
            epsilon = 0.000001
        );
        assert_ulps_eq!(
            post["c"],
            Gaussian::new(25.00000315330858, 3.880150268080797),
            epsilon = 0.000001
        );
    }
}