use std::collections::HashSet;

use crate::{message::DiffMessages, utils, variable::TeamVariable, Gaussian, Player, N00};

pub struct Game {
    teams: Vec<Vec<Player>>,
    result: Vec<u16>,
    p_draw: f64,
    pub likelihoods: Vec<Vec<Gaussian>>,
    pub evidence: f64,
}

impl Game {
    pub fn new(teams: Vec<Vec<Player>>, result: Vec<u16>, p_draw: f64) -> Self {
        if !result.is_empty() {
            debug_assert!(
                teams.len() == result.len(),
                "len(result) and (len(teams) != len(result))"
            );
        }

        debug_assert!(p_draw >= 0.0 && p_draw < 1.0, "0.0 <= p_draw < 1.0");

        if p_draw == 0.0 {
            debug_assert!(
                result.iter().collect::<HashSet<_>>().len() == result.len(),
                "(p_draw == 0.0) and (len(result) > 0) and (len(set(result)) != len(result))"
            );
        }

        let mut this = Self {
            teams,
            result,
            p_draw,
            likelihoods: Vec::new(),
            evidence: 0.0,
        };

        this.compute_likelihoods();

        this
    }

    fn performance(&self, index: usize) -> Gaussian {
        self.teams[index]
            .iter()
            .fold(N00, |sum, p| sum + p.performance())
    }

    fn partial_evidence(&mut self, d: &[DiffMessages], margin: &[f64], tie: &[bool], e: usize) {
        let mu = d[e].prior.mu();
        let sigma = d[e].prior.sigma();

        if tie[e] {
            self.evidence *= utils::cdf(margin[e], mu, sigma) - utils::cdf(-margin[e], mu, sigma)
        } else {
            self.evidence *= 1.0 - utils::cdf(margin[e], mu, sigma);
        }
    }

    fn graphical_model(
        &mut self,
    ) -> (
        Vec<usize>,
        Vec<TeamVariable>,
        Vec<DiffMessages>,
        Vec<bool>,
        Vec<f64>,
    ) {
        if self.result.is_empty() {
            self.result = (0..self.teams.len() as u16).rev().collect::<Vec<_>>();
        }

        let r = &self.result;
        let o = utils::sortperm(r, true);

        let t = (0..self.teams.len())
            .map(|e| TeamVariable {
                prior: self.teams[o[e]]
                    .iter()
                    .fold(N00, |sum, p| sum + p.performance()),
                ..Default::default()
            })
            .collect::<Vec<_>>();

        let d = t
            .windows(2)
            .map(|window| DiffMessages {
                prior: window[0].prior - window[1].prior,
                ..Default::default()
            })
            .collect::<Vec<_>>();

        let tie = (0..d.len())
            .map(|e| r[o[e]] == r[o[e + 1]])
            .collect::<Vec<_>>();

        let margin = (0..d.len())
            .map(|e| {
                if self.p_draw == 0.0 {
                    0.0
                } else {
                    let a: f64 = self.teams[o[e]].iter().map(|a| a.beta.powi(2)).sum();
                    let b: f64 = self.teams[o[e + 1]].iter().map(|a| a.beta.powi(2)).sum();

                    utils::compute_margin(self.p_draw, (a + b).sqrt())
                }
            })
            .collect::<Vec<_>>();

        self.evidence = 1.0;

        (o, t, d, tie, margin)
    }

    fn likelihood_analitico(&mut self) -> Vec<Vec<Gaussian>> {
        let (o, t, d, tie, margin) = self.graphical_model();

        self.partial_evidence(&d, &margin, &tie, 0);

        let d = d[0].prior;
        let (mu_trunc, sigma_trunc) = utils::trunc(d.mu(), d.sigma(), margin[0], tie[0]);

        let (delta_div, theta_div_pow2) = if d.sigma() == sigma_trunc {
            (
                d.sigma().powi(2) * mu_trunc - sigma_trunc.powi(2) * d.mu(),
                f64::INFINITY,
            )
        } else {
            (
                (d.sigma().powi(2) * mu_trunc - sigma_trunc.powi(2) * d.mu())
                    / (d.sigma().powi(2) - sigma_trunc.powi(2)),
                (sigma_trunc.powi(2) * d.sigma().powi(2))
                    / (d.sigma().powi(2) - sigma_trunc.powi(2)),
            )
        };

        let mut res = Vec::new();

        for i in 0..t.len() {
            let mut team = Vec::new();

            for j in 0..self.teams[o[i]].len() {
                //
                let mu = if d.sigma() == sigma_trunc {
                    0.0
                } else {
                    self.teams[o[i]][j].prior.mu()
                        + (delta_div - d.mu()) * (-1.0f64).powi(if i == 1 { 1 } else { 0 })
                };

                let sigma_analitico = (theta_div_pow2 + d.sigma().powi(2)
                    - self.teams[o[i]][j].prior.sigma().powi(2))
                .sqrt();

                team.push(Gaussian::new(mu, sigma_analitico));
            }

            res.push(team);
        }

        if o[0] >= o[1] {
            res.swap(0, 1);
        }

        res
    }

