use crate::*;
use crate::mock::*;
use strum::IntoEnumIterator;

use sp_std::{
    prelude::Vec,
    collections::{btree_set::BTreeSet, btree_map::BTreeMap},
};
use frame_support::{assert_ok, assert_err, traits::Hooks};
use sp_runtime::testing::UintAuthorityId;

use rand::Rng;
use rand_chacha::rand_core::SeedableRng;
use rand_chacha::ChaCha20Rng;

use frost_core::Field;
use frost_secp256k1::{
    Secp256K1Sha256, Secp256K1ScalarField, Identifier as Secp256K1Identifier,
    keys::{
        VerifyingShare as Secp256K1VerifyingShare,
        SigningShare as Secp256K1SigningShare,
    },
};

const MAX_SIGNERS: u16 = 10;
const MIN_SIGNERS: u16 = 5;

const DUMMY_HASH: ExodusHash = ExodusHash::repeat_byte(0x69);

#[derive(Debug, PartialEq, Eq)]
pub enum TestError {
    CallMismatch,
    PhaseMismatch,
    WorkerExecutionFailed,
    PoolIsEmpty,
    TransactionDecodeFailed,
    UnexpectedRuntimeCall,
    MempoolRejectedTransaction,
    OnChainApplicationFailed,
    InvalidStorageAfter,
}

fn decode_transaction<Runtime, Call, Signature>(raw_tx: &[u8]) -> Option<Call>
where
    Runtime: SendTransactionTypes<Call>,
    <Runtime as SendTransactionTypes<Call>>::Extrinsic: Decode,
    <Runtime as SendTransactionTypes<Call>>::Extrinsic: Into<Extrinsic<Call>>,
{
    let extrinsic = <Runtime as SendTransactionTypes<Call>>::Extrinsic::decode(&mut &raw_tx[..]).ok()?;
    let test_extrinsic: Extrinsic<Call> = extrinsic.into();
    Some(test_extrinsic.call)
}

fn run_increase_block_by(ext: &mut sp_io::TestExternalities, increase_by: u64) {
    ext.execute_with(|| {
        (0..increase_by).for_each(|_| {
            Exodus::on_finalize(System::block_number());
            System::on_finalize(System::block_number());

            System::set_block_number(System::block_number() + 1);

            System::on_initialize(System::block_number());
            Exodus::on_initialize(System::block_number());
        });
    });

    ext.persist_offchain_overlay();
}

fn run_to_next_round(ext: &mut sp_io::TestExternalities, curve: NetworkCurve) {
    let target_block = ext.execute_with(|| {
        let stored = QualificationDkgState::<Runtime>::get(curve).get_block();
        let delay = <<Runtime as Config>::DkgRoundPeriod as frame_support::traits::Get<
            frame_system::pallet_prelude::BlockNumberFor<Runtime>
        >>::get();

        stored + delay
    });

    ext.execute_with(|| {
        while System::block_number() < target_block {
            Exodus::on_finalize(System::block_number());
            System::on_finalize(System::block_number());

            System::set_block_number(System::block_number() + 1);

            System::on_initialize(System::block_number());
            Exodus::on_initialize(System::block_number());
        }
    });

    ext.persist_offchain_overlay();
}

fn generate_mock_secp256k1_shares(
    indices: &[AuthIndex],
) -> BTreeMap<Secp256K1Identifier, Secp256K1VerifyingShare> {
    let mut shares = BTreeMap::new();
    let mut rng = ChaCha20Rng::from_seed(Default::default());

    for &i in indices {
        let random_scalar = <Secp256K1ScalarField as Field>::random(&mut rng);
        let secret_share = Secp256K1SigningShare::new(random_scalar);
        let verifying_share = Secp256K1VerifyingShare::from(secret_share);

        let identifier = Secp256K1Identifier::try_from(i + 1).unwrap();
        shares.insert(identifier, verifying_share);
    }

    shares
}

fn generate_mock_secp256k1_scalars(
    indices: &[AuthIndex],
) -> BTreeMap<Secp256K1Identifier, <Secp256K1ScalarField as Field>::Scalar> {
    let mut scalars = BTreeMap::new();
    let mut rng = ChaCha20Rng::from_seed(Default::default());

    for &i in indices {
        let random_scalar = <Secp256K1ScalarField as Field>::random(&mut rng);
        let identifier = Secp256K1Identifier::try_from(i + 1).unwrap();
        scalars.insert(identifier, random_scalar);
    }

    scalars
}

fn execute_dkg_happy_path<F1, F2, F3>(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    curve: NetworkCurve,
    authority: u64,
    expected_phase: DkgPhase,
    mut execute_transaction: F1,
    mut apply_transaction: F2,
    mut check_storage_changed: F3,
) -> Result<(), TestError>
where
	F1: FnMut(
        DkgStorage<BlockNumberFor<Runtime>, NetworkCurve>,
        &QualifyingState<Runtime>,
        UintAuthorityId,
    ) -> Result<ExodusOk<NetworkCurve>, ExodusError>,

	F2: FnMut(Call<Runtime>) -> Result<(), TestError>,
	F3: FnMut(AuthIndex, NetworkCurve) -> Result<(), TestError>,
{
    use sp_runtime::transaction_validity::TransactionSource;

    let state = ext.execute_with(|| QualificationDkgState::<Runtime>::get(&curve));

    if state.get_phase() != expected_phase {
        return Err(TestError::PhaseMismatch);
    }

    let current_block = ext.execute_with(|| Exodus::current_block_number());
    let block_longevity = ext.execute_with(|| {
        < <Runtime as Config>::UnsignedLongevity as frame_support::traits::Get<u64> >::get()
    });

    let authority_index = authority as AuthIndex;
    let authority_key = authority.into();

    let dkg_storage = DkgStorage::default()
        .with_network_curve(curve)
        .with_current_block(current_block)
        .with_block_longevity(block_longevity)
        .with_authority_index(authority_index);

    let round_result = ext.execute_with(|| {
        execute_transaction(dkg_storage, &state, authority_key)
    });

    if round_result.is_err() {
        return Err(TestError::WorkerExecutionFailed);
    }

    let mut write_pool = tx_pool_state.write();
    let raw_tx = write_pool.transactions.pop().ok_or(TestError::PoolIsEmpty)?;
    let runtime_call = decode_transaction::<Runtime, RuntimeCall, ()>(&raw_tx)
        .ok_or(TestError::TransactionDecodeFailed)?;

    if let RuntimeCall::Exodus(inner_exodus_call) = runtime_call {
        ext.execute_with(|| {
            let is_mempool_ok = Exodus::validate_unsigned(TransactionSource::InBlock, &inner_exodus_call).is_ok();
            if !is_mempool_ok {
                return Err(TestError::MempoolRejectedTransaction);
            }

            apply_transaction(inner_exodus_call)
                .map_err(|_| TestError::OnChainApplicationFailed)?;

            Ok(())
        })?;

    } else {
        return Err(TestError::UnexpectedRuntimeCall);
    }
    write_pool.transactions.clear();
    ext.execute_with(|| check_storage_changed(authority_index, curve))?;

    Ok(())
}

fn execute_exodus_happy_path<F1, F2>(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    curve: NetworkCurve,
    authority: u64,
    mut apply_transaction: F1,
    mut check_storage_changed: F2,
) -> Result<(), TestError>
where
    F1: FnMut(Call<Runtime>) -> Result<(), TestError>,
	F2: FnMut(ExodusOk<NetworkCurve>) -> Result<(), TestError>,
{
    use sp_runtime::transaction_validity::TransactionSource;

    let state = ext.execute_with(|| QualificationDkgState::<Runtime>::get(&curve));
    if state.is_dkg_pending() {
        return Err(TestError::PhaseMismatch);
    }

    let active_authorities = ext.execute_with(|| {
        ActiveDkgAuthorities::<Runtime>::get(&curve)
    });

    let current_block = ext.execute_with(|| Exodus::current_block_number());
    let authority_index = authority as AuthIndex;
    let authority_key = authority.into();

    let exodus_process_result = ext.execute_with(|| {
        Exodus::exodus_run_process(
            authority_index,
            authority_key,
            curve,
            current_block,
            &active_authorities,
        )
    });

    match exodus_process_result {
        Ok(ExodusOk::ExodusEmptyRequests(_, _)) => return Ok(()),
        Err(_) => return Err(TestError::WorkerExecutionFailed),
        _ => {}
    }

    let mut write_pool = tx_pool_state.write();
    let raw_tx = write_pool.transactions.pop().ok_or(TestError::PoolIsEmpty)?;
    let runtime_call = decode_transaction::<Runtime, RuntimeCall, ()>(&raw_tx)
        .ok_or(TestError::TransactionDecodeFailed)?;

    if let RuntimeCall::Exodus(inner_exodus_call) = runtime_call {
        ext.execute_with(|| {
            let is_mempool_ok = Exodus::validate_unsigned(TransactionSource::InBlock, &inner_exodus_call).is_ok();
            if !is_mempool_ok {
                return Err(TestError::MempoolRejectedTransaction);
            }

            apply_transaction(inner_exodus_call)
                .map_err(|_| TestError::OnChainApplicationFailed)?;

            Ok(())
        })?;

    } else {
        return Err(TestError::UnexpectedRuntimeCall);
    }
    write_pool.transactions.clear();
    ext.execute_with(|| {
        check_storage_changed(exodus_process_result.unwrap())
    })?;

    Ok(())
}

