use attesters::Attesters;
use errors::EpochProcessingError as Error;
use fnv::FnvHashSet;
use integer_sqrt::IntegerSquareRoot;
use rayon::prelude::*;
use ssz::TreeHash;
use std::collections::HashMap;
use std::iter::FromIterator;
use types::{validator_registry::get_active_validator_indices, *};
use winning_root::{winning_root, WinningRoot};

pub mod attester_sets;
pub mod attesters;
pub mod errors;
pub mod inclusion_distance;
pub mod tests;
pub mod winning_root;

/// Maps a shard to a winning root.
///
/// It is generated during crosslink processing and later used to reward/penalize validators.
pub type WinningRootHashSet = HashMap<u64, WinningRoot>;

/// Performs per-epoch processing on some BeaconState.
///
/// Mutates the given `BeaconState`, returning early if an error is encountered. If an error is
/// returned, a state might be "half-processed" and therefore in an invalid state.
///
/// Spec v0.4.0
pub fn per_epoch_processing(state: &mut BeaconState, spec: &ChainSpec) -> Result<(), Error> {
    let previous_epoch = state.previous_epoch(spec);

    // Ensure all of the caches are built.
    state.build_epoch_cache(RelativeEpoch::Previous, spec)?;
    state.build_epoch_cache(RelativeEpoch::Current, spec)?;
    state.build_epoch_cache(RelativeEpoch::Next, spec)?;

    let active_validator_indices = calculate_active_validator_indices(&state, spec);

    let current_total_balance = state.get_total_balance(&active_validator_indices[..], spec);

    let previous_total_balance = state.get_total_balance(
        &get_active_validator_indices(&state.validator_registry, previous_epoch)[..],
        spec,
    );

    let attesters = calculate_attester_sets(&state, &active_validator_indices, spec)?;

    process_eth1_data(state, spec);

    process_justification(
        state,
        current_total_balance,
        previous_total_balance,
        attesters.balances.previous_epoch_boundary,
        attesters.balances.current_epoch_boundary,
        spec,
    );

    // Crosslinks
    let winning_root_for_shards = process_crosslinks(state, spec)?;

    // Rewards and Penalities
    process_rewards_and_penalities(
        state,
        &attesters,
        previous_total_balance,
        &winning_root_for_shards,
        spec,
    )?;

    // Ejections
    state.process_ejections(spec);

    // Validator Registry
    process_validator_registry(state, spec)?;

    // Final updates
    update_active_tree_index_roots(state, spec)?;
    update_latest_slashed_balances(state, spec);
    clean_attestations(state, spec);

    // Rotate the epoch caches to suit the epoch transition.
    state.advance_caches();

    Ok(())
}

/// Returns a list of active validator indices for the state's current epoch.
///
/// Spec v0.4.0
pub fn calculate_active_validator_indices(state: &BeaconState, spec: &ChainSpec) -> Vec<usize> {
    get_active_validator_indices(
        &state.validator_registry,
        state.slot.epoch(spec.slots_per_epoch),
    )
}

/// Calculates various sets of attesters, including:
///
/// - current epoch attesters
/// - current epoch boundary attesters
/// - previous epoch attesters
/// - etc.
///
/// Spec v0.4.0
pub fn calculate_attester_sets(
    state: &BeaconState,
    active_validator_indices: &[usize],
    spec: &ChainSpec,
) -> Result<Attesters, BeaconStateError> {
    let mut attesters = Attesters::empty(state.validator_registry.len());
    attesters.process_active_validator_indices(&active_validator_indices);
    attesters.process_attestations(&state, &state.latest_attestations, spec)?;
    Ok(attesters)
}