    fn likelihood_teams(&mut self) -> Vec<Gaussian> {
        let (o, mut t, mut d, tie, margin) = self.graphical_model();

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

        while ((step.0 > 1e-6) || (step.1 > 1e-6)) && i < 10 {
            step = (0.0, 0.0);

            for e in 0..d.len() - 1 {
                d[e].prior = t[e].posterior_win() - t[e + 1].posterior_lose();

                if i == 0 {
                    let mu = d[e].prior.mu();
                    let sigma = d[e].prior.sigma();

                    if tie[e] {
                        self.evidence *=
                            utils::cdf(margin[e], mu, sigma) - utils::cdf(-margin[e], mu, sigma)
                    } else {
                        self.evidence *= 1.0 - utils::cdf(margin[e], mu, sigma);
                    }
                }

                d[e].likelihood = utils::approx(d[e].prior, margin[e], tie[e]) / d[e].prior;

                let likelihood_lose = t[e].posterior_win() - d[e].likelihood;
                let delta = t[e + 1].likelihood_lose.delta(likelihood_lose);

                step = (
                    if step.0 > delta.0 { step.0 } else { delta.0 },
                    if step.1 > delta.1 { step.1 } else { delta.1 },
                );

                t[e + 1].likelihood_lose = likelihood_lose;
            }

            for e in (1..d.len()).rev() {
                d[e].prior = t[e].posterior_win() - t[e + 1].posterior_lose();

                if i == 0 && e == d.len() - 1 {
                    self.partial_evidence(&d, &margin, &tie, e);
                }

                d[e].likelihood = utils::approx(d[e].prior, margin[e], tie[e]) / d[e].prior;

                let likelihood_win = t[e + 1].posterior_lose() + d[e].likelihood;
                let delta = t[e].likelihood_win.delta(likelihood_win);

                step = (
                    if step.0 > delta.0 { step.0 } else { delta.0 },
                    if step.1 > delta.1 { step.1 } else { delta.1 },
                );

                t[e].likelihood_win = likelihood_win;
            }

            i += 1;
        }

        if d.len() == 1 {
            self.partial_evidence(&d, &margin, &tie, 0);

            d[0].prior = t[0].posterior_win() - t[1].posterior_lose();
            d[0].likelihood = utils::approx(d[0].prior, margin[0], tie[0]) / d[0].prior;
        }

        let t_e = t.len();
        let d_e = d.len();

        t[0].likelihood_win = t[1].posterior_lose() + d[0].likelihood;
        t[t_e - 1].likelihood_lose = t[t_e - 2].posterior_win() - d[d_e - 1].likelihood;

        (0..t.len())
            .map(|e| t[o[e]].likelihood())
            .collect::<Vec<_>>()
    }

    fn compute_likelihoods(&mut self) {
        if self.teams.len() > 2 {
            let m_t_ft = self.likelihood_teams();

            self.likelihoods = (0..self.teams.len())
                .map(|e| {
                    (0..self.teams[e].len())
                        .map(|i| m_t_ft[e] - self.performance(e).exclude(self.teams[e][i].prior))
                        .collect::<Vec<_>>()
                })
                .collect::<Vec<_>>();
        } else {
            self.likelihoods = self.likelihood_analitico();
        }
    }

    pub fn posteriors(&self) -> Vec<Vec<Gaussian>> {
        (0..self.teams.len())
            .map(|e| {
                (0..self.teams[e].len())
                    .map(|i| self.likelihoods[e][i] * self.teams[e][i].prior)
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    }
}

#[cfg(test)]
mod tests {
    use crate::{Gaussian, Player, GAMMA, N_INF};

    use super::*;

    #[test]
    fn test_1vs1() {
        let t_a = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );

        let t_b = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );

        let g = Game::new(vec![vec![t_a], vec![t_b]], vec![0, 1], 0.0);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][0];

        assert_eq!(a.mu(), 20.79477925612302);
        assert_eq!(b.mu(), 29.205220743876975);
        assert_eq!(a.sigma(), 7.194481422570443);

        let t_a = Player::new(Gaussian::new(29.0, 1.0), 25.0 / 6.0, GAMMA, N_INF);
        let t_b = Player::new(Gaussian::new(25.0, 25.0 / 3.0), 25.0 / 6.0, GAMMA, N_INF);