fn execute_round0_happy_path(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    curve: NetworkCurve,
    authority: u64,
) -> Result<(), TestError> {
    let execute_transaction = |sb, state: &_, key| Exodus::dkg_run_round0(sb, &state, key);

    let apply_transaction = |call| match call {
        Call::register_round0_package { dkg_package, signature } => {
            Exodus::register_round0_package(RuntimeOrigin::none(), dkg_package, signature)
                .map_err(|_| TestError::OnChainApplicationFailed)
        },
        _ => Err(TestError::CallMismatch),
    };

    let check_storage_changed = |_, curve| {
        for (_, hash) in Round0Packages::<Runtime>::iter_prefix(&curve) {
            if hash == ExodusHash::default() || hash == DUMMY_HASH {
                return Err(TestError::InvalidStorageAfter);
            }
        }

        Ok(())
    };

    execute_dkg_happy_path(
        ext,
        &tx_pool_state,
        curve,
        authority,
        DkgPhase::Round0,
        execute_transaction,
        apply_transaction,
        check_storage_changed,
    )
}

fn execute_round1_happy_path(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    curve: NetworkCurve,
    authority: u64,
) -> Result<(), TestError> {
    let execute_transaction = |sb, state: &_, key| Exodus::dkg_run_round1(sb, &state, key);

    let apply_transaction = |call| match call {
        Call::register_round1_package { dkg_package, signature } => {
            Exodus::register_round1_package(RuntimeOrigin::none(), dkg_package, signature)
                .map_err(|_| TestError::OnChainApplicationFailed)
        },
        _ => Err(TestError::CallMismatch),
    };

    let check_storage_changed = |idx, curve| {
        let participants = QualificationAuthorities::<Runtime>::decode_len(&curve).unwrap();
        let round1_packages = Round1Packages::<Runtime>::get(curve);

        let round1_contains_idx = round1_packages.contains(idx);
        let round1_valid_len = round1_packages.active_bits_count() == participants as u32;

        if !round1_contains_idx || !round1_valid_len {
            return Err(TestError::InvalidStorageAfter);
        }

        Ok(())
    };

    execute_dkg_happy_path(
        ext,
        &tx_pool_state,
        curve,
        authority,
        DkgPhase::Round1,
        execute_transaction,
        apply_transaction,
        check_storage_changed,
    )
}

fn execute_round2_happy_path(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    curve: NetworkCurve,
    authority: u64,
) -> Result<(), TestError> {
    let execute_transaction = |sb, state: &_, key| Exodus::dkg_run_round2(sb, &state, key);

    let apply_transaction = |call| match call {
        Call::register_encrypted_round2_packages { dkg_package, signature } => {
            Exodus::register_encrypted_round2_packages(RuntimeOrigin::none(), dkg_package, signature)
                .map_err(|_| TestError::OnChainApplicationFailed)
        },
        _ => Err(TestError::CallMismatch),
    };

    let check_storage_changed = |idx, curve| {
        let participants = QualificationAuthorities::<Runtime>::decode_len(&curve).unwrap();
        let round2_packages = Round2Packages::<Runtime>::get(curve);

        let round2_contains_idx = round2_packages.contains(idx);
        let round2_valid_len = round2_packages.active_bits_count() == participants as u32;

        if !round2_contains_idx || !round2_valid_len {
            return Err(TestError::InvalidStorageAfter);
        }

        Ok(())
    };

    execute_dkg_happy_path(
        ext,
        &tx_pool_state,
        curve,
        authority,
        DkgPhase::Round2,
        execute_transaction,
        apply_transaction,
        check_storage_changed,
    )
}

fn execute_round3_happy_path(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    curve: NetworkCurve,
    authority: u64,
) -> Result<(), TestError> {
    let execute_transaction = |sb, state: &_, key| Exodus::dkg_run_round3(sb, &state, key);

    let apply_transaction = |call| match call {
        Call::register_round3_complaints { dkg_package, signature } => {
            Exodus::register_round3_complaints(RuntimeOrigin::none(), dkg_package, signature)
                .map_err(|_| TestError::OnChainApplicationFailed)
        },
        _ => Err(TestError::CallMismatch),
    };

    let check_storage_changed = |idx, curve| {
        let participants = QualificationAuthorities::<Runtime>::decode_len(&curve).unwrap();
        let complaints = Complaints::<Runtime>::get(curve);

        let complaints_valid = complaints.get(&idx)
            .map(|bitmap| {
                bitmap.active_bits_count() <= participants as u32
            })
            .unwrap_or(false);

        if !complaints_valid {
            return Err(TestError::InvalidStorageAfter);
        }

        Ok(())
    };

    execute_dkg_happy_path(
        ext,
        &tx_pool_state,
        curve,
        authority,
        DkgPhase::Round3,
        execute_transaction,
        apply_transaction,
        check_storage_changed,
     )
}

fn execute_round4_happy_path(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    curve: NetworkCurve,
    authority: u64,
) -> Result<(), TestError> {
    let execute_transaction = |sb, state: &_, key| Exodus::dkg_run_round4(sb, &state, key);

    let apply_transaction = |call| match call {
        Call::register_round4_justifications { dkg_package, signature } => {
            Exodus::register_round4_justifications(RuntimeOrigin::none(), dkg_package, signature)
                .map_err(|_| TestError::OnChainApplicationFailed)
        },
        _ => Err(TestError::CallMismatch),
    };

    let check_storage_changed = |idx, curve| {
        let curr_dkg_index = QualificationDkgState::<Runtime>::get(curve).get_dkg_index();
        let participants = QualificationAuthorities::<Runtime>::decode_len(&curve).unwrap();
        let justifications = Justifications::<Runtime>::get(curve);

        let justifications_valid = justifications.get(&idx)
            .map(|bitmap| {
                bitmap.active_bits_count() <= participants as u32
            })
            .unwrap_or(false);

        if !justifications_valid {
            return Err(TestError::InvalidStorageAfter);
        }

        let non_empty_past_verifications = (0..curr_dkg_index).any(|dkg_index| {
            Verifications::<Runtime>::iter_key_prefix((curve, dkg_index))
                .next()
                .is_some()
        });

        if non_empty_past_verifications {
            return Err(TestError::InvalidStorageAfter);
        }

        Ok(())
    };

    execute_dkg_happy_path(
        ext,
        &tx_pool_state,
        curve,
        authority,
        DkgPhase::Round4,
        execute_transaction,
        apply_transaction,
        check_storage_changed,
    )
}

fn execute_round5_happy_path(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    curve: NetworkCurve,
    authority: u64,
) -> Result<(), TestError> {
    let execute_transaction = |sb, state: &_, key| Exodus::dkg_run_round5(sb, &state, key);

    let apply_transaction = |call| match call {
        Call::register_round5_verifying_package { dkg_package, signature } => {
            Exodus::register_round5_verifying_package(RuntimeOrigin::none(), dkg_package, signature)
                .map_err(|_| TestError::OnChainApplicationFailed)
        },
        _ => Err(TestError::CallMismatch),
    };

    let check_storage_changed = |_, curve| {
        let curr_dkg_index = QualificationDkgState::<Runtime>::get(curve).get_dkg_index();

        let non_empty_verification_states = (0..curr_dkg_index).any(|dkg_index| {
            VerifyingKeyConsensus::<Runtime>::contains_key(curve, dkg_index)
        });

        if non_empty_verification_states {
            return Err(TestError::InvalidStorageAfter);
        }

        Ok(())
    };

    execute_dkg_happy_path(
        ext,
        &tx_pool_state,
        curve,
        authority,
        DkgPhase::Round5,
        execute_transaction,
        apply_transaction,
        check_storage_changed,
    )
}

fn execute_round6_happy_path(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    curve: NetworkCurve,
    authority: u64,
) -> Result<(), TestError> {
    let execute_transaction = |sb, state: &_, key| Exodus::dkg_run_round6(sb, &state, key);

    let apply_transaction = |call| match call {
        Call::register_round6_public_share_package { dkg_package, signature } => {
            Exodus::register_round6_public_share_package(RuntimeOrigin::none(), dkg_package, signature)
                .map_err(|_| TestError::OnChainApplicationFailed)
        },
        _ => Err(TestError::CallMismatch),
    };

    let check_storage_changed = |idx, curve| {
        let curr_dkg_index = QualificationDkgState::<Runtime>::get(curve).get_dkg_index();
        let verifying_share_key = (curve, curr_dkg_index, idx);

        if VerifyingShares::<Runtime>::get(verifying_share_key).is_none() {
            return Err(TestError::InvalidStorageAfter);
        }

        Ok(())
    };

    execute_dkg_happy_path(
        ext,
        &tx_pool_state,
        curve,
        authority,
        DkgPhase::Round6,
        execute_transaction,
        apply_transaction,
        check_storage_changed,
    )
}