/// Spec v0.4.0
pub fn process_eth1_data(state: &mut BeaconState, spec: &ChainSpec) {
    let next_epoch = state.next_epoch(spec);
    let voting_period = spec.epochs_per_eth1_voting_period;

    if next_epoch % voting_period == 0 {
        for eth1_data_vote in &state.eth1_data_votes {
            if eth1_data_vote.vote_count * 2 > voting_period {
                state.latest_eth1_data = eth1_data_vote.eth1_data.clone();
            }
        }
        state.eth1_data_votes = vec![];
    }
}

/// Update the following fields on the `BeaconState`:
///
/// - `justification_bitfield`.
/// - `finalized_epoch`
/// - `justified_epoch`
/// - `previous_justified_epoch`
///
/// Spec v0.4.0
pub fn process_justification(
    state: &mut BeaconState,
    current_total_balance: u64,
    previous_total_balance: u64,
    previous_epoch_boundary_attesting_balance: u64,
    current_epoch_boundary_attesting_balance: u64,
    spec: &ChainSpec,
) {
    let previous_epoch = state.previous_epoch(spec);
    let current_epoch = state.current_epoch(spec);

    let mut new_justified_epoch = state.justified_epoch;
    state.justification_bitfield <<= 1;

    // If > 2/3 of the total balance attested to the previous epoch boundary
    //
    // - Set the 2nd bit of the bitfield.
    // - Set the previous epoch to be justified.
    if (3 * previous_epoch_boundary_attesting_balance) >= (2 * previous_total_balance) {
        state.justification_bitfield |= 2;
        new_justified_epoch = previous_epoch;
    }
    // If > 2/3 of the total balance attested to the previous epoch boundary
    //
    // - Set the 1st bit of the bitfield.
    // - Set the current epoch to be justified.
    if (3 * current_epoch_boundary_attesting_balance) >= (2 * current_total_balance) {
        state.justification_bitfield |= 1;
        new_justified_epoch = current_epoch;
    }

    // If:
    //
    // - All three epochs prior to this epoch have been justified.
    // - The previous justified justified epoch was three epochs ago.
    //
    // Then, set the finalized epoch to be three epochs ago.
    if ((state.justification_bitfield >> 1) % 8 == 0b111)
        & (state.previous_justified_epoch == previous_epoch - 2)
    {
        state.finalized_epoch = state.previous_justified_epoch;
    }
    // If:
    //
    // - Both two epochs prior to this epoch have been justified.
    // - The previous justified epoch was two epochs ago.
    //
    // Then, set the finalized epoch to two epochs ago.
    if ((state.justification_bitfield >> 1) % 4 == 0b11)
        & (state.previous_justified_epoch == previous_epoch - 1)
    {
        state.finalized_epoch = state.previous_justified_epoch;
    }
    // If:
    //
    // - This epoch and the two prior have been justified.
    // - The presently justified epoch was two epochs ago.
    //
    // Then, set the finalized epoch to two epochs ago.
    if (state.justification_bitfield % 8 == 0b111) & (state.justified_epoch == previous_epoch - 1) {
        state.finalized_epoch = state.justified_epoch;
    }
    // If:
    //
    // - This epoch and the epoch prior to it have been justified.
    // - Set the previous epoch to be justified.
    //
    // Then, set the finalized epoch to be the previous epoch.
    if (state.justification_bitfield % 4 == 0b11) & (state.justified_epoch == previous_epoch) {
        state.finalized_epoch = state.justified_epoch;
    }

    state.previous_justified_epoch = state.justified_epoch;
    state.justified_epoch = new_justified_epoch;
}

