use std::collections::HashMap;

#[derive(Debug)]
pub(crate) struct ScoredVarId {
    pub(crate) variation_id: u32,
    pub(crate) score: f64,
}

#[derive(Debug)]
#[cfg_attr(test, derive(PartialEq))]
pub(crate) struct VarGroup {
    pub(crate) ids: Vec<u32>,
}

impl VarGroup {
    pub fn new(mut ids: Vec<u32>) -> Self {
        ids.sort();
        return VarGroup { ids };
    }
}

#[derive(Debug)]
#[cfg_attr(test, derive(PartialEq))]
pub(crate) struct CBScores {
    // keys = experiment IDs / values = list of variation IDs ordered descending
    // by score (there may be several variation ids with same score)
    values: HashMap<u32, Vec<VarGroup>>,
}

impl CBScores {
    pub fn new(cbs_map: &HashMap<u32, Vec<ScoredVarId>>) -> Self {
        return CBScores {
            values: cbs_map.iter().map(|(exp_id, scores)| (*exp_id, Self::extract_var_ids(scores))).collect(),
        };
    }

    fn extract_var_ids(scores: &[ScoredVarId]) -> Vec<VarGroup> {
        let mut sorted: Vec<&ScoredVarId> = scores.iter().collect();
        sorted.sort_by(|a, b| b.score.total_cmp(&a.score)); // highest score first (total order; handles NaN)

        sorted
            .chunk_by(|a, b| a.score.total_cmp(&b.score).is_eq())
            .map(|chunk| {
                let ids = chunk.iter().map(|s| s.variation_id).collect();
                VarGroup::new(ids)
            })
            .collect()
    }

    pub fn values(&self) -> &HashMap<u32, Vec<VarGroup>> {
        return &self.values;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rstest::rstest;

    #[rstest]
    #[case(
        HashMap::from([
            (1, vec![
                ScoredVarId { variation_id: 104, score: 5.0 },
                ScoredVarId { variation_id: 103, score: 3.5 },
                ScoredVarId { variation_id: 102, score: 4.0 },
                ScoredVarId { variation_id: 101, score: 3.5 },
            ])
        ]),
        HashMap::from([
            (1, vec![
                VarGroup::new(vec![104]),
                VarGroup::new(vec![102]),
                VarGroup::new(vec![101, 103]),
            ])
        ])
    )]
    #[case(
        HashMap::from([
            (2, vec![
                ScoredVarId { variation_id: 201, score: 2.0 },
                ScoredVarId { variation_id: 203, score: 2.0 },
                ScoredVarId { variation_id: 202, score: 2.0 },
            ])
        ]),
        HashMap::from([(2, vec![VarGroup::new(vec![201, 202, 203])])])
    )]
    #[case(
        HashMap::from([
            (1, vec![
                ScoredVarId { variation_id: 101, score: 3.0 },
                ScoredVarId { variation_id: 102, score: 4.0 },
            ]),
            (2, vec![
                ScoredVarId { variation_id: 201, score: 5.0 },
                ScoredVarId { variation_id: 202, score: 5.0 },
            ])
        ]),
        HashMap::from([
            (1, vec![
                VarGroup::new(vec![102]),
                VarGroup::new(vec![101]),
            ]),
            (2, vec![VarGroup::new(vec![201, 202])])
        ])
    )]
    fn test_new(#[case] cbs_map: HashMap<u32, Vec<ScoredVarId>>, #[case] expected: HashMap<u32, Vec<VarGroup>>) {
        let cbs = CBScores::new(&cbs_map);
        assert_eq!(expected.len(), cbs.values().len());
        for (expected_key, expected_var_groups) in expected.iter() {
            let actual_var_groups = cbs.values().get(expected_key).unwrap();
            assert_eq!(expected_var_groups, actual_var_groups);
        }
    }
}