fn run_succesfull_dkg_session_partially(
    ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    raw_authorities: &[u64],
    curve: NetworkCurve,
) -> (u64, u64) {
    let authorities: Vec<UintAuthorityId> = raw_authorities.iter().map(|&id| id.into()).collect();
    let prev_dkg_index = ext.execute_with(|| QualificationDkgState::<Runtime>::get(curve).get_dkg_index());
    ext.execute_with(|| Exodus::start_dkg_qualification(authorities, NetworkCurve::iter()));

    let dkg_index = ext.execute_with(|| QualificationDkgState::<Runtime>::get(curve).get_dkg_index());
    let (start_block, round_delay) = ext.execute_with(|| {
        let round_delay = <<Runtime as Config>::DkgRoundPeriod as frame_support::traits::Get<
            frame_system::pallet_prelude::BlockNumberFor<Runtime>
            >>::get();

        let start_block = System::block_number();

        (start_block, round_delay)
    });

    assert_eq!(prev_dkg_index + 1, dkg_index);

    ext.execute_with(|| {
        let dummy_oversized_bitmap = ParticipantsBitmap::<Runtime>::all_ones(1337);
        let mut map: BitmapByAuthority<Runtime> = Default::default();

        for authority in raw_authorities.iter() {
            let auth_idx = *authority as AuthIndex;
            Round0Packages::<Runtime>::insert(curve, auth_idx, DUMMY_HASH);
            map.try_insert(auth_idx, dummy_oversized_bitmap.clone()).unwrap();
        }

        Round1Packages::<Runtime>::insert(curve, dummy_oversized_bitmap.clone());
        Round2Packages::<Runtime>::insert(curve, dummy_oversized_bitmap.clone());

        let prev_verifications_key1 = (curve, prev_dkg_index, DUMMY_HASH);
        let prev_verifications_key2 = (curve, prev_dkg_index, ExodusHash::default());

        Verifications::<Runtime>::insert(prev_verifications_key1, dummy_oversized_bitmap.clone());
        Verifications::<Runtime>::insert(prev_verifications_key2, dummy_oversized_bitmap.clone());

        let dummy_verifying_key = BoundedVec::<u8, ConstU32<ELEMENT_MAX_BYTES>>::try_from(
            vec![69u8; ELEMENT_MAX_BYTES as usize],
        ).unwrap();

        let dummy_consensus_state = ConsensusState::new(
            dummy_verifying_key,
            dummy_oversized_bitmap,
            DUMMY_HASH,
        );

        VerifyingKeyConsensus::<Runtime>::insert(curve, prev_dkg_index, dummy_consensus_state);
    });

    let dkg_rounds = [
        execute_round0_happy_path,
        execute_round1_happy_path,
        execute_round2_happy_path,
        execute_round3_happy_path,
        execute_round4_happy_path,
        execute_round5_happy_path,
        execute_round6_happy_path,
    ];

    for dkg_round in dkg_rounds {
        run_to_next_round(ext, curve);
        for authority in raw_authorities.iter() {
            assert_ok!(dkg_round(ext, &tx_pool_state, curve, *authority));
        }
    }

    run_to_next_round(ext, curve);

    (start_block, round_delay)
}

fn run_initial_dkg_session(raw_authorities: &[u64], curve: NetworkCurve) -> (
	sp_io::TestExternalities,
    std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
) {
    let (mut ext, tx_pool_state) = new_test_ext();

    let prev_verifying_key = ext.execute_with(|| ActiveVerifyingKey::<Runtime>::get(curve));
    assert!(prev_verifying_key.is_empty());

    let (start_block, round_delay) =
        run_succesfull_dkg_session_partially(&mut ext, &tx_pool_state, &raw_authorities, curve);

    run_to_next_round(&mut ext, curve);

    let curr_verifying_key = ext.execute_with(|| ActiveVerifyingKey::<Runtime>::get(curve));
    assert!(!curr_verifying_key.is_empty());

    let state = ext.execute_with(|| QualificationDkgState::<Runtime>::get(curve));
    let authorities_len = raw_authorities.len();

    for authority in raw_authorities {
        let member = *authority as AuthIndex;
        assert!(state.contains_index(member));
    }
    assert_eq!(state.get_block(), start_block + round_delay * 9);
    assert_eq!(state.get_dkg_index(), 1);
    assert_eq!(state.get_phase(), DkgPhase::Finalized);

    assert_eq!(state.count_ones::<usize>(), authorities_len);
    assert_eq!(state.get_indexes().active_bits_count(), authorities_len as u32);

    let active = ext.execute_with(|| ActiveDkgAuthorities::<Runtime>::get(curve));
    assert_eq!(active.get_indexes(), state.get_indexes());
    assert_eq!(active.get_dkg_index(), state.get_dkg_index());

    let active_authorities = ext.execute_with(|| ActiveAuthorities::<Runtime>::get(curve));
    let qualification_authorities = ext.execute_with(|| QualificationAuthorities::<Runtime>::get(curve));
    assert_eq!(active_authorities, qualification_authorities);

    (ext, tx_pool_state)
}