/// Updates the following fields on the `BeaconState`:
///
/// - `latest_crosslinks`
///
/// Also returns a `WinningRootHashSet` for later use during epoch processing.
///
/// Spec v0.4.0
pub fn process_crosslinks(
    state: &mut BeaconState,
    spec: &ChainSpec,
) -> Result<WinningRootHashSet, Error> {
    let current_epoch_attestations: Vec<&PendingAttestation> = state
        .latest_attestations
        .par_iter()
        .filter(|a| a.data.slot.epoch(spec.slots_per_epoch) == state.current_epoch(spec))
        .collect();

    let previous_epoch_attestations: Vec<&PendingAttestation> = state
        .latest_attestations
        .par_iter()
        .filter(|a| a.data.slot.epoch(spec.slots_per_epoch) == state.previous_epoch(spec))
        .collect();

    let mut winning_root_for_shards: WinningRootHashSet = HashMap::new();

    let previous_and_current_epoch_slots: Vec<Slot> = state
        .previous_epoch(spec)
        .slot_iter(spec.slots_per_epoch)
        .chain(state.current_epoch(spec).slot_iter(spec.slots_per_epoch))
        .collect();

    for slot in previous_and_current_epoch_slots {
        // Clone removes the borrow which becomes an issue when mutating `state.balances`.
        let crosslink_committees_at_slot =
            state.get_crosslink_committees_at_slot(slot, spec)?.clone();

        for (crosslink_committee, shard) in crosslink_committees_at_slot {
            let shard = shard as u64;

            let winning_root = winning_root(
                state,
                shard,
                &current_epoch_attestations[..],
                &previous_epoch_attestations[..],
                spec,
            )?;

            if let Some(winning_root) = winning_root {
                let total_committee_balance = state.get_total_balance(&crosslink_committee, spec);

                // TODO: I think this has a bug.
                if (3 * winning_root.total_attesting_balance) >= (2 * total_committee_balance) {
                    state.latest_crosslinks[shard as usize] = Crosslink {
                        epoch: state.current_epoch(spec),
                        crosslink_data_root: winning_root.crosslink_data_root,
                    }
                }
                winning_root_for_shards.insert(shard, winning_root);
            }
        }
    }

    Ok(winning_root_for_shards)
}

