use crate::{
    approx, compute_margin, evidence,
    gaussian::Gaussian,
    message::{DiffMessage, TeamMessage},
    player::Player,
    sort_perm, tuple_gt, tuple_max, N00, N_INF,
};

pub struct Game {
    teams: Vec<Vec<Player>>,
    result: Vec<f64>,
    weights: Vec<Vec<f64>>,
    p_draw: f64,
    pub(crate) likelihoods: Vec<Vec<Gaussian>>,
    pub(crate) evidence: f64,
}

impl Game {
    pub fn new(
        teams: Vec<Vec<Player>>,
        mut result: Vec<f64>,
        mut weights: Vec<Vec<f64>>,
        p_draw: f64,
    ) -> Self {
        assert!(
            (result.is_empty() || result.len() == teams.len()),
            "result must be empty or the same length as teams"
        );

        assert!(
            (weights.is_empty() || weights.len() == teams.len()),
            "weights must be empty or the same length as teams"
        );

        assert!(
            weights.is_empty()
                || weights
                    .iter()
                    .zip(teams.iter())
                    .all(|(w, t)| w.len() == t.len()),
            "weights must be empty or has the same dimensions as teams"
        );

        assert!(
            (0.0..1.0).contains(&p_draw),
            "draw probability.must be >= 0.0 and < 1.0"
        );

        assert!(
            p_draw > 0.0 || {
                let mut r = result.clone();
                r.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
                r.windows(2).all(|w| w[0] != w[1])
            },
            "draw must be > 0.0 if there is teams with draw"
        );

        if result.is_empty() {
            result = (0..teams.len()).rev().map(|i| i as f64).collect::<Vec<_>>();
        }

        if weights.is_empty() {
            weights = teams
                .iter()
                .map(|team| vec![1.0; team.len()])
                .collect::<Vec<_>>();
        }

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

        this.likelihoods();

        this
    }

    fn likelihoods(&mut self) -> &Vec<Vec<Gaussian>> {
        let m_t_ft = self.likelihood_teams();

        self.likelihoods = self
            .teams
            .iter()
            .zip(self.weights.iter())
            .zip(m_t_ft)
            .map(|((p, w), m)| {
                let performance = p.iter().zip(w.iter()).fold(N00, |p, (player, &weight)| {
                    p + (player.performance() * weight)
                });

                p.iter()
                    .zip(w.iter())
                    .map(|(p, &w)| {
                        ((m - performance.exclude(p.performance() * w)) * (1.0 / w))
                            .forget(p.beta, 1)
                    })
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>();

        &self.likelihoods
    }

    fn likelihood_teams(&mut self) -> Vec<Gaussian> {
        let o = sort_perm(&self.result, true);

        let mut t = o
            .iter()
            .map(|&e| {
                let performance = self.teams[e]
                    .iter()
                    .zip(self.weights[e].iter())
                    .fold(N00, |p, (player, &weight)| {
                        p + (player.performance() * weight)
                    });

                TeamMessage {
                    prior: performance,
                    likelihood_lose: N_INF,
                    likelihood_win: N_INF,
                    likelihood_draw: N_INF,
                }
            })
            .collect::<Vec<_>>();

        let mut d = t
            .windows(2)
            .map(|w| DiffMessage {
                prior: w[0].prior - w[1].prior,
                likelihood: N_INF,
            })
            .collect::<Vec<_>>();

        let tie = o
            .windows(2)
            .map(|e| self.result[e[0]] == self.result[e[1]])
            .collect::<Vec<_>>();

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

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

        self.evidence = 1.0;

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

        while tuple_gt(step, 1e-6) && iter < 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 iter == 0 {
                    self.evidence *= evidence(&d, &margin, &tie, e);
                }

                d[e].likelihood = approx(d[e].prior, margin[e], tie[e]) / d[e].prior;
                let likelihood_lose = t[e].posterior_win() - d[e].likelihood;
                step = tuple_max(step, t[e + 1].likelihood_lose.delta(likelihood_lose));
                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 iter == 0 && e == d.len() - 1 {
                    self.evidence *= evidence(&d, &margin, &tie, e);
                }

                d[e].likelihood = approx(d[e].prior, margin[e], tie[e]) / d[e].prior;
                let likelihood_win = t[e + 1].posterior_lose() + d[e].likelihood;
                step = tuple_max(step, t[e].likelihood_win.delta(likelihood_win));
                t[e].likelihood_win = likelihood_win;
            }

            iter += 1;
        }

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

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

        let t_end = t.len() - 1;
        let d_end = d.len() - 1;

        t[0].likelihood_win = t[1].posterior_lose() + d[0].likelihood;
        t[t_end].likelihood_lose = t[t_end - 1].posterior_win() - d[d_end].likelihood;

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

    pub fn posteriors(&self) -> Vec<Vec<Gaussian>> {
        self.likelihoods
            .iter()
            .zip(self.teams.iter())
            .map(|(l, t)| {
                l.iter()
                    .zip(t.iter())
                    .map(|(&l, p)| l * p.prior)
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    }
}

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

    use crate::{Gaussian, Player, GAMMA};

    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);
        let t_b = Player::new(Gaussian::new(25.0, 25.0 / 3.0), 25.0 / 6.0, 25.0 / 300.0);

        let g = Game::new(vec![vec![t_a], vec![t_b]], vec![0.0, 1.0], vec![], 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);
        let t_b = Player::new(Gaussian::new(25.0, 25.0 / 3.0), 25.0 / 6.0, GAMMA);

        let g = Game::new(vec![vec![t_a], vec![t_b]], vec![0.0, 1.0], vec![], 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);
        let t_b = Player::new(Gaussian::new(15.568, 0.51), 1.0, 0.2125);

        let g = Game::new(vec![vec![t_a], vec![t_b]], vec![0.0, 1.0], vec![], 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 teams = vec![
            vec![Player::new(
                Gaussian::new(25.0, 25.0 / 3.0),
                25.0 / 6.0,
                25.0 / 300.0,
            )],
            vec![Player::new(
                Gaussian::new(25.0, 25.0 / 3.0),
                25.0 / 6.0,
                25.0 / 300.0,
            )],
            vec![Player::new(
                Gaussian::new(25.0, 25.0 / 3.0),
                25.0 / 6.0,
                25.0 / 300.0,
            )],
        ];

        let g = Game::new(teams.clone(), vec![1.0, 2.0, 0.0], vec![], 0.0);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][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);

        let g = Game::new(teams.clone(), vec![], vec![], 0.0);
        let p = g.posteriors();

        let a = p[0][0];
        let b = p[1][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);

        let g = Game::new(teams, vec![1.0, 2.0, 0.0], vec![], 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);
        let t_b = Player::new(Gaussian::new(25.0, 25.0 / 3.0), 25.0 / 6.0, 25.0 / 300.0);

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

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

        assert_eq!(a.mu, 24.999999999999996);
        assert_eq!(a.sigma, 6.469480769842277);
        assert_eq!(b.mu, 24.999999999999996);
        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);
        let t_b = Player::new(Gaussian::new(29.0, 2.0), 25.0 / 6.0, 25.0 / 300.0);

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

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