fn run_dkg_session(
	ext: &mut sp_io::TestExternalities,
    tx_pool_state: &std::sync::Arc<parking_lot::RwLock<sp_core::offchain::testing::PoolState>>,
    prev_authorities: Vec<u64>,
    raw_authorities: Vec<u64>,
    curve: NetworkCurve,
) {
    let prev_active_authorities = ext.execute_with(|| ActiveAuthorities::<Runtime>::get(curve));
    let prev_verifying_key = ext.execute_with(|| ActiveVerifyingKey::<Runtime>::get(curve));
    assert!(!prev_verifying_key.is_empty());

    let dkg_index = ext.execute_with(|| {
        QualificationDkgState::<Runtime>::get(curve).get_dkg_index()
    });
    let next_dkg_index = dkg_index.saturating_add(1);

    let (start_block, round_delay) =
        run_succesfull_dkg_session_partially(ext, tx_pool_state, &raw_authorities, curve);

    let register_nonce = |call| {
        match call {
            Call::register_nonce_commitment { exodus_package, signature } => {
                Exodus::register_nonce_commitment(RuntimeOrigin::none(), exodus_package, signature)
                   .map_err(|_| TestError::OnChainApplicationFailed)
            }
            _ => Err(TestError::CallMismatch),
        }
    };

    let register_group = |call| {
        match call {
            Call::register_group_commitment { exodus_package, signature } => {
                Exodus::register_group_commitment(RuntimeOrigin::none(), exodus_package, signature)
                    .map_err(|_| TestError::OnChainApplicationFailed)
            }
            _ => Err(TestError::CallMismatch),
        }
    };

    let register_share = |call| {
        match call {
            Call::register_signature_share { exodus_package, signature } => {
                Exodus::register_signature_share(RuntimeOrigin::none(), exodus_package, signature)
                    .map_err(|_| TestError::OnChainApplicationFailed)
            }
            _ => Err(TestError::CallMismatch),
        }
    };

    let check_nonce = |exodus_ok| {
        match exodus_ok {
            ExodusOk::ExodusNonceCommitted(exodus, roast, idx, _) => {
                if !NonceCommitments::<Runtime>::contains_key((exodus, roast), idx) {
                    return Err(TestError::InvalidStorageAfter);
                }

                if !RoastSessions::<Runtime>::contains_key(exodus, idx) {
                    return Err(TestError::InvalidStorageAfter);
                }

                let roast_state = RoastSessionStates::<Runtime>::get(exodus, roast);

                if !roast_state.nonce_committers.contains(idx)
                    || roast_state.group_committers.contains(idx)
                    || roast_state.partial_signers.contains(idx) {
                        return Err(TestError::InvalidStorageAfter);
                    }

                Ok(())
            }
            _ => Err(TestError::InvalidStorageAfter),
        }
    };

    let check_group = |exodus_ok| {
        match exodus_ok {
            ExodusOk::ExodusGroupCommitted(exodus, roast, idx, _) => {
                if !RoastSessions::<Runtime>::contains_key(exodus, idx) {
                    return Err(TestError::InvalidStorageAfter);
                }

                if !GroupCommitters::<Runtime>::iter_prefix_values((exodus, roast))
                    .any(|participants| participants.contains(idx)) {
                        return Err(TestError::InvalidStorageAfter);
                    }

                let roast_state = RoastSessionStates::<Runtime>::get(exodus, roast);

                let is_nonce_committer = roast_state.nonce_committers.contains(idx);
                let is_group_committer = roast_state.group_committers.contains(idx);

                if !is_nonce_committer || !is_group_committer {
            		return Err(TestError::InvalidStorageAfter);
            	}

                Ok(())
            }
            _ => Err(TestError::InvalidStorageAfter),
        }
    };

    let check_share = |exodus_ok| {
        match exodus_ok {
            ExodusOk::ExodusSignedPartially(exodus, roast, idx, curve) => {
                if RoastSessions::<Runtime>::contains_key(exodus, idx) {
                    return Err(TestError::InvalidStorageAfter);
                }

                if RoastSessions::<Runtime>::contains_key(exodus, idx) {
                    return Err(TestError::InvalidStorageAfter);
                }

                if ExodusRequests::<Runtime>::contains_key(curve, exodus) {
                    let roast_state = RoastSessionStates::<Runtime>::get(exodus, roast);
                    if !roast_state.partial_signers.contains(idx)
                        || !roast_state.nonce_committers.contains(idx)
                        || !roast_state.group_committers.contains(idx) {
                            return Err(TestError::InvalidStorageAfter);
                        }

                } else {
                    if RoastSessionStates::<Runtime>::contains_key(exodus, roast) {
                        return Err(TestError::InvalidStorageAfter);
                    }

                    if SignatureScalars::<Runtime>::contains_key(exodus, roast) {
                        return Err(TestError::InvalidStorageAfter);
                    }

                    if GroupCommitmentsConsensus::<Runtime>::contains_key(exodus, roast) {
                        return Err(TestError::InvalidStorageAfter);
                    }
                }

                Ok(())
            }
            _ => Err(TestError::InvalidStorageAfter),
        }
    };

    let threshold = ext.execute_with(|| {
        let len = ActiveAuthorities::<Runtime>::decode_len(&curve)
            .unwrap_or(0);
        get_byzantium_threshold(len) as u16
    });

    let mut exodus_steps_outer = BTreeMap::new();
    let mut roast_sessions_outer = BTreeMap::new();
    let mut roast_participants_outer = BTreeMap::new();
    let mut malicious_participants_outer = BTreeMap::new();
    let mut throttled_participants_outer = BTreeMap::new();

    let malicious = prev_authorities.first().copied().unwrap();
    let throttled = prev_authorities.last().copied().unwrap();
    let prev_len = prev_authorities.len() as AuthIndex;

    let mut start = 0;
    loop {
        if start >= 100 { break; }
        start += 1;

        let mut current_pending_requests: Vec<ExodusSession> = ext.execute_with(|| {
            ExodusRequests::<Runtime>::iter_keys()
                .filter_map(|(current_curve, exodus_session)| {
                    curve.eq(&current_curve).then(|| exodus_session)
                })
                .collect()
        });

        current_pending_requests.sort();

        if let Some(exodus_session) = current_pending_requests.first() {
            if !current_pending_requests.contains(exodus_session) { continue; }

            let exodus_steps: &mut BTreeMap<u64, RoastSession> = exodus_steps_outer.entry(*exodus_session).or_default();
            let roast_sessions: &mut BTreeMap<u64, RoastSession> = roast_sessions_outer.entry(*exodus_session).or_default();
            let roast_participants: &mut BTreeMap<RoastSession, (u16, u16, u16)> = roast_participants_outer.entry(*exodus_session).or_default();

            let malicious_participants: &mut BTreeSet<u64> = malicious_participants_outer.entry(*exodus_session).or_default();
            let throttled_participants: &mut BTreeSet<u64> = throttled_participants_outer.entry(*exodus_session).or_default();

            for auth in prev_authorities.iter() {
                let request_exists = ext.execute_with(|| {
                    ExodusRequests::<Runtime>::contains_key(&curve, &exodus_session)
                });

                if !request_exists { break; }
                if malicious_participants.contains(auth) { continue; }
                if throttled_participants.contains(auth) { continue; }

                let mut current_step = exodus_steps.get(&auth).copied().unwrap_or(0);
                let mut roast_session = roast_sessions.get(&auth).copied().unwrap_or(0);

                let mut participants = roast_participants
                    .get(&roast_session)
                    .copied()
                    .unwrap_or((0, 0, 0));

                let mut has_enough_nonces = participants.0 >= threshold;
                let mut has_enough_groups = participants.1 >= threshold;

                if has_enough_nonces && current_step == 0 {
                    roast_session += 1;

                    participants = roast_participants
                        .get(&roast_session)
                        .copied()
                        .unwrap_or((0, 0, 0));

                    has_enough_nonces = participants.0 >= threshold;
                    has_enough_groups = participants.1 >= threshold;
                }

                let result = match current_step {
                    0 => {
                        current_step += 1;
                        participants.0 += 1;
                        execute_exodus_happy_path(ext, tx_pool_state, curve, *auth, register_nonce, check_nonce)
                    },
                    1 if has_enough_nonces => {
                        participants.1 += 1;
                        let result = execute_exodus_happy_path(ext, tx_pool_state, curve, *auth, register_group, check_group);

                        if malicious == *auth && prev_len > threshold {
                            ext.execute_with(|| {
                                let raw_vec = vec![69u8; ELEMENT_MAX_BYTES as usize];
                                let fake_group_commitment =
                                    BoundedVec::<u8, ConstU32<{ ELEMENT_MAX_BYTES }>>::try_from(raw_vec)
                                        .unwrap();

                                let empty_participants = ParticipantsBitmap::<Runtime>::empty(prev_len);
                                let mut self_participant = empty_participants.clone();
                                self_participant.insert(malicious);

                                let (correct_commitment, correct_root) = GroupCommitmentsConsensus::<Runtime>::mutate(
                                    &exodus_session,
                                    &roast_session,
                                    |consensus| {
                                        let binding_root = consensus.merkle_root;
                                        let group_commitment = consensus
                                            .element_bytes
                                            .clone();

                                        *consensus = ConsensusState::new(
                                            fake_group_commitment.clone(),
                                            consensus.participants.clone(),
                                            DUMMY_HASH,
                                        );

                                        (group_commitment, binding_root)
                                    });

                                let prev_group_key = (exodus_session, roast_session, correct_root, correct_commitment);
                                GroupCommitters::<Runtime>::insert(prev_group_key, empty_participants);

                                let fake_group_key = (exodus_session, roast_session, DUMMY_HASH, fake_group_commitment);
                                GroupCommitters::<Runtime>::insert(fake_group_key, self_participant);

                                malicious_participants.insert(malicious);
                            });
                        } else {
                            current_step += 1;
                        }

                        result
                    }
                    2 if has_enough_groups => {
                        if prev_len > threshold + 1 && *auth == throttled {
                            throttled_participants.insert(*auth);
                            continue;
                        }

                        current_step = 0;
                        participants.2 += 1;
                        execute_exodus_happy_path(ext, tx_pool_state, curve, *auth, register_share, check_share).unwrap();

                        exodus_steps.insert(*auth, current_step);
                        roast_sessions.insert(*auth, roast_session + 1);
                        roast_participants.insert(roast_session, participants);

                        continue;
                    }
                    _ => continue,
                };

                result.unwrap();
                exodus_steps.insert(*auth, current_step);
                roast_sessions.insert(*auth, roast_session);
                roast_participants.insert(roast_session, participants);
            }

            run_increase_block_by(ext, 1);

        } else {
            break;
        }
    }

    run_to_next_round(ext, curve);

    let curr_verifying_key = ext.execute_with(|| ActiveVerifyingKey::<Runtime>::get(curve));
    assert!(curr_verifying_key != prev_verifying_key);

    let state = ext.execute_with(|| QualificationDkgState::<Runtime>::get(curve));
    let authorities_len = raw_authorities.len();

    for authority in raw_authorities {
        let member = authority as AuthIndex;
        assert!(state.contains_index(member));
    }
    assert_eq!(state.get_block(), start_block + round_delay * 9);
    assert_eq!(state.get_dkg_index(), next_dkg_index);
    assert_eq!(state.get_phase(), DkgPhase::Finalized);

    assert_eq!(state.count_ones::<usize>(), authorities_len);
    assert_eq!(state.get_indexes().active_bits_count(), authorities_len as u32);

    let active = ext.execute_with(|| ActiveDkgAuthorities::<Runtime>::get(curve));
    assert_eq!(active.get_indexes(), state.get_indexes());
    assert_eq!(active.get_dkg_index(), state.get_dkg_index());

    let active_authorities = ext.execute_with(|| ActiveAuthorities::<Runtime>::get(curve));
    let qualification_authorities = ext.execute_with(|| QualificationAuthorities::<Runtime>::get(curve));
    assert_eq!(active_authorities, qualification_authorities);
    assert_ne!(prev_active_authorities, active_authorities);
}

#[test]
fn test_all_validators_can_execute_initial_dkg() {
    let curve = NetworkCurve::Secp256k1;
    let raw_authorities = FixedAuthorities::get();

    run_initial_dkg_session(&raw_authorities, curve);
}

#[test]
fn test_all_validators_can_execute_dkg() {
    let curve = NetworkCurve::Secp256k1;
    let authorities_1 = FixedAuthorities::get();

    let (mut ext, tx_pool_state) = run_initial_dkg_session(&authorities_1, curve);

    ext.execute_with(|| {
        Exodus::do_register_evm_bridge_out_exodus(
            31_337,
            69,
            Default::default(),
            EvmAddress::from_low_u64_be(1),
        ).unwrap();
    });

    let mut authorities_2 = authorities_1.clone();
    let new_authority = authorities_2.iter().max().unwrap() + 1;
    authorities_2.push(new_authority);
    authorities_2.push(new_authority + 1);
    authorities_2.push(new_authority + 2);

    run_dkg_session(&mut ext, &tx_pool_state, authorities_1, authorities_2.clone(), curve);

    ext.execute_with(|| {
        Exodus::do_register_evm_bridge_out_exodus(
            31_337,
            420,
            Perbill::from_percent(50),
            EvmAddress::from_low_u64_be(2),
        ).unwrap();
    });

    let mut authorities_3 = FixedAuthorities::get();
    authorities_3.push(new_authority);

    run_dkg_session(&mut ext, &tx_pool_state, authorities_2, authorities_3, curve);

    let rotations_count = ext.execute_with(|| {
        ExodusSignedRotations::<Runtime>::iter().count()
    });

    let bridges_count = ext.execute_with(|| {
        ExodusSignedBridges::<Runtime>::iter().count()
    });

    assert_eq!(rotations_count, 3);
    assert_eq!(bridges_count, 2);
}

