use anyhow::{Context, Result};
use statrs::function::factorial::binomial;
use std::collections::{BTreeSet, HashMap};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Coalition {
    members: BTreeSet<u64>,
}

impl Coalition {
    pub fn new(members: Vec<u64>) -> Self {
        Coalition {
            members: members.into_iter().collect(),
        }
    }

    fn size(&self) -> usize {
        self.members.len()
    }

    fn contains(&self, player: u64) -> bool {
        self.members.contains(&player)
    }

    fn subtract(&self, player: u64) -> Self {
        let mut new_members = self.members.clone();
        new_members.remove(&player);
        Coalition {
            members: new_members,
        }
    }
}

#[derive(Debug)]
pub struct Shapley {
    coalition_worth: HashMap<Coalition, f64>,
    coalition_size_weights: HashMap<usize, f64>,
}

impl Shapley {
    pub fn new(players: Vec<u64>, mut coalition_worth: HashMap<Coalition, f64>) -> Self {
        let n = players.len();
        // Ensure empty coalition exists (value = 0.0)
        coalition_worth.entry(Coalition::new(vec![])).or_insert(0.0);

        // Precompute combinatorial weights
        let mut weights = HashMap::new();
        for s in 0..n {
            // s = coalition size without player i
            let binom = binomial((n - 1) as u64, s as u64);
            weights.insert(s, (1.0 / n as f64) / binom);
        }

        Shapley {
            coalition_worth,
            coalition_size_weights: weights,
        }
    }

    pub fn shapley_value(&self, player: u64) -> Result<f64> {
        let mut total_contribution = 0.0;
        let mut total_weight = 0.0;

        for (coalition, &value_with) in &self.coalition_worth {
            // Consider only coalitions containing player
            if coalition.contains(player) {
                let without = coalition.subtract(player);
                let value_without = self
                    .coalition_worth
                    .get(&without)
                    .context(format!("Missing value for coalition {:?}", without))?;

                let s = without.size();
                let weight = self
                    .coalition_size_weights
                    .get(&s)
                    .context("Missing combinatorial weight")?;

                total_contribution += weight * (value_with - value_without);
                total_weight += weight;
            }
        }

        if total_weight.abs() < f64::EPSILON {
            anyhow::bail!(
                "Total weight is zero for player {} (insufficient data)",
                player
            );
        }

        Ok(total_contribution / total_weight)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use maplit::hashmap;
    use statrs::assert_almost_eq;

    #[test]
    fn test_empty_coalition_handling() {
        let coalition_worth = hashmap! {
            Coalition::new(vec![1]) => 5.0,
        };
        let shapley = Shapley::new(vec![1], coalition_worth);
        // Should contain empty coalition automatically
        assert_eq!(
            shapley.coalition_worth.get(&Coalition::new(vec![])),
            Some(&0.0)
        );
    }

    #[test]
    fn test_missing_data() {
        let coalition_worth = hashmap! {
            Coalition::new(vec![1, 2]) => 10.0,
        };
        let shapley = Shapley::new(vec![1, 2], coalition_worth);
        assert!(shapley.shapley_value(1).is_err());
    }

    #[test]
    fn test_simple_shapley_value() {
        let coalition_worth = hashmap! {
            Coalition::new(vec![]) => 0.0,
            Coalition::new(vec![1]) => 10.0,
            Coalition::new(vec![2]) => 20.0,
            Coalition::new(vec![1, 2]) => 30.0,
        };

        let shapley = Shapley::new(vec![1, 2], coalition_worth);
        assert_eq!(shapley.shapley_value(1).unwrap(), 10.0);
        assert_eq!(shapley.shapley_value(2).unwrap(), 20.0);
    }

    // https://gtl.csa.iisc.ac.in/gametheory/ln/web-cp5-shapley.pdf
    // 2.1 Example 1: Divide the Dollar Game
    #[test]
    fn test_divide_dollar_game() {
        let coalition_worth = hashmap! {
            Coalition::new(vec![1]) => 0.0,
            Coalition::new(vec![2]) => 0.0,
            Coalition::new(vec![3]) => 0.0,
            Coalition::new(vec![2, 3]) => 0.0,

            Coalition::new(vec![1, 2]) => 300.0,
            Coalition::new(vec![1, 3]) => 300.0,
            Coalition::new(vec![1, 2, 3]) => 300.0,
        };

        let shapley = Shapley::new(vec![1, 2, 3], coalition_worth);
        assert_almost_eq!(shapley.shapley_value(1).unwrap(), 200.0, 1e-10);
        assert_almost_eq!(shapley.shapley_value(2).unwrap(), 50.0, 1e-10);
        assert_almost_eq!(shapley.shapley_value(3).unwrap(), 50.0, 1e-10);
    }

    // https://gtl.csa.iisc.ac.in/gametheory/ln/web-cp5-shapley.pdf
    // 2.4 Example 4: A Logistics Game(the doc maybe give wrong result)
    #[test]
    fn test_logistics_game() {
        let coalition_worth = hashmap! {
            Coalition::new(vec![1]) => 0.0,
            Coalition::new(vec![2]) => 0.0,
            Coalition::new(vec![3]) => 0.0,
            Coalition::new(vec![4]) => 0.0,

            Coalition::new(vec![1, 2]) => 0.0,
            Coalition::new(vec![1, 3]) => 0.0,
            Coalition::new(vec![1, 4]) => 0.0,
            Coalition::new(vec![2, 3]) => 0.0,
            Coalition::new(vec![2, 4]) => 0.0,
            Coalition::new(vec![3, 4]) => 0.0,

            Coalition::new(vec![1, 2, 3]) => 0.0,
            Coalition::new(vec![2, 3, 4]) => 0.0,

            Coalition::new(vec![1, 2, 4]) => 45.0,
            Coalition::new(vec![1, 3, 4]) => 40.0,
            Coalition::new(vec![1, 2, 3, 4]) => 65.0,
        };

        let shapley = Shapley::new(vec![1, 2, 3, 4], coalition_worth);
        assert_almost_eq!(shapley.shapley_value(1).unwrap(), 23.333333333333332, 1e-10);
        assert_almost_eq!(shapley.shapley_value(2).unwrap(), 10.0, 1e-10);
        assert_almost_eq!(shapley.shapley_value(3).unwrap(), 8.333333333333332, 1e-10);
        assert_almost_eq!(shapley.shapley_value(4).unwrap(), 23.333333333333332, 1e-10);
    }
}