/// Updates the following fields on the BeaconState:
///
/// - `validator_balances`
///
/// Spec v0.4.0
pub fn process_rewards_and_penalities(
    state: &mut BeaconState,
    attesters: &Attesters,
    previous_total_balance: u64,
    winning_root_for_shards: &WinningRootHashSet,
    spec: &ChainSpec,
) -> Result<(), Error> {
    let next_epoch = state.next_epoch(spec);

    /*
    let previous_epoch_attestations: Vec<&PendingAttestation> = state
        .latest_attestations
        .par_iter()
        .filter(|a| a.data.slot.epoch(spec.slots_per_epoch) == state.previous_epoch(spec))
        .collect();
    */

    let base_reward_quotient = previous_total_balance.integer_sqrt() / spec.base_reward_quotient;

    if base_reward_quotient == 0 {
        return Err(Error::BaseRewardQuotientIsZero);
    }
    if previous_total_balance == 0 {
        return Err(Error::PreviousTotalBalanceIsZero);
    }

    /*
    // Map is ValidatorIndex -> ProposerIndex
    let mut inclusion_slots: FnvHashMap<usize, (Slot, usize)> = FnvHashMap::default();
    for a in &previous_epoch_attestations {
        let participants =
            state.get_attestation_participants(&a.data, &a.aggregation_bitfield, spec)?;
        let inclusion_distance = (a.inclusion_slot - a.data.slot).as_u64();
        for participant in participants {
            if let Some((existing_distance, _)) = inclusion_slots.get(&participant) {
                if *existing_distance <= inclusion_distance {
                    continue;
                }
            }
            let proposer_index = state
                .get_beacon_proposer_index(a.data.slot, spec)
                .map_err(|_| Error::UnableToDetermineProducer)?;
            inclusion_slots.insert(
                participant,
                (Slot::from(inclusion_distance), proposer_index),
            );
        }
    }
    */

    // Justification and finalization

    let epochs_since_finality = next_epoch - state.finalized_epoch;

    state.validator_balances = state
        .validator_balances
        .par_iter()
        .enumerate()
        .map(|(index, &balance)| {
            let mut balance = balance;
            let status = &attesters.statuses[index];

            if epochs_since_finality <= 4 {
                let base_reward = state.base_reward(index, base_reward_quotient, spec);

                // Expected FFG source
                if status.is_previous_epoch {
                    safe_add_assign!(
                        balance,
                        base_reward * attesters.balances.previous_epoch / previous_total_balance
                    );
                } else if status.is_active {
                    safe_sub_assign!(balance, base_reward);
                }

                // Expected FFG target
                if status.is_previous_epoch_boundary {
                    safe_add_assign!(
                        balance,
                        base_reward * attesters.balances.previous_epoch_boundary
                            / previous_total_balance
                    );
                } else if status.is_active {
                    safe_sub_assign!(balance, base_reward);
                }

                // Expected beacon chain head
                if status.is_previous_epoch_head {
                    safe_add_assign!(
                        balance,
                        base_reward * attesters.balances.previous_epoch_head
                            / previous_total_balance
                    );
                } else if status.is_active {
                    safe_sub_assign!(balance, base_reward);
                };
            } else {
                let inactivity_penalty = state.inactivity_penalty(
                    index,
                    epochs_since_finality,
                    base_reward_quotient,
                    spec,
                );

                if status.is_active {
                    if !status.is_previous_epoch {
                        safe_sub_assign!(balance, inactivity_penalty);
                    }
                    if !status.is_previous_epoch_boundary {
                        safe_sub_assign!(balance, inactivity_penalty);
                    }
                    if !status.is_previous_epoch_head {
                        safe_sub_assign!(balance, inactivity_penalty);
                    }

                    if state.validator_registry[index].slashed {
                        let base_reward = state.base_reward(index, base_reward_quotient, spec);
                        safe_sub_assign!(balance, 2 * inactivity_penalty + base_reward);
                    }
                }
            }

            balance
        })
        .collect();

    // Attestation inclusion

    for (index, _validator) in state.validator_registry.iter().enumerate() {
        let status = &attesters.statuses[index];

        if status.is_previous_epoch {
            let proposer_index = status.inclusion_info.proposer_index;
            let inclusion_distance = status.inclusion_info.distance;

            let base_reward = state.base_reward(proposer_index, base_reward_quotient, spec);

            if inclusion_distance > 0 && inclusion_distance < Slot::max_value() {
                safe_add_assign!(
                    state.validator_balances[proposer_index],
                    base_reward * spec.min_attestation_inclusion_delay
                        / inclusion_distance.as_u64()
                )
            }
        }
    }

    //Crosslinks

    for slot in state.previous_epoch(spec).slot_iter(spec.slots_per_epoch) {
        // Clone removes the borrow which becomes an issue when mutating `state.balances`.
        let crosslink_committees_at_slot =
            state.get_crosslink_committees_at_slot(slot, spec)?.clone();

        for (crosslink_committee, shard) in crosslink_committees_at_slot {
            let shard = shard as u64;

            // Note: I'm a little uncertain of the logic here -- I am waiting for spec v0.5.0 to
            // clear it up.
            //
            // What happens here is:
            //
            // - If there was some crosslink root elected by the super-majority of this committee,
            // then we reward all who voted for that root and penalize all that did not.
            // - However, if there _was not_ some super-majority-voted crosslink root, then penalize
            // all the validators.
            //
            // I'm not quite sure that the second case (no super-majority crosslink) is correct.
            if let Some(winning_root) = winning_root_for_shards.get(&shard) {
                // Hash set de-dedups and (hopefully) offers a speed improvement from faster
                // lookups.
                let attesting_validator_indices: FnvHashSet<usize> =
                    FnvHashSet::from_iter(winning_root.attesting_validator_indices.iter().cloned());

                for &index in &crosslink_committee {
                    let base_reward = state.base_reward(index, base_reward_quotient, spec);

                    let total_balance = state.get_total_balance(&crosslink_committee, spec);

                    if attesting_validator_indices.contains(&index) {
                        safe_add_assign!(
                            state.validator_balances[index],
                            base_reward * winning_root.total_attesting_balance / total_balance
                        );
                    } else {
                        safe_sub_assign!(state.validator_balances[index], base_reward);
                    }
                }
            } else {
                for &index in &crosslink_committee {
                    let base_reward = state.base_reward(index, base_reward_quotient, spec);

                    safe_sub_assign!(state.validator_balances[index], base_reward);
                }
            }
        }
    }

    Ok(())
}