#[test]
fn test_exodus_dkg_robustness() {
    let (mut ext, tx_pool_state) = new_test_ext();

    let curve = NetworkCurve::Secp256k1;
    let raw_authorities = FixedAuthorities::get();
    let authorities: Vec<UintAuthorityId> = raw_authorities.iter().map(|&id| id.into()).collect();

    let mut excluded_authorities = sp_std::collections::btree_set::BTreeSet::new();

    let dkg_rounds = [
        execute_round0_happy_path,
        execute_round1_happy_path,
        execute_round2_happy_path,
        execute_round3_happy_path,
        execute_round4_happy_path,
        execute_round5_happy_path,
        execute_round6_happy_path,
    ];

    let mut dkg_index = ext.execute_with(|| {
        Exodus::start_dkg_qualification(authorities, NetworkCurve::iter());
        QualificationDkgState::<Runtime>::get(curve).get_dkg_index()
    });

    let (start_block, round_delay) = ext.execute_with(|| {
        let round_delay = <<Runtime as Config>::DkgRoundPeriod as frame_support::traits::Get<
            frame_system::pallet_prelude::BlockNumberFor<Runtime>
            >>::get();

        let start_block = System::block_number();

        (start_block, round_delay)
    });

    let mut failed_at_round = None;
    let mut total_round_idx = 0;
    let mut round_idx = 0;

    loop {
        run_to_next_round(&mut ext, curve);
        total_round_idx += 1;

        let phase = ext.execute_with(|| {
            QualificationDkgState::<Runtime>::get(curve).get_phase()
        });

        match phase {
            DkgPhase::Finalized => break,
            DkgPhase::Round7 => continue,
            DkgPhase::Vacant => {
                failed_at_round = Some(round_idx);
                dkg_index += 1;
                round_idx = 0;

                excluded_authorities.clear();
                continue;
            }
            _ => {}
        }

        let is_fast_track = failed_at_round
            .map(|failed_round| round_idx < failed_round)
            .unwrap_or(false);

        let mut maybe_lost_authority = None;
        if !is_fast_track {
            maybe_lost_authority = Some((total_round_idx % raw_authorities.len()) as u64);
        }

        for authority in raw_authorities.iter() {
            if is_fast_track {
                assert_ok!(dkg_rounds[round_idx](&mut ext, &tx_pool_state, curve, *authority));
                continue;
            }

            if let Some(lost_auth) = maybe_lost_authority {
                if *authority == lost_auth { continue; }
            }

            if excluded_authorities.contains(authority) {
                if round_idx != 0 {
                    assert_err!(
                        dkg_rounds[round_idx](&mut ext, &tx_pool_state, curve, *authority),
                        TestError::PoolIsEmpty,
                    );
                }
                continue;
            }

            assert_ok!(dkg_rounds[round_idx](&mut ext, &tx_pool_state, curve, *authority));
        }

        if !is_fast_track {
            if let Some(lost_auth) = maybe_lost_authority {
                excluded_authorities.insert(lost_auth);
            }
        }

        round_idx += 1;
    }

    let state = ext.execute_with(|| QualificationDkgState::<Runtime>::get(curve));
    let authorities_len = raw_authorities.len();

    let qualification_count = raw_authorities
        .iter()
        .filter(|&authority| {
            let member = *authority as AuthIndex;
            state.contains_index(member)
        })
        .count();

    let estimated_block = start_block + round_delay * (total_round_idx as u64);
    assert_eq!(state.get_block(), estimated_block);
    assert_eq!(state.get_dkg_index(), dkg_index);
    assert_eq!(state.get_phase(), DkgPhase::Finalized);

    let bft_threshold = get_byzantium_threshold(authorities_len);
    assert!(qualification_count >= bft_threshold);
    assert!(state.count_ones::<usize>() >= bft_threshold);
    assert_eq!(state.get_indexes().active_bits_count(), authorities_len as u32);

    let active = ext.execute_with(|| ActiveDkgAuthorities::<Runtime>::get(curve));
    assert_eq!(active.get_indexes(), state.get_indexes());
    assert_eq!(active.get_dkg_index(), state.get_dkg_index());

    let active_authorities = ext.execute_with(|| ActiveAuthorities::<Runtime>::get(curve));
    let qualification_authorities = ext.execute_with(|| QualificationAuthorities::<Runtime>::get(curve));
    assert_eq!(active_authorities, qualification_authorities);
}

#[test]
fn test_prepare_round4_abort_on_byzantine_attack() {
    let (mut ext, tx_pool_state) = new_test_ext();

    let curve = NetworkCurve::Secp256k1;
    let raw_authorities = FixedAuthorities::get();

    let authorities: Vec<UintAuthorityId> = raw_authorities.iter().map(|&id| id.into()).collect();
    let authorities_len = authorities.len();

    ext.execute_with(|| Exodus::start_dkg_qualification(authorities, NetworkCurve::iter()));

    let dkg_rounds = [
        execute_round0_happy_path,
        execute_round1_happy_path,
        execute_round2_happy_path,
        execute_round3_happy_path,
    ];

    for dkg_round in dkg_rounds {
        run_to_next_round(&mut ext, curve);
        for authority in raw_authorities.iter() {
            assert_ok!(dkg_round(&mut ext, &tx_pool_state, curve, *authority));
        }
    }

    let bft_threshold = get_byzantium_threshold(authorities_len);
    ext.execute_with(|| {
        let mut complaints_map = BitmapByAuthority::<Runtime>::default();
        for authority in 0..=bft_threshold {
            let auth_index = authority as AuthIndex;
            let complaints = ParticipantsBitmap::<Runtime>::all_ones(authorities_len);
            complaints_map.try_insert(auth_index, complaints).unwrap();
        }
        Complaints::<Runtime>::insert(curve, complaints_map);
    });

    run_to_next_round(&mut ext, curve);

    let state = ext.execute_with(|| QualificationDkgState::<Runtime>::get(curve));
    assert!(state.is_dkg_vacant());
    assert!(state.get_indexes().is_empty());
}

#[test]
fn test_prepare_round5_abort_on_byzantine_attack() {
    let (mut ext, tx_pool_state) = new_test_ext();

    let curve = NetworkCurve::Secp256k1;
    let raw_authorities = FixedAuthorities::get();

    let authorities: Vec<UintAuthorityId> = raw_authorities.iter().map(|&id| id.into()).collect();
    let authorities_len = authorities.len();

    ext.execute_with(|| Exodus::start_dkg_qualification(authorities, NetworkCurve::iter()));

    let dkg_rounds = [
        execute_round0_happy_path,
        execute_round1_happy_path,
        execute_round2_happy_path,
        execute_round3_happy_path,
        execute_round4_happy_path,
    ];

    for dkg_round in dkg_rounds {
        run_to_next_round(&mut ext, curve);
        for authority in raw_authorities.iter() {
            assert_ok!(dkg_round(&mut ext, &tx_pool_state, curve, *authority));
        }
    }

    let bft_threshold = get_byzantium_threshold(authorities_len);
    ext.execute_with(|| {
        let mut justifications_map = BitmapByAuthority::<Runtime>::default();
        for authority in 0..=bft_threshold {
            let auth_index = authority as AuthIndex;
            let justifications = ParticipantsBitmap::<Runtime>::all_ones(authorities_len);
            justifications_map.try_insert(auth_index, justifications).unwrap();
        }
        Justifications::<Runtime>::insert(curve, justifications_map);
    });

    run_to_next_round(&mut ext, curve);

    let state = ext.execute_with(|| QualificationDkgState::<Runtime>::get(curve));
    let indexes = state.get_indexes();

    assert!(state.is_dkg_vacant());
    assert!(indexes.is_empty());
    assert_eq!(indexes.active_bits_count(), authorities_len as u32);
}

#[test]
fn test_prepare_round5_kick_on_empty_justifications() {
    let (mut ext, tx_pool_state) = new_test_ext();

    let curve = NetworkCurve::Secp256k1;
    let raw_authorities = FixedAuthorities::get();

    let authorities: Vec<UintAuthorityId> = raw_authorities.iter().map(|&id| id.into()).collect();
    let authorities_len = authorities.len();

    ext.execute_with(|| Exodus::start_dkg_qualification(authorities, NetworkCurve::iter()));

    let dkg_rounds = [
        execute_round0_happy_path,
        execute_round1_happy_path,
        execute_round2_happy_path,
        execute_round3_happy_path,
        execute_round4_happy_path,
    ];

    for dkg_round in dkg_rounds {
        run_to_next_round(&mut ext, curve);
        for authority in raw_authorities.iter() {
            assert_ok!(dkg_round(&mut ext, &tx_pool_state, curve, *authority));
        }
    }

    ext.execute_with(|| {
        let mut complaints_map = BitmapByAuthority::<Runtime>::default();
        let mut justifications_map = BitmapByAuthority::<Runtime>::default();

        let auth: Vec<AuthIndex> = raw_authorities.iter().map(|&id| id as AuthIndex).collect();

        let mut complaints_from_1 = ParticipantsBitmap::<Runtime>::empty(authorities_len);
        complaints_from_1.insert(auth[0]);
        complaints_from_1.insert(auth[3]);
        complaints_map.try_insert(auth[1], complaints_from_1).unwrap();

        let mut complaints_from_2 = ParticipantsBitmap::<Runtime>::empty(authorities_len);
        complaints_from_2.insert(auth[0]);
        complaints_map.try_insert(auth[2], complaints_from_2).unwrap();

        complaints_map.try_insert(auth[0], ParticipantsBitmap::<Runtime>::empty(authorities_len)).unwrap();
        complaints_map.try_insert(auth[3], ParticipantsBitmap::<Runtime>::empty(authorities_len)).unwrap();

        let mut justifications_from_0 = ParticipantsBitmap::<Runtime>::empty(authorities_len);
        justifications_from_0.insert(auth[1]);
        justifications_from_0.insert(auth[2]);
        justifications_map.try_insert(auth[0], justifications_from_0).unwrap();

        let mut justifications_from_3 = ParticipantsBitmap::<Runtime>::empty(authorities_len);
        justifications_from_3.insert(auth[0]);
        justifications_map.try_insert(auth[3], justifications_from_3).unwrap();

        justifications_map.try_insert(auth[1], ParticipantsBitmap::<Runtime>::empty(authorities_len)).unwrap();
        justifications_map.try_insert(auth[2], ParticipantsBitmap::<Runtime>::empty(authorities_len)).unwrap();

        Complaints::<Runtime>::insert(curve, complaints_map);
        Justifications::<Runtime>::insert(curve, justifications_map);
    });

    run_to_next_round(&mut ext, curve);

    let state = ext.execute_with(|| QualificationDkgState::<Runtime>::get(curve));
    let indexes = state.get_indexes();

    for authority in raw_authorities.iter() {
        let auth_index = *authority as AuthIndex;

        if auth_index == 3 {
            assert!(!indexes.contains(auth_index));
            continue;
        }

        assert!(indexes.contains(auth_index));
    }

    assert_eq!(state.get_phase(), DkgPhase::Round5);
    assert_eq!(indexes.active_bits_count(), authorities_len as u32);
}

