#[macro_use]
extern crate log;
extern crate statrs;

#[cfg(test)]
#[macro_use]
extern crate approx;

#[cfg(test)]
extern crate env_logger;

mod factor_graph;
mod gaussian;
mod math;
mod matrix;

use factor_graph::*;
use gaussian::Gaussian;
use matrix::Matrix;

/// Default initial mean of ratings.
pub const MU: f64 = 25.0;

/// Default initial standard deviation of ratings.
pub const SIGMA: f64 = MU / 3.0;

///  Default distance that guarantees about 76% chance of winning.
const BETA: f64 = SIGMA / 2.0;

///  Default dynamic factor.
const TAU: f64 = SIGMA / 100.0;

///  Default draw probability of the game.
const DRAW_PROBABILITY: f64 = 0.10;

///  A basis to check reliability of the result.
const DELTA: f64 = 0.0001;

pub trait Rateable {
    fn mu(&self) -> f64;
    fn sigma(&self) -> f64;
}

#[derive(Debug, PartialEq)]
pub struct Rating {
    pub mu: f64,
    pub sigma: f64,
}

impl Rating {
    pub fn new(mu: f64, sigma: f64) -> Rating {
        Rating { mu, sigma }
    }
}

impl Rateable for Rating {
    fn mu(&self) -> f64 {
        self.mu
    }

    fn sigma(&self) -> f64 {
        self.sigma
    }
}

impl Rateable for Gaussian {
    fn mu(&self) -> f64 {
        self.mu()
    }

    fn sigma(&self) -> f64 {
        self.sigma()
    }
}

fn draw_margin(p: f64, beta: f64, total_players: f64) -> f64 {
    math::icdf((p + 1.0) / 2.0) * total_players.sqrt() * beta
}