/// Peforms a validator registry update, if required.
///
/// Spec v0.4.0
pub fn process_validator_registry(state: &mut BeaconState, spec: &ChainSpec) -> Result<(), Error> {
    let current_epoch = state.current_epoch(spec);
    let next_epoch = state.next_epoch(spec);

    state.previous_shuffling_epoch = state.current_shuffling_epoch;
    state.previous_shuffling_start_shard = state.current_shuffling_start_shard;

    state.previous_shuffling_seed = state.current_shuffling_seed;

    let should_update_validator_registy = if state.finalized_epoch
        > state.validator_registry_update_epoch
    {
        (0..state.get_current_epoch_committee_count(spec)).all(|i| {
            let shard = (state.current_shuffling_start_shard + i as u64) % spec.shard_count;
            state.latest_crosslinks[shard as usize].epoch > state.validator_registry_update_epoch
        })
    } else {
        false
    };

    if should_update_validator_registy {
        state.update_validator_registry(spec);

        state.current_shuffling_epoch = next_epoch;
        state.current_shuffling_start_shard = (state.current_shuffling_start_shard
            + state.get_current_epoch_committee_count(spec) as u64)
            % spec.shard_count;
        state.current_shuffling_seed = state.generate_seed(state.current_shuffling_epoch, spec)?
    } else {
        let epochs_since_last_registry_update =
            current_epoch - state.validator_registry_update_epoch;
        if (epochs_since_last_registry_update > 1)
            & epochs_since_last_registry_update.is_power_of_two()
        {
            state.current_shuffling_epoch = next_epoch;
            state.current_shuffling_seed =
                state.generate_seed(state.current_shuffling_epoch, spec)?
        }
    }

    state.process_slashings(spec);
    state.process_exit_queue(spec);

    Ok(())
}

/// Updates the state's `latest_active_index_roots` field with a tree hash the active validator
/// indices for the next epoch.
///
/// Spec v0.4.0
pub fn update_active_tree_index_roots(
    state: &mut BeaconState,
    spec: &ChainSpec,
) -> Result<(), Error> {
    let next_epoch = state.next_epoch(spec);

    let active_tree_root = get_active_validator_indices(
        &state.validator_registry,
        next_epoch + Epoch::from(spec.activation_exit_delay),
    )
    .hash_tree_root();

    state.latest_active_index_roots[(next_epoch.as_usize()
        + spec.activation_exit_delay as usize)
        % spec.latest_active_index_roots_length] = Hash256::from_slice(&active_tree_root[..]);

    Ok(())
}

/// Advances the state's `latest_slashed_balances` field.
///
/// Spec v0.4.0
pub fn update_latest_slashed_balances(state: &mut BeaconState, spec: &ChainSpec) {
    let current_epoch = state.current_epoch(spec);
    let next_epoch = state.next_epoch(spec);

    state.latest_slashed_balances[next_epoch.as_usize() % spec.latest_slashed_exit_length] =
        state.latest_slashed_balances[current_epoch.as_usize() % spec.latest_slashed_exit_length];
}

/// Removes all pending attestations from the previous epoch.
///
/// Spec v0.4.0
pub fn clean_attestations(state: &mut BeaconState, spec: &ChainSpec) {
    let current_epoch = state.current_epoch(spec);

    state.latest_attestations = state
        .latest_attestations
        .iter()
        .filter(|a| a.data.slot.epoch(spec.slots_per_epoch) >= current_epoch)
        .cloned()
        .collect();
}