#[test]
fn test_merkle_tree_success_power_of_two_for_shares() {
    let network_curve = NetworkCurve::Secp256k1;
    let indices = (1..=32).collect::<Vec<AuthIndex>>();
    let shares = generate_mock_secp256k1_shares(&indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &shares,
        |share| share.serialize().map_err(|_| ExodusError::SerializationError),
    ).unwrap();

    let root = *tree.last().unwrap();

    for (identifier, _) in shares.iter() {
        let auth_index = NetworkCurve::convert_identifier_to_index(
            &identifier.serialize(),
        ).unwrap();

        let (shares_bytes, proof) =
            NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
                &shares,
                &tree,
                auth_index,
                |share| share.serialize().map_err(|_| ExodusError::SerializationError)
            ).unwrap();

        let result = network_curve.verify_merkle_proof(
            &shares_bytes,
            &proof,
            root,
            auth_index,
        );

        assert!(result.is_ok());
    }
}

#[test]
fn test_merkle_tree_success_power_of_two_for_scalar() {
    let network_curve = NetworkCurve::Secp256k1;
    let indices = (1..=32).collect::<Vec<AuthIndex>>();
    let scalars = generate_mock_secp256k1_scalars(&indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &scalars,
        |scalar| {
            let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
            Ok(serialized.as_ref().to_vec())
        }
    ).unwrap();

    let root = *tree.last().unwrap();

    for (identifier, _) in scalars.iter() {
        let auth_index = NetworkCurve::convert_identifier_to_index(
            &identifier.serialize(),
        ).unwrap();

        let (shares_bytes, proof) =
            NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
                &scalars,
                &tree,
                auth_index,
                |scalar| {
                    let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
                    Ok(serialized.as_ref().to_vec())
                }
            ).unwrap();

        let result = network_curve.verify_merkle_proof(
            &shares_bytes,
            &proof,
            root,
            auth_index,
        );

        assert!(result.is_ok());
    }
}

#[test]
fn test_random_sparse_combinations_stress_for_shares() {
    let network_curve = NetworkCurve::Secp256k1;
    let mut rng = ChaCha20Rng::from_entropy();

    let mut unique_indices = BTreeSet::new();
    while unique_indices.len() < 15 {
        unique_indices.insert(rng.gen_range(1..150) as AuthIndex);
    }
    let indices: Vec<AuthIndex> = unique_indices.into_iter().collect();
    let shares = generate_mock_secp256k1_shares(&indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &shares,
        |share| share.serialize().map_err(|_| ExodusError::SerializationError),
    ).unwrap();

    let root = *tree.last().unwrap();

    for &auth_index in &indices {
        let (shares_bytes, proof) =
            NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
                &shares,
                &tree,
                auth_index,
                |share| share.serialize().map_err(|_| ExodusError::SerializationError),
            ).unwrap();

        let result = network_curve.verify_merkle_proof(
            &shares_bytes,
            &proof,
            root,
            auth_index,
        );
        assert!(result.is_ok());
    }
}

#[test]
fn test_random_sparse_combinations_stress_for_scalars() {
    let network_curve = NetworkCurve::Secp256k1;
    let mut rng = ChaCha20Rng::from_entropy();

    let mut unique_indices = BTreeSet::new();
    while unique_indices.len() < 15 {
        unique_indices.insert(rng.gen_range(1..150) as AuthIndex);
    }
    let indices: Vec<AuthIndex> = unique_indices.into_iter().collect();
    let scalars = generate_mock_secp256k1_scalars(&indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &scalars,
        |scalar| {
            let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
            Ok(serialized.as_ref().to_vec())
        }
    ).unwrap();

    let root = *tree.last().unwrap();

    for &auth_index in &indices {
        let (scalars_bytes, proof) =
            NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
                &scalars,
                &tree,
                auth_index,
                |scalar| {
                    let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
                    Ok(serialized.as_ref().to_vec())
                }
            ).unwrap();

        let result = network_curve.verify_merkle_proof(
            &scalars_bytes,
            &proof,
            root,
            auth_index,
        );
        assert!(result.is_ok());
    }
}

#[test]
fn test_participants_with_high_index_gap_for_shares() {
    let network_curve = NetworkCurve::Secp256k1;
    let indices = vec![1, 1023];
    let shares = generate_mock_secp256k1_shares(&indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &shares,
        |share| share.serialize().map_err(|_| ExodusError::SerializationError),
    ).unwrap();

    let root = *tree.last().unwrap();

    let (shares_bytes, proof) =
        NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
            &shares,
            &tree,
            1023,
            |share| share.serialize().map_err(|_| ExodusError::SerializationError),
        ).unwrap();

    assert_eq!(proof.len(), 10);

    let result = network_curve.verify_merkle_proof(
        &shares_bytes,
        &proof,
        root,
        1023,
    );
    assert!(result.is_ok());
}

#[test]
fn test_participants_with_high_index_gap_for_scalars() {
    let network_curve = NetworkCurve::Secp256k1;
    let indices = vec![1, 1023];
    let scalars = generate_mock_secp256k1_scalars(&indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &scalars,
        |scalar| {
            let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
            Ok(serialized.as_ref().to_vec())
        }
    ).unwrap();

    let root = *tree.last().unwrap();

    let (scalars_bytes, proof) =
        NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
            &scalars,
            &tree,
            1023,
            |scalar| {
                let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
                Ok(serialized.as_ref().to_vec())
            }
        ).unwrap();

    assert_eq!(proof.len(), 10);

    let result = network_curve.verify_merkle_proof(
        &scalars_bytes,
        &proof,
        root,
        1023,
    );
    assert!(result.is_ok());
}

#[test]
fn test_regression_consecutive_vs_position_for_shares() {
    let network_curve = NetworkCurve::Secp256k1;
    let custom_indices = vec![2, 3];
    let shares = generate_mock_secp256k1_shares(&custom_indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &shares,
        |scalar| scalar.serialize().map_err(|_| ExodusError::SerializationError),
    ).unwrap();

    let root = *tree.last().unwrap();

    let (shares_bytes, proof) =
        NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
            &shares,
            &tree,
            2,
            |scalar| scalar.serialize().map_err(|_| ExodusError::SerializationError),
        ).unwrap();

    let result = network_curve.verify_merkle_proof(&shares_bytes, &proof, root, 2);
    assert!(result.is_ok());
}

#[test]
fn test_regression_consecutive_vs_position_for_scalars() {
    let network_curve = NetworkCurve::Secp256k1;
    let custom_indices = vec![2, 3];
    let scalars = generate_mock_secp256k1_scalars(&custom_indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &scalars,
        |scalar| {
            let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
            Ok(serialized.as_ref().to_vec())
        }
    ).unwrap();

    let root = *tree.last().unwrap();

    let (scalars_bytes, proof) =
        NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
            &scalars,
            &tree,
            2,
            |scalar| {
                let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
                Ok(serialized.as_ref().to_vec())
            }
        ).unwrap();

    let result = network_curve.verify_merkle_proof(&scalars_bytes, &proof, root, 2);
    assert!(result.is_ok());
}

#[test]
fn test_attack_with_unknown_authority_index_for_shares() {
    let network_curve = NetworkCurve::Secp256k1;
    let custom_indices = vec![1, 2, 3];
    let shares = generate_mock_secp256k1_shares(&custom_indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &shares,
        |scalar| scalar.serialize().map_err(|_| ExodusError::SerializationError),
    ).unwrap();

    let root = *tree.last().unwrap();

    let (shares_bytes, proof) =
        NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
            &shares,
            &tree,
            2,
            |scalar| scalar.serialize().map_err(|_| ExodusError::SerializationError),
        ).unwrap();

    let result = network_curve.verify_merkle_proof(&shares_bytes, &proof, root, 4);
    assert!(result.is_err());
}

#[test]
fn test_attack_with_unknown_authority_index_for_scalars() {
    let network_curve = NetworkCurve::Secp256k1;
    let custom_indices = vec![1, 2, 3];
    let scalars = generate_mock_secp256k1_scalars(&custom_indices);

    let tree = NetworkCurve::build_merkle_tree::<Secp256K1Sha256, _, _>(
        &scalars,
        |scalar| {
            let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
            Ok(serialized.as_ref().to_vec())
        }
    ).unwrap();

    let root = *tree.last().unwrap();

    let (scalars_bytes, proof) =
        NetworkCurve::generate_merkle_proof_from_tree::<Secp256K1Sha256, _, _>(
            &scalars,
            &tree,
            2,
            |scalar| {
                let serialized = <Secp256K1ScalarField as Field>::serialize(scalar);
                Ok(serialized.as_ref().to_vec())
            }
        ).unwrap();

    let result = network_curve.verify_merkle_proof(&scalars_bytes, &proof, root, 4);
    assert!(result.is_err());
}