pub fn rate<R>(ratings: &[(R, u16)], ranks: &[u16]) -> Vec<Rating>
where
    R: Rateable,
{
    // TODO Validate rating_groups is orderded in teams.
    // TODO Validate that teams are orderd after rank.

    let tau_sqr = TAU.powi(2);
    let beta_sqr = BETA.powi(2);

    let mut variable_arena = VariableArena::new();

    let rating_count = ratings.len();
    let team_count = ranks.len();

    let rating_vars = (0..rating_count)
        .map(|_| variable_arena.create())
        .collect::<Vec<_>>();
    let perf_vars = ratings
        .iter()
        .map(|(_, team)| (variable_arena.create(), *team))
        .collect::<Vec<_>>();

    let team_perf_vars = (0..team_count)
        .map(|_| variable_arena.create())
        .collect::<Vec<_>>();
    let team_diff_vars = (0..team_count - 1)
        .map(|_| variable_arena.create())
        .collect::<Vec<_>>();

    let mut factor_id = 0;

    let rating_layer = rating_vars
        .iter()
        .zip(ratings.iter().map(|(rating, _)| rating))
        .map(|(rating_var, rating)| {
            let gaussian =
                Gaussian::from_mu_sigma(rating.mu(), (rating.sigma().powi(2) + tau_sqr).sqrt());

            factor_id += 1;

            PriorFactor::new(&mut variable_arena, factor_id, *rating_var, gaussian)
        })
        .collect::<Vec<_>>();

    let perf_layer = rating_vars
        .iter()
        .zip(perf_vars.iter().map(|(variable, _)| variable))
        .map(|(rating_var, perf)| {
            factor_id += 1;

            LikelihoodFactor::new(&mut variable_arena, factor_id, *rating_var, *perf, beta_sqr)
        })
        .collect::<Vec<_>>();

    let team_perf_layer = team_perf_vars
        .iter()
        .enumerate()
        .map(|(i, variable)| {
            factor_id += 1;

            let team = perf_vars
                .iter()
                .filter(|(_, team)| *team as usize == i)
                .map(|(variable, _)| *variable)
                .collect::<Vec<_>>();

            let team_count = team.len();

            SumFactor::new(
                &mut variable_arena,
                factor_id,
                *variable,
                team,
                vec![1.0; team_count],
            )
        })
        .collect::<Vec<_>>();

    let team_diff_layer = team_diff_vars
        .iter()
        .zip(team_perf_vars.windows(2))
        .map(|(variable, teams)| {
            factor_id += 1;

            SumFactor::new(
                &mut variable_arena,
                factor_id,
                *variable,
                teams.to_vec(),
                vec![1.0, -1.0],
            )
        })
        .collect::<Vec<_>>();

    let trunc_layer = team_diff_vars
        .iter()
        .enumerate()
        .map(|(i, variable)| {
            factor_id += 1;

            let player_count = perf_vars
                .iter()
                .filter(|(_, team)| *team as usize == i || *team as usize == i + 1)
                .count();

            TruncateFactor::new(
                &mut variable_arena,
                factor_id,
                *variable,
                draw_margin(DRAW_PROBABILITY, BETA, player_count as f64),
                ranks[i] == ranks[i + 1],
            )
        })
        .collect::<Vec<_>>();;

    for factor in &rating_layer {
        factor.start(&mut variable_arena);
    }

    for factor in &perf_layer {
        factor.update_value(&mut variable_arena);
    }

    for factor in &team_perf_layer {
        factor.update_sum(&mut variable_arena);
    }

    for _ in 0..5 {
        for factor in &team_diff_layer {
            factor.update_sum(&mut variable_arena);
        }

        for factor in &trunc_layer {
            factor.update(&mut variable_arena);
        }

        for factor in &team_diff_layer {
            factor.update_term(&mut variable_arena, 0);
            factor.update_term(&mut variable_arena, 1);
        }
    }

    for factor in &team_perf_layer {
        factor.update_all_terms(&mut variable_arena);
    }

    for factor in &perf_layer {
        factor.update_mean(&mut variable_arena);
    }

    rating_vars
        .iter()
        .map(|variable| variable_arena.get(*variable).unwrap().get_value())
        .map(|value| Rating {
            mu: value.mu(),
            sigma: value.sigma(),
        })
        .collect::<Vec<_>>()
}

pub fn quality<R>(rating_groups: &[&[R]]) -> f64
where
    R: Rateable,
{
    let flatten_ratings = rating_groups
        .iter()
        .flat_map(|group| group.iter())
        .collect::<Vec<_>>();

    let flatten_weights = vec![1.0; flatten_ratings.len()].into_boxed_slice();

    let length = flatten_ratings.len();

    let mut mean_matrix = Matrix::new(length, 1);

    for (i, rating) in flatten_ratings.iter().enumerate() {
        mean_matrix[(i, 0)] = rating.mu();
    }

    let mut variance_matrix = Matrix::new(length, length);

    for (i, rating) in flatten_ratings.iter().enumerate() {
        variance_matrix[(i, i)] = rating.sigma().powi(2);
    }

    let mut rotated_a_matrix = Matrix::new(rating_groups.len() - 1, length);

    let mut t = 0;
    let mut x = 0;

    for (row, group) in rating_groups.windows(2).enumerate() {
        let current = group[0];
        let next = group[1];

        for n in t..t + current.len() {
            rotated_a_matrix[(row, n)] = flatten_weights[n];
            t += 1;
            x += 1;
        }

        for n in x..x + next.len() {
            rotated_a_matrix[(row, n)] = -flatten_weights[n];
            x += 1;
        }
    }

    let a_matrix = rotated_a_matrix.transpose();

    let _ata = BETA.powi(2) * &rotated_a_matrix * &a_matrix;
    let _atsa = &rotated_a_matrix * &variance_matrix * &a_matrix;

    let start = mean_matrix.transpose() * &a_matrix;
    let middle = &_ata + &_atsa;
    let end = &rotated_a_matrix * &mean_matrix;

    let e_arg = (-0.5 * &start * &middle.inverse() * &end).determinant();
    let s_arg = _ata.determinant() / middle.determinant();

    e_arg.exp() * s_arg.sqrt()
}

