use types::{
    AttesterMap,
    ProposerMap,
    ShardAndCommittee,
};

#[derive(Debug, PartialEq)]
pub enum AttesterAndProposerMapError {
    NoShardAndCommitteeForSlot,
    NoAvailableProposer,
}

/// Generate a map of `(slot, shard) |--> committee`.
///
/// The attester map is used to optimise the lookup of a committee.
pub fn generate_attester_and_proposer_maps(
    shard_and_committee_for_slots: &Vec<Vec<ShardAndCommittee>>,
    start_slot: u64)
    -> Result<(AttesterMap, ProposerMap), AttesterAndProposerMapError>
{
    let mut attester_map = AttesterMap::new();
    let mut proposer_map = ProposerMap::new();
    for (i, slot) in shard_and_committee_for_slots.iter().enumerate() {
        /*
         * Store the proposer for the block.
         */
        let slot_number = (i as u64).saturating_add(start_slot);
        let first_committee = &slot.get(0)
            .ok_or(AttesterAndProposerMapError::NoShardAndCommitteeForSlot)?
            .committee;
        let proposer_index = (slot_number as usize).checked_rem(first_committee.len())
            .ok_or(AttesterAndProposerMapError::NoAvailableProposer)?;
        proposer_map.insert(slot_number, first_committee[proposer_index]);

        /*
         * Loop through the shards and extend the attester map.
         */
        for shard_and_committee in slot {
            let committee = shard_and_committee.committee.clone();
            attester_map.insert((slot_number, shard_and_committee.shard), committee);
        }
    };
    Ok((attester_map, proposer_map))
}

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

    fn sac_generator(shard_count: u16,
                     slot_count: usize,
                     sac_per_slot: usize,
                     committee_size: usize)
        -> Vec<Vec<ShardAndCommittee>>
    {
        let mut shard = 0;
        let mut validator = 0;
        let mut cycle = vec![];

        for _ in 0..slot_count {
            let mut slot: Vec<ShardAndCommittee> = vec![];
            for _ in 0..sac_per_slot {
                let mut sac = ShardAndCommittee {
                    shard: shard % shard_count,
                    committee: vec![],
                };
                for _ in 0..committee_size {
                    sac.committee.push(validator);
                    validator += 1;
                }
                slot.push(sac);
                shard += 1;
            }
            cycle.push(slot);
        }
        cycle
    }

    #[test]
    fn test_attester_proposer_maps_empty_slots() {
        let sac = sac_generator(4, 4, 0, 1);
        let result = generate_attester_and_proposer_maps(&sac, 0);
        assert_eq!(result, Err(AttesterAndProposerMapError::NoShardAndCommitteeForSlot));
    }

    #[test]
    fn test_attester_proposer_maps_empty_committees() {
        let sac = sac_generator(4, 4, 1, 0);
        let result = generate_attester_and_proposer_maps(&sac, 0);
        assert_eq!(result, Err(AttesterAndProposerMapError::NoAvailableProposer));
    }

    #[test]
    fn test_attester_proposer_maps_scenario_a() {
        let sac = sac_generator(4, 4, 1, 1);
        let (a, p) = generate_attester_and_proposer_maps(&sac, 0).unwrap();

        assert_eq!(*p.get(&0).unwrap(), 0);
        assert_eq!(*p.get(&1).unwrap(), 1);
        assert_eq!(*p.get(&2).unwrap(), 2);
        assert_eq!(*p.get(&3).unwrap(), 3);

        assert_eq!(*a.get(&(0, 0)).unwrap(), vec![0]);
        assert_eq!(*a.get(&(1, 1)).unwrap(), vec![1]);
        assert_eq!(*a.get(&(2, 2)).unwrap(), vec![2]);
        assert_eq!(*a.get(&(3, 3)).unwrap(), vec![3]);
    }

    #[test]
    fn test_attester_proposer_maps_scenario_b() {
        let sac = sac_generator(4, 4, 1, 4);
        let (a, p) = generate_attester_and_proposer_maps(&sac, 0).unwrap();

        assert_eq!(*p.get(&0).unwrap(), 0);
        assert_eq!(*p.get(&1).unwrap(), 5);
        assert_eq!(*p.get(&2).unwrap(), 10);
        assert_eq!(*p.get(&3).unwrap(), 15);

        assert_eq!(*a.get(&(0, 0)).unwrap(), vec![0, 1, 2, 3]);
        assert_eq!(*a.get(&(1, 1)).unwrap(), vec![4, 5, 6, 7]);
        assert_eq!(*a.get(&(2, 2)).unwrap(), vec![8, 9, 10, 11]);
        assert_eq!(*a.get(&(3, 3)).unwrap(), vec![12, 13, 14, 15]);
    }
}