#[test]
fn curve_wrapper_for_dkg_works() {
    let mut rng = ChaCha20Rng::from_seed(Default::default());
	let curve = NetworkCurve::Secp256k1;
    let dkg_index = 69 as DkgIndex;

    // NOTE: iterate over all NetworkCurve enum entries.
    // For now it is fine, because we are using only one curve

    let mut lost_participants         = BTreeSet::new();
    let mut round1_secret_packages    = BTreeMap::new();
    let mut received_round1_packages  = BTreeMap::new();
    let mut round2_secret_packages    = BTreeMap::new();
    let mut round2_encrypted_ciphers  = BTreeMap::new();
    let mut round2_decrypted_packages = BTreeMap::new();
    let mut secret_key_packages       = BTreeMap::new();
    let mut public_verifying_shares   = BTreeMap::new();
    let mut public_verifying_keys     = BTreeMap::new();

    // throttle first participant
    lost_participants.insert(0u16);

    for index in 0..MAX_SIGNERS {
        if lost_participants.contains(&index) { continue; }

        let (round1_secret_package_bytes, round1_package_bytes) =
            curve.dkg_part1(index, MAX_SIGNERS, MIN_SIGNERS, &mut rng).unwrap();

        let number_of_coefficients =
            curve.dkg_verify_proof_of_knowledge(index, &round1_package_bytes).unwrap();
        assert!(number_of_coefficients == MIN_SIGNERS);

        round1_secret_packages.insert(index, round1_secret_package_bytes);
        received_round1_packages.insert(index, round1_package_bytes);
    }

    assert!(round1_secret_packages.len() == (MAX_SIGNERS - 1) as usize);
    assert!(received_round1_packages.len() == (MAX_SIGNERS - 1) as usize);

    // throttle last participant
    lost_participants.insert(MAX_SIGNERS - 1);

    for sender_index in 0..MAX_SIGNERS {
        if lost_participants.contains(&sender_index) { continue; }

        let sender_round1_package = received_round1_packages
            .remove(&sender_index)
            .unwrap();

        round1_secret_packages.get_mut(&sender_index)
            .map(|round1_secret_package| {
                curve.dkg_narrow_round_max_signers(
                    round1_secret_package,
                    received_round1_packages.len() as AuthIndex + 1,
                    1u8, // narrow max signers for round1
                )
                    .map(|new_secret_package| {
                        *round1_secret_package = new_secret_package;
                    })
                    .ok()
            })
            .flatten()
            .unwrap();

        let round1_secret_package_bytes = round1_secret_packages
            .get(&sender_index)
            .unwrap();

        let (round2_secret_package_bytes, round2_packages_bytes) =
            curve.dkg_part2(&round1_secret_package_bytes, &received_round1_packages)
                .unwrap();

        assert_eq!(
            round2_packages_bytes.len(),
            (MAX_SIGNERS as usize) - lost_participants.len(),
        );

        round2_secret_packages.insert(sender_index, round2_secret_package_bytes);

        let insertion_outcome = received_round1_packages
            .insert(sender_index, sender_round1_package);
        assert!(insertion_outcome.is_none());

        for (receiver_index, round2_package_bytes) in round2_packages_bytes {
            let additional_info = curve.ecdh_prepare_additional_info(
                sender_index,
                receiver_index,
                dkg_index,
                sp_std::marker::PhantomData::<EncryptionData<Runtime>>,
            );

            let round1_receiver_package_bytes = received_round1_packages
                .get(&receiver_index)
                .unwrap();

            let encryption_cipher = curve.ecdh_get_cipher_from_keys(
                &round1_receiver_package_bytes,
                &round1_secret_package_bytes,
                &additional_info,
                sp_std::marker::PhantomData::<EncryptionData<Runtime>>,
            ).unwrap();

            let encryption_data: EncryptionData<Runtime> = curve
                .ecdh_encrypt_package(
                    &encryption_cipher,
                    &additional_info,
                    &round2_package_bytes,
                    &mut rng,
                ).unwrap();

            round2_encrypted_ciphers
                .entry(receiver_index)
                .or_insert_with(BTreeMap::new)
                .insert(sender_index, encryption_data);
        }
    }

    // throttle someone in the middle
    let malicious_index = MAX_SIGNERS / 2;
    lost_participants.insert(malicious_index);

    for receiver_index in 0..MAX_SIGNERS {
        if lost_participants.contains(&receiver_index) { continue; }

        let round1_secret_package_bytes = round1_secret_packages
            .get(&receiver_index)
            .unwrap();

        let my_encrypted_messages = round2_encrypted_ciphers
            .get(&receiver_index)
            .unwrap();

        assert_eq!(
            my_encrypted_messages.len(),
            (MAX_SIGNERS as usize) - lost_participants.len(),
        );

        for (&sender_index, encrypted_message) in my_encrypted_messages.iter() {
            if lost_participants.contains(&sender_index) { continue; }

            let sender_round1_package_bytes = received_round1_packages
                .get(&sender_index)
                .unwrap();

            let additional_info = curve.ecdh_prepare_additional_info(
                sender_index,
                receiver_index,
                dkg_index,
                sp_std::marker::PhantomData::<EncryptionData<Runtime>>,
            );

            let encryption_cipher = curve.ecdh_get_cipher_from_keys(
                &sender_round1_package_bytes,
                &round1_secret_package_bytes,
                &additional_info,
                sp_std::marker::PhantomData::<EncryptionData<Runtime>>,
            ).unwrap();

            let decrypted_private_package = curve.ecdh_decrypt_package(
                &encryption_cipher,
                &encrypted_message,
                &additional_info,
            ).unwrap();

            curve.dkg_verify_private_package(
                receiver_index,
                &sender_round1_package_bytes,
                &decrypted_private_package,
            ).unwrap();

            round2_decrypted_packages
                .entry(receiver_index)
                .or_insert_with(BTreeMap::new)
                .insert(sender_index, decrypted_private_package);
        }

        let decrypted_packages_len = round2_decrypted_packages
            .get(&receiver_index)
            .map(|inner_map| inner_map.len())
            .unwrap_or_default();

        assert_eq!(
            decrypted_packages_len,
            (MAX_SIGNERS as usize) - 1 - lost_participants.len(),
        );
    }

    assert_eq!(
        round2_decrypted_packages.len(),
        (MAX_SIGNERS as usize) - lost_participants.len(),
    );

    for index in 0..MAX_SIGNERS {
        if lost_participants.contains(&index) { continue; }

        let round2_packages_bytes = round2_decrypted_packages
            .get(&index)
            .unwrap();

        round2_secret_packages.get_mut(&index)
            .map(|round2_secret_package| {
                curve.dkg_narrow_round_max_signers(
                    round2_secret_package,
                    MAX_SIGNERS - lost_participants.len() as AuthIndex,
                    2u8, // narrow max signers for round2
                )
                    .map(|new_secret_package| {
                        *round2_secret_package = new_secret_package;
                    })
                    .ok()
            })
            .flatten()
            .unwrap();

        let round2_secret_package_bytes = round2_secret_packages
            .get(&index)
            .unwrap();

        let round1_packages = received_round1_packages
            .iter()
            .filter(|(&i, _)| i != index && !lost_participants.contains(&i))
            .map(|(i, package)| (*i, package.clone()))
            .collect::<BTreeMap<AuthIndex, Vec<u8>>>();

        let (secret_key_package, verifying_key, merkle_proof, merkle_root) =
            curve.dkg_part3(
                index,
                round2_secret_package_bytes,
                &round1_packages,
                round2_packages_bytes,
            ).unwrap();

        let verifying_share = curve.dkg_derive_verifying_share(&secret_key_package).unwrap();

        curve.verify_merkle_proof(&verifying_share, &merkle_proof, merkle_root, index).unwrap();

        secret_key_packages.insert(index, secret_key_package);
        public_verifying_shares.insert(index, verifying_share);
        public_verifying_keys.insert(index, verifying_key);
    }

    let group_public_key = public_verifying_keys.values().next().unwrap();
    let all_equal = public_verifying_keys.values().all(|val| val == group_public_key);
    assert!(all_equal, "Consensus on group pubkey not reached");
    assert!(public_verifying_shares.keys().count() == (MAX_SIGNERS as usize) - lost_participants.len());

    let mut nonces_map = BTreeMap::new();
    let mut commitments_map = BTreeMap::new();
    let mut signing_packages = BTreeMap::new();

    let mut group_commitments = BTreeMap::new();
    let mut self_binding_factors = BTreeMap::new();
    let mut binding_factors_proofs = BTreeMap::new();
    let mut binding_factor_roots = BTreeMap::new();
    let mut session_commitments = BTreeMap::new();

    let message = "message to sign".as_bytes();
    let mut number_of_valid_signatures = 0;
    let mut session_id = 0;
    let mut member_id = 0;

    let max_possible_sessions = MAX_SIGNERS
        .saturating_sub(MIN_SIGNERS)
        .saturating_add(1)
        .saturating_sub(lost_participants.len() as AuthIndex);

    loop {
        if member_id >= MAX_SIGNERS { member_id = 0; }
        if lost_participants.contains(&member_id) {
            member_id += 1;
            continue;
        }

        let session_members = nonces_map.get(&session_id)
            .map(|members: &BTreeMap<AuthIndex, Vec<u8>>| members.len())
            .unwrap_or_default();

        if session_members >= MIN_SIGNERS as usize { session_id += 1; }
        if session_id == max_possible_sessions { break; }

        let secret_share = secret_key_packages.get(&member_id).unwrap();
        let (nonce, hiding, binding) = curve
            .generate_nonce_commitment(&secret_share, &mut rng)
            .unwrap();

        nonces_map
            .entry(session_id)
            .or_insert_with(BTreeMap::new)
            .insert(member_id, nonce);

        commitments_map
            .entry(session_id)
            .or_insert_with(BTreeMap::new)
            .insert(member_id, (hiding, binding));

        member_id += 1;
    }

    commitments_map.iter().for_each(|(&session_id, raw_commitments)| {
        let commitments = raw_commitments
            .iter()
            .map(|(&auth_index, (hiding, binding))| {
                (auth_index, (hiding.as_ref(), binding.as_ref()))
            })
            .collect::<BTreeMap<AuthIndex, _>>();

        raw_commitments.keys().for_each(|&auth_index| {
            let (group_commitment, self_factor, factors_proof, factor_root) = curve
                .generate_group_commitment(
                    auth_index,
                    &commitments,
                    &group_public_key,
                    &message,
                )
                .unwrap();

            group_commitments
                .entry(session_id)
                .or_insert_with(BTreeMap::new)
                .insert(auth_index, group_commitment);

            self_binding_factors
                .entry(session_id)
                .or_insert_with(BTreeMap::new)
                .insert(auth_index, self_factor);

            binding_factors_proofs
                .entry(session_id)
                .or_insert_with(BTreeMap::new)
                .insert(auth_index, factors_proof);

            binding_factor_roots
                .entry(session_id)
                .or_insert_with(BTreeMap::new)
                .insert(auth_index, factor_root);
        });

        let signing_package = curve
            .init_signing_package(&commitments, &message)
            .unwrap();

        signing_packages.insert(session_id, signing_package);
        session_commitments.insert(session_id, raw_commitments);
    });

    group_commitments.iter().for_each(|(_, commitments)| {
        let first_commitment = commitments.values().next().unwrap();
        let all_equal = commitments.values().all(|val| val == first_commitment);
        assert!(all_equal, "Group commitments inside ROAST session differs.");
    });

    for session_id in commitments_map.keys() {
        let signing_package = match signing_packages.get(session_id) {
            Some(signing_package) => signing_package,
            None => continue,
        };

        let nonces = match nonces_map.get(&session_id) {
            Some(nonces) => nonces,
            None => continue,
        };

        let current_raw_commitments = session_commitments.get(session_id).unwrap();
        let participants_list: Vec<AuthIndex> = current_raw_commitments.keys().cloned().collect();
        let mut accumulated_signature_scalar = vec![0u8; 32];

        for (sender_index, signer_nonces) in nonces.iter() {
            let secret_share = secret_key_packages.get(&sender_index).unwrap();
            let verifying_share = public_verifying_shares.get(sender_index).unwrap();

            let group_commitment = group_commitments.get(session_id).unwrap().get(sender_index).unwrap();
            let self_factor = self_binding_factors.get(session_id).unwrap().get(sender_index).unwrap();
            let factors_proof = binding_factors_proofs.get(session_id).unwrap().get(sender_index).unwrap();
            let factor_root = binding_factor_roots.get(session_id).unwrap().get(sender_index).unwrap();

            curve.verify_merkle_proof(
                self_factor.as_ref(),
                factors_proof.as_ref(),
                *factor_root,
                *sender_index,
            ).unwrap();

            let signature_share = curve.partial_sign_message(
                signing_package,
                signer_nonces,
                secret_share,
            ).unwrap();

            let (raw_hiding, raw_binding) = session_commitments
                .get(session_id)
                .unwrap()
                .get(sender_index)
                .unwrap();

            curve.verify_signature_share(
                *sender_index,
                participants_list.iter().cloned(),
                signature_share.as_ref(),
                self_factor.as_ref(),
                raw_hiding.as_ref(),
                raw_binding.as_ref(),
                group_commitment.as_ref(),
                verifying_share,
                &group_public_key,
                &message,
            ).unwrap();

            accumulated_signature_scalar = curve.accumulate_signature_scalar(
                &accumulated_signature_scalar,
                signature_share.as_ref(),
            ).unwrap();
        }

        let any_sender = nonces.keys().next().unwrap();
        let current_session_group_commitment = group_commitments
            .get(session_id)
            .unwrap()
            .get(any_sender)
            .unwrap();

        if curve
            .is_signature_valid(
                &group_public_key,
                current_session_group_commitment.as_ref(),
                &accumulated_signature_scalar,
                &message,
            )
            .is_ok() {
                number_of_valid_signatures += 1;
            }
    }

    assert_eq!(number_of_valid_signatures, commitments_map.len());
}