        let g = Game::new(vec![vec![t_a], vec![t_b]], vec![0, 1], 0.0);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][0];

        assert_eq!(a.mu(), 28.896475351225412);
        assert_eq!(a.sigma(), 0.9966043313004235);
        assert_eq!(b.mu(), 32.18921172045737);
        assert_eq!(b.sigma(), 6.062063735879715);

        let t_a = Player::new(Gaussian::new(1.139, 0.531), 1.0, 0.2125, N_INF);
        let t_b = Player::new(Gaussian::new(15.568, 0.51), 1.0, 0.2125, N_INF);

        let g = Game::new(vec![vec![t_a], vec![t_b]], vec![0, 1], 0.0);

        assert_eq!(g.likelihoods[0][0].sigma(), f64::INFINITY);
        assert_eq!(g.likelihoods[1][0].sigma(), f64::INFINITY);
        assert_eq!(g.likelihoods[0][0].mu(), 0.0);
    }

    #[test]
    fn test_1vs1vs1() {
        let t_a = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );
        let t_b = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );
        let t_c = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );

        let g = Game::new(vec![vec![t_a], vec![t_b], vec![t_c]], vec![1, 2, 0], 0.0);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][0];
        let c = p[2][0];

        assert_eq!(a.mu(), 25.00000000000592);
        assert_eq!(a.sigma(), 6.238469796269066);
        assert_eq!(b.mu(), 31.31135822129149);
        assert_eq!(b.sigma(), 6.69881865477675);
        assert_eq!(c.mu(), 18.688641778702593);
        assert_eq!(c.sigma(), 6.698818654778007);

        let t_a = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );
        let t_b = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );
        let t_c = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );

        let g = Game::new(vec![vec![t_a], vec![t_b], vec![t_c]], vec![2, 1, 0], 0.0);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][0];
        let c = p[2][0];

        assert_eq!(a.mu(), 31.31135822129149);
        assert_eq!(a.sigma(), 6.69881865477675);
        assert_eq!(b.mu(), 25.00000000000592);
        assert_eq!(b.sigma(), 6.238469796269066);
        assert_eq!(c.mu(), 18.688641778702593);
        assert_eq!(c.sigma(), 6.698818654778007);

        let t_a = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );
        let t_b = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );
        let t_c = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );

        let g = Game::new(vec![vec![t_a], vec![t_b], vec![t_c]], vec![1, 2, 0], 0.5);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][0];
        let c = p[2][0];

        assert_eq!(a.mu(), 24.999999999511545);
        assert_eq!(a.sigma(), 6.092561128305945);
        assert_eq!(b.mu(), 33.37931495595287);
        assert_eq!(b.sigma(), 6.483575782278924);
        assert_eq!(c.mu(), 16.62068504453558);
        assert_eq!(c.sigma(), 6.483575782198122);
    }

    #[test]
    fn test_1vs1_draw() {
        let t_a = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );

        let t_b = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );

        let g = Game::new(vec![vec![t_a], vec![t_b]], vec![0, 0], 0.25);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][0];

        assert_eq!(a.mu(), 25.0);
        assert_eq!(a.sigma(), 6.469480769842277);
        assert_eq!(b.mu(), 25.0);
        assert_eq!(b.sigma(), 6.469480769842277);

        let t_a = Player::new(Gaussian::new(25.0, 3.0), 25.0 / 6.0, 25.0 / 300.0, N_INF);

        let t_b = Player::new(Gaussian::new(29.0, 2.0), 25.0 / 6.0, 25.0 / 300.0, N_INF);

        let g = Game::new(vec![vec![t_a], vec![t_b]], vec![0, 0], 0.25);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][0];

        assert_eq!(a.mu(), 25.736001810566616);
        assert_eq!(a.sigma(), 2.709956162204711);
        assert_eq!(b.mu(), 28.67288808419261);
        assert_eq!(b.sigma(), 1.9164711604544398);
    }

    #[test]
    fn test_1vs1vs1_draw() {
        let t_a = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );
        let t_b = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );
        let t_c = Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            25.0 / 300.0,
            N_INF,
        );

        let g = Game::new(vec![vec![t_a], vec![t_b], vec![t_c]], vec![0, 0, 0], 0.25);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][0];
        let c = p[2][0];

        assert_eq!(a.mu(), 24.999999999999996);
        assert_eq!(a.sigma(), 5.729068664890827);
        assert_eq!(b.mu(), 25.000000000000004);
        assert_eq!(b.sigma(), 5.707423522433266);
        assert_eq!(c.mu(), 24.999999999999996);
        assert_eq!(c.sigma(), 5.729068664890825);

        let t_a = Player::new(Gaussian::new(25.0, 3.0), 25.0 / 6.0, 25.0 / 300.0, N_INF);
        let t_b = Player::new(Gaussian::new(25.0, 3.0), 25.0 / 6.0, 25.0 / 300.0, N_INF);
        let t_c = Player::new(Gaussian::new(29.0, 2.0), 25.0 / 6.0, 25.0 / 300.0, N_INF);

        let g = Game::new(vec![vec![t_a], vec![t_b], vec![t_c]], vec![0, 0, 0], 0.25);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][0];
        let c = p[2][0];

        assert_eq!(a.mu(), 25.48850755025261);
        assert_eq!(a.sigma(), 2.638208444298423);
        assert_eq!(b.mu(), 25.51067170990121);
        assert_eq!(b.sigma(), 2.6287517663583633);
        assert_eq!(c.mu(), 28.555920328820523);
        assert_eq!(c.sigma(), 1.8856891308577184);
    }
}