        assert_eq!(a.mu, 25.73600181056662);
        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);
        let t_b = Player::new(Gaussian::new(25.0, 25.0 / 3.0), 25.0 / 6.0, 25.0 / 300.0);
        let t_c = Player::new(Gaussian::new(25.0, 25.0 / 3.0), 25.0 / 6.0, 25.0 / 300.0);

        let g = Game::new(
            vec![vec![t_a], vec![t_b], vec![t_c]],
            vec![0.0, 0.0, 0.0],
            vec![],
            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);
        let t_b = Player::new(Gaussian::new(25.0, 3.0), 25.0 / 6.0, 25.0 / 300.0);
        let t_c = Player::new(Gaussian::new(29.0, 2.0), 25.0 / 6.0, 25.0 / 300.0);

        let g = Game::new(
            vec![vec![t_a], vec![t_b], vec![t_c]],
            vec![0.0, 0.0, 0.0],
            vec![],
            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.6382084442984226);
        assert_eq!(b.mu, 25.510671709901217);
        assert_eq!(b.sigma, 2.6287517663583633);
        assert_eq!(c.mu, 28.555920328820527);
        assert_eq!(c.sigma, 1.8856891308577184);
    }

    #[test]
    fn test_1vs1_weighted() {
        let w_a = vec![1.0];
        let w_b = vec![2.0];

        let t_a = vec![Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            0.0,
        )];
        let t_b = vec![Player::new(
            Gaussian::new(25.0, 25.0 / 3.0),
            25.0 / 6.0,
            0.0,
        )];

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

        assert_ulps_eq!(
            p[0][0],
            Gaussian::new(30.625173, 7.765472),
            epsilon = 0.000001
        );
        assert_ulps_eq!(
            p[1][0],
            Gaussian::new(13.749653, 5.733840),
            epsilon = 0.000001
        );

        let w_a = vec![1.0];
        let w_b = vec![0.7];

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

        assert_ulps_eq!(
            p[0][0],
            Gaussian::new(27.630080, 7.206676),
            epsilon = 0.000001
        );
        assert_ulps_eq!(
            p[1][0],
            Gaussian::new(23.158943, 7.801628),
            epsilon = 0.000001
        );

        let w_a = vec![1.6];
        let w_b = vec![0.7];

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

        assert_ulps_eq!(
            p[0][0],
            Gaussian::new(26.142438, 7.573088),
            epsilon = 0.000001
        );
        assert_ulps_eq!(
            p[1][0],
            Gaussian::new(24.500183, 8.193278),
            epsilon = 0.000001
        );

        let w_a = vec![1.0];
        let w_b = vec![0.0];

        let t_a = vec![Player::new(Gaussian::new(2.0, 6.0), 1.0, 0.0)];
        let t_b = vec![Player::new(Gaussian::new(2.0, 6.0), 1.0, 0.0)];

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

        assert_ulps_eq!(
            p[0][0],
            Gaussian::new(5.55706798109105, 4.0528268357577995),
            epsilon = 0.000001
        );
        assert_ulps_eq!(p[1][0], Gaussian::new(2.0, 6.0), epsilon = 0.000001);

        let w_a = vec![1.0];
        let w_b = vec![-1.0];

        let t_a = vec![Player::new(Gaussian::new(2.0, 6.0), 1.0, 0.0)];
        let t_b = vec![Player::new(Gaussian::new(2.0, 6.0), 1.0, 0.0)];

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

        assert_ulps_eq!(p[0][0], p[1][0], epsilon = 0.000001);
    }
}