#[test]
fn test_common_way_of_frost_usage() {
    use frost_secp256k1 as frost;

    let mut rng = ChaCha20Rng::from_seed(Default::default());

    // Key generation, Round 1

    // Keep track of each participant's round 1 secret package.  In
    // practice each participant will keep its copy; no one will have
    // all the participant's packages.
    let mut round1_secret_packages = BTreeMap::new();

    // Keep track of all round 1 packages sent to the given
    // participant.  This is used to simulate the broadcast; in
    // practice the packages will be sent through some communication
    // channel.
    let mut received_round1_packages = BTreeMap::new();

    // For each participant, perform the first part of the DKG
    // protocol. In practice, each participant will perform this on
    // their own environments.
    for participant_index in 1..=MAX_SIGNERS {
        let participant_identifier = participant_index.try_into().expect("should be nonzero");
        let (round1_secret_package, round1_package) = frost::keys::dkg::part1(
            participant_identifier,
            MAX_SIGNERS,
            MIN_SIGNERS,
            &mut rng,
        ).unwrap();

        // store the participants secret Store the participant's
        // secret package for later use. In practice each participant
        // will store it in their own environment.
        round1_secret_packages.insert(
            participant_identifier,
            round1_secret_package,
        );

        // "Send" the round 1 package to all other participants. In
        // this test this is simulated using a BTreeMap; in practice
        // this will be sent through some communication channel.
        for receiver_participant_index in 1..=MAX_SIGNERS {
            if receiver_participant_index == participant_index {
                continue;
            }

            let receiver_participant_identifier: frost::Identifier = receiver_participant_index
                .try_into()
                .expect("should be nonzero");

            received_round1_packages
                .entry(receiver_participant_identifier)
                .or_insert_with(BTreeMap::new)
                .insert(participant_identifier, round1_package.clone());
        }
    }

    // Ket generation, Round 2

    // Keep track of each participant's round 2 secret package. In
    // practice each participant will keep its copy; no one will have
    // all the participant's packages.
    let mut round2_secret_packages = BTreeMap::new();

    // Keep track of all round 2 packages sent to the given
    // participant. This is used to simulate the broadcast; in
    // practice the packages will be sent through some communication
    // channel.
    let mut received_round2_packages = BTreeMap::new();

    // For each participant, perform the second part of the DKG
    // protocol. In practice, each participant will perform this on
    // their own environments.
    for participant_index in 1..=MAX_SIGNERS {
        let participant_identifier = participant_index.try_into().expect("should be nonzero");
        let round1_secret_package = round1_secret_packages
            .remove(&participant_identifier)
            .unwrap();

        let round1_package = &received_round1_packages[&participant_identifier];
        let (round2_secret_package, round2_packages) =
            frost::keys::dkg::part2(round1_secret_package, round1_package).unwrap();

        // Store the participant's secret package for later use. In
        // practice each participant will store in their own
        // environment.
        round2_secret_packages.insert(
            participant_identifier,
            round2_secret_package,
        );

        // "Send" the round 2 package to all other participants. In
        // this test this is simulated using a BTreeMap; in practice
        // this will be sent through some communication channel. Note
        // that, in contrast to the previous part, here each other
        // participant gets its own specific package.
        for (receiver_identifier, round2_package) in round2_packages {
            received_round2_packages
                .entry(receiver_identifier)
                .or_insert_with(BTreeMap::new)
                .insert(participant_identifier, round2_package);
        }
    }

    // Key generation, final computation

    // Keep track of each participant's long-lived key package. In
    // practice each participant will keep it's own copy; no one will
    // have all the participant's packages.
    let mut key_packages = BTreeMap::new();

    // Keep track of each participant's public key package.  In
    // practice, if there is a Coordinator, only they need to store
    // the set. If there is not, then all candidates must store their
    // own sets. All participants will have the same exact public key
    // package.
    let mut pubkey_packages = BTreeMap::new();

    // For each participant, perform the third part of the DKG
    // protocol. In practice, each participant will perform this on
    // their own environments.
    for participant_index in 1..=MAX_SIGNERS {
        let participant_identifier = participant_index.try_into().expect("should be nonzero");

        let round2_secret_package = &round2_secret_packages[&participant_identifier];
        let round1_packages = &received_round1_packages[&participant_identifier];
        let round2_packages = &received_round2_packages[&participant_identifier];

        let (key_package, pubkey_package) = frost::keys::dkg::part3(
            round2_secret_package,
            round1_packages,
            round2_packages,
        ).unwrap();

        key_packages.insert(participant_identifier, key_package);
        pubkey_packages.insert(participant_identifier, pubkey_package);
    }

    let pubkey_package = pubkey_packages.values().next().expect("should be one");
    let mut nonces_map = BTreeMap::new();
    let mut commitments_map = BTreeMap::new();

    // Round 1, generate nonces and signing commitments
    // In practice, each iteration of this loop will be executed by its respective participant.
    for participant_index in 1..=MIN_SIGNERS {
        let participant_identifier = participant_index.try_into().expect("should be nonzero");

        let key_package = &key_packages[&participant_identifier];

        // Generate one (1) nonce and one SigningCommitments instance
        // for each participant, up to _threshold_.
        let (nonces, commitments) = frost::round1::commit(
            key_package.signing_share(),
            &mut rng,
        );

        // In practice, the nonces must be kept by the participant to
        // use in the next round, while the commitment must be sent to
        // the coordinator (or to every other participant if there is
        // no coordinator) using an authenticated channel.
        nonces_map.insert(participant_identifier, nonces);
        commitments_map.insert(participant_identifier, commitments);
    }

    // This is what the signature aggregator / coordinator needs to do:
    // - decide what message to sign
    // - take one (unused) commitment per signing participant
    let mut signature_shares = BTreeMap::new();
    let message = "message to sign".as_bytes();
    let signing_package = frost::SigningPackage::new(commitments_map, message);

    // Round 2, each participant generates their signature shares

    // In practice, each iteration of this loop will be executed by
    // its respective participant.
    for participant_identifier in nonces_map.keys() {
        let key_package = &key_packages[participant_identifier];

        let nonces = &nonces_map[participant_identifier];

        // Each participant generates their signature share.
        let signature_share = frost::round2::sign(
            &signing_package,
            nonces,
            key_package,
        ).unwrap();

        // In practice, the signature share must be sent to the
        // Coordinator using an authenticated channel.
        signature_shares.insert(*participant_identifier, signature_share);
    }

    // Aggregation: collects signing shares from all participants
    // Aggregate (also verifies the signature shares)
    let group_signature = frost::aggregate(
        &signing_package,
        &signature_shares,
        &pubkey_package,
    ).unwrap();

    // Check that the threshold signature can be verified by the group
    // public key (the verification key).
    let is_signature_valid = pubkey_package
        .verifying_key()
        .verify(message, &group_signature)
        .is_ok();

    assert!(is_signature_valid);
}