#[cfg(test)]
mod tests {
    use approx::{AbsDiffEq, RelativeEq};
    use env_logger;

    use super::*;

    const EPSILON: f64 = 2e-14;

    impl AbsDiffEq for Rating {
        type Epsilon = f64;

        fn default_epsilon() -> Self::Epsilon {
            f64::default_epsilon()
        }

        fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
            self.mu.abs_diff_eq(&other.mu, epsilon) && self.sigma.abs_diff_eq(&other.sigma, epsilon)
        }
    }

    impl RelativeEq for Rating {
        fn default_max_relative() -> Self::Epsilon {
            f64::default_max_relative()
        }

        fn relative_eq(
            &self,
            other: &Self,
            epsilon: Self::Epsilon,
            max_relative: Self::Epsilon,
        ) -> bool {
            self.mu.relative_eq(&other.mu, epsilon, max_relative)
                && self.sigma.relative_eq(&other.sigma, epsilon, max_relative)
        }
    }

    #[test]
    fn test_quality_1vs1() {
        let alice = Rating::new(MU, SIGMA);
        let bob = Rating::new(MU, SIGMA);

        assert_relative_eq!(
            quality(&[&[alice], &[bob]]),
            0.4472135954999579,
            epsilon = EPSILON
        );
    }

    #[test]
    fn test_rate_4_free_for_all() {
        let alice = Rating::new(MU, SIGMA);
        let bob = Rating::new(MU, SIGMA);
        let chris = Rating::new(MU, SIGMA);
        let darren = Rating::new(MU, SIGMA);

        let expected_ratings = vec![
            Rating::new(33.20778932559525, 6.347937214998893),
            Rating::new(27.401497882797486, 5.787057812482782),
            Rating::new(22.598576351652632, 5.7871159419307645),
            Rating::new(16.79337409436942, 6.348053083319977),
        ];

        let ratings = rate(
            &[(alice, 0), (bob, 1), (chris, 2), (darren, 3)],
            &[0, 1, 2, 3],
        );

        for (rating, expected) in ratings.iter().zip(expected_ratings.iter()) {
            assert_relative_eq!(rating, expected, epsilon = EPSILON);
        }
    }

    #[test]
    fn test_rate_1vs1_draw() {
        let alice = Rating::new(MU, SIGMA);
        let bob = Rating::new(MU, SIGMA);

        let expected_ratings = vec![
            Rating::new(25.0, 6.457515683245051),
            Rating::new(25.0, 6.457515683245051),
        ];

        let ratings = rate(&[(alice, 0), (bob, 1)], &[0, 0]);

        for (rating, expected) in ratings.iter().zip(expected_ratings.iter()) {
            assert_relative_eq!(rating, expected, epsilon = EPSILON);
        }
    }

    #[test]
    fn test_rate_2vs2() {
        let _ = env_logger::try_init();

        let alice = Rating::new(MU, SIGMA);
        let bob = Rating::new(MU, SIGMA);
        let chris = Rating::new(MU, SIGMA);
        let darren = Rating::new(MU, SIGMA);

        let expected_ratings = vec![
            Rating::new(28.108322399069035, 7.77436345109384),
            Rating::new(28.108322399069035, 7.77436345109384),
            Rating::new(21.891677600930958, 7.77436345109384),
            Rating::new(21.891677600930958, 7.77436345109384),
        ];

        let ratings = rate(&[(alice, 0), (bob, 0), (chris, 1), (darren, 1)], &[0, 1]);

        for (rating, expected) in ratings.iter().zip(expected_ratings.iter()) {
            assert_relative_eq!(rating, expected, epsilon = EPSILON);
        }
    }
}