use super::{exp_acc, ElementHasher, HashFn};
use core::{convert::TryInto, ops::Range};
use winterfell::math::{fields::f64::BaseElement, FieldElement, StarkField};

mod digest;
pub use digest::ElementDigest;

#[cfg(test)]
mod tests;

mod mds_freq;
use mds_freq::mds_multiply_freq;

// CONSTANTS
// ================================================================================================

/// Sponge state is set to 12 field elements or 768 bytes; 8 elements are reserved for rate and
/// the remaining 4 elements are reserved for capacity.
const STATE_WIDTH: usize = 12;

/// The rate portion of the state is located in elements 4 through 11.
const RATE_RANGE: Range<usize> = 4..12;
const RATE_WIDTH: usize = RATE_RANGE.end - RATE_RANGE.start;

const INPUT1_RANGE: Range<usize> = 4..8;
const INPUT2_RANGE: Range<usize> = 8..12;

/// The capacity portion of the state is located in elements 0, 1, 2, and 3.
const CAPACITY_RANGE: Range<usize> = 0..4;

/// The output of the hash function is a digest which consists of 4 field elements or 32 bytes.
///
/// The digest is returned from state elements 4, 5, 6, and 7 (the first four elements of the
/// rate portion).
const DIGEST_RANGE: Range<usize> = 4..8;
const DIGEST_SIZE: usize = DIGEST_RANGE.end - DIGEST_RANGE.start;

/// The number of rounds is set to 7 to target 128-bit security level
const NUM_ROUNDS: usize = 7;

/// S-Box and Inverse S-Box powers;
///
/// The constants are defined for tests only because the exponentiations in the code are unrolled
/// for efficiency reasons.
#[cfg(test)]
const ALPHA: u64 = 7;
#[cfg(test)]
const INV_ALPHA: u64 = 10540996611094048183;

// HASHER IMPLEMENTATION
// ================================================================================================

/// Implementation of [Hasher] trait for Rescue Prime Optimized (RPO) hash function with 256-bit output.
///
/// The hash function is implemented according to the Rescue Prime Optimized
/// [specifications](https://github.org/aszepieniec/rpo/)
///
/// The parameters used to instantiate the function are:
/// * Field: 64-bit prime field with modulus 2^64 - 2^32 + 1.
/// * State width: 12 field elements.
/// * Capacity size: 4 field elements.
/// * Number of founds: 7.
/// * S-Box degree: 7.
///
/// The above parameters target 128-bit security level. The digest consists of four field elements
/// and it can be serialized into 32 bytes (256 bits).
///
/// ## Hash output consistency
/// Functions [hash_elements()](Rpo::hash_elements), [merge()](Rpo::merge), and
/// [merge_with_int()](Rpo::merge_with_int) are internally consistent. That is, computing
/// a hash for the same set of elements using these functions will always produce the same
/// result. For example, merging two digests using [merge()](Rpo::merge) will produce the
/// same result as hashing 8 elements which make up these digests using
/// [hash_elements()](Rpo::hash_elements) function.
///
/// However, [hash()](Rpo::hash) function is not consistent with functions mentioned above.
/// For example, if we take two field elements, serialize them to bytes and hash them using
/// [hash()](Rpo::hash), the result will differ from the result obtained by hashing these
/// elements directly using [hash_elements()](Rpo::hash_elements) function. The reason for
/// this difference is that [hash()](Rpo::hash) function needs to be able to handle
/// arbitrary binary strings, which may or may not encode valid field elements - and thus,
/// deserialization procedure used by this function is different from the procedure used to
/// deserialize valid field elements.
///
/// Thus, if the underlying data consists of valid field elements, it might make more sense
/// to deserialize them into field elements and then hash them using
/// [hash_elements()](Rpo::hash_elements) function rather then hashing the serialized bytes
/// using [hash()](Rpo::hash) function.
pub struct Rpo();

impl HashFn for Rpo {
    type Digest = ElementDigest;

    fn hash(bytes: &[u8]) -> Self::Digest {
        // compute the number of elements required to represent the string; we will be processing
        // the string in 7-byte chunks, thus the number of elements will be equal to the number
        // of such chunks (including a potential partial chunk at the end).
        let num_elements = if bytes.len() % 7 == 0 {
            bytes.len() / 7
        } else {
            bytes.len() / 7 + 1
        };

        // initialize state to all zeros, except for the first element of the capacity part, which
        // is set to the number of elements to be hashed. this is done so that adding zero elements
        // at the end of the list always results in a different hash.
        let mut state = [BaseElement::ZERO; STATE_WIDTH];
        state[CAPACITY_RANGE.start] = BaseElement::new(num_elements as u64);

        // break the string into 7-byte chunks, convert each chunk into a field element, and
        // absorb the element into the rate portion of the state. we use 7-byte chunks because
        // every 7-byte chunk is guaranteed to map to some field element.
        let mut i = 0;
        let mut buf = [0_u8; 8];
        for chunk in bytes.chunks(7) {
            if i < num_elements - 1 {
                buf[..7].copy_from_slice(chunk);
            } else {
                // if we are dealing with the last chunk, it may be smaller than 7 bytes long, so
                // we need to handle it slightly differently. we also append a byte with value 1
                // to the end of the string; this pads the string in such a way that adding
                // trailing zeros results in different hash
                let chunk_len = chunk.len();
                buf = [0_u8; 8];
                buf[..chunk_len].copy_from_slice(chunk);
                buf[chunk_len] = 1;
            }

            // convert the bytes into a field element and absorb it into the rate portion of the
            // state; if the rate is filled up, apply the Rescue permutation and start absorbing
            // again from zero index.
            state[RATE_RANGE.start + i] += BaseElement::new(u64::from_le_bytes(buf));
            i += 1;
            if i % RATE_WIDTH == 0 {
                Self::apply_permutation(&mut state);
                i = 0;
            }
        }

        // if we absorbed some elements but didn't apply a permutation to them (would happen when
        // the number of elements is not a multiple of RATE_WIDTH), apply the Rescue permutation.
        // we don't need to apply any extra padding because we injected total number of elements
        // in the input list into the capacity portion of the state during initialization.
        if i > 0 {
            Self::apply_permutation(&mut state);
        }

        // return the first 4 elements of the state as hash result
        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
    }

    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
        // initialize the state by copying the digest elements into the rate portion of the state
        // (8 total elements), and set the capacity elements to 0.
        let mut state = [BaseElement::ZERO; STATE_WIDTH];
        state[RATE_RANGE].copy_from_slice(Self::Digest::digests_as_elements(values));

        // apply the RPO permutation and return the first four elements of the state
        Self::apply_permutation(&mut state);
        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
    }

    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
        // initialize the state as follows:
        // - seed is copied into the first 4 elements of the rate portion of the state.
        // - if the value fits into a single field element, copy it into the fifth rate element
        //   and set the sixth rate element and first capacity element to 1.
        // - if the value doesn't fit into a single field element, split it into two field
        //   elements, copy them into rate elements 5 and 6, and set the seventh rate element
        //   and first capacity element to 1.
        let mut state = [BaseElement::ZERO; STATE_WIDTH];
        state[INPUT1_RANGE].copy_from_slice(seed.as_elements());
        state[INPUT2_RANGE.start] = BaseElement::new(value);
        if value < BaseElement::MODULUS {
            state[INPUT2_RANGE.start + 1] = BaseElement::ONE;
        } else {
            state[INPUT2_RANGE.start + 1] = BaseElement::new(value / BaseElement::MODULUS);
            state[INPUT2_RANGE.start + 2] = BaseElement::ONE;            
        }

        // common padding for both cases
        state[CAPACITY_RANGE.start] = BaseElement::ONE;
        
        // apply the RPO permutation and return the first four elements of the state
        Self::apply_permutation(&mut state);
        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
    }
}

impl ElementHasher for Rpo {
    type BaseField = BaseElement;

    fn hash_elements<E: FieldElement<BaseField = Self::BaseField>>(elements: &[E]) -> Self::Digest {
        // convert the elements into a list of base field elements
        let elements = E::as_base_elements(elements);

        // initialize state to all zeros, except for the first element of the capacity part, which
        // is set to 1 if the number of elements is not a multiple of RATE_WIDTH.
        let mut state = [BaseElement::ZERO; STATE_WIDTH];
        if elements.len() % RATE_WIDTH != 0 {
            state[CAPACITY_RANGE.start] = BaseElement::ONE;
        }

        // absorb elements into the state one by one until the rate portion of the state is filled
        // up; then apply the Rescue permutation and start absorbing again; repeat until all
        // elements have been absorbed
        let mut i = 0;
        for &element in elements.iter() {
            state[RATE_RANGE.start + i] = element;
            i += 1;
            if i % RATE_WIDTH == 0 {
                Self::apply_permutation(&mut state);
                i = 0;
            }
        }

        // if we absorbed some elements but didn't apply a permutation to them (would happen when
        // the number of elements is not a multiple of RATE_WIDTH), apply the RPO permutation after
        // padding by appending a 1 followed by as many 0 as necessary to make the input length a
        // multiple of the RATE_WIDTH.
        if i > 0 {
            state[RATE_RANGE.start + i] = BaseElement::ONE;
            i += 1;
            while i != RATE_WIDTH {
                state[RATE_RANGE.start + i] = BaseElement::ZERO;
                i += 1;
            }
            Self::apply_permutation(&mut state);
        }

        // return the first 4 elements of the state as hash result
        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
    }
}

// HASH FUNCTION IMPLEMENTATION
// ================================================================================================

impl Rpo {
    // RESCUE PERMUTATION
    // --------------------------------------------------------------------------------------------

    /// Applies RPO permutation to the provided state.
    pub fn apply_permutation(state: &mut [BaseElement; STATE_WIDTH]) {
        for i in 0..NUM_ROUNDS {
            Self::apply_round(state, i);
        }
    }

    /// RPO round function.
    #[inline(always)]
    pub fn apply_round(state: &mut [BaseElement; STATE_WIDTH], round: usize) {
        // apply first half of RPO round
        Self::apply_mds(state);
        Self::add_constants(state, &ARK1[round]);
        Self::apply_sbox(state);

        // apply second half of RPO round
        Self::apply_mds(state);
        Self::add_constants(state, &ARK2[round]);
        Self::apply_inv_sbox(state);
    }

    // HELPER FUNCTIONS
    // --------------------------------------------------------------------------------------------

    #[inline(always)]
    fn apply_mds(state: &mut [BaseElement; STATE_WIDTH]) {
        let mut result = [BaseElement::ZERO; STATE_WIDTH];

        // Using the linearity of the operations we can split the state into a low||high decomposition
        // and operate on each with no overflow and then combine/reduce the result to a field element.
        // The no overflow is guaranteed by the fact that the MDS matrix is a small powers of two in
        // frequency domain.
        let mut state_l = [0u64; STATE_WIDTH];
        let mut state_h = [0u64; STATE_WIDTH];

        for r in 0..STATE_WIDTH {
            let s = state[r].inner();
            state_h[r] = s >> 32;
            state_l[r] = (s as u32) as u64;
        }

        let state_h = mds_multiply_freq(state_h);
        let state_l = mds_multiply_freq(state_l);

        for r in 0..STATE_WIDTH {
            let s = state_l[r] as u128 + ((state_h[r] as u128) << 32);
            let s_hi = (s >> 64) as u64;
            let s_lo = s as u64;
            let z = (s_hi << 32) - s_hi;
            let (res, over) = s_lo.overflowing_add(z);

            result[r] =
                BaseElement::from_mont(res.wrapping_add(0u32.wrapping_sub(over as u32) as u64));
        }
        *state = result;
    }

    #[inline(always)]
    fn add_constants(state: &mut [BaseElement; STATE_WIDTH], ark: &[BaseElement; STATE_WIDTH]) {
        state.iter_mut().zip(ark).for_each(|(s, &k)| *s += k);
    }

    #[inline(always)]
    fn apply_sbox(state: &mut [BaseElement; STATE_WIDTH]) {
        state[0] = state[0].exp7();
        state[1] = state[1].exp7();
        state[2] = state[2].exp7();
        state[3] = state[3].exp7();
        state[4] = state[4].exp7();
        state[5] = state[5].exp7();
        state[6] = state[6].exp7();
        state[7] = state[7].exp7();
        state[8] = state[8].exp7();
        state[9] = state[9].exp7();
        state[10] = state[10].exp7();
        state[11] = state[11].exp7();
    }

    #[inline(always)]
    fn apply_inv_sbox(state: &mut [BaseElement; STATE_WIDTH]) {
        // compute base^10540996611094048183 using 72 multiplications per array element
        // 10540996611094048183 = b1001001001001001001001001001000110110110110110110110110110110111

        // compute base^10
        let mut t1 = *state;
        t1.iter_mut().for_each(|t| *t = t.square());

        // compute base^100
        let mut t2 = t1;
        t2.iter_mut().for_each(|t| *t = t.square());

        // compute base^100100
        let t3 = exp_acc::<BaseElement, STATE_WIDTH, 3>(t2, t2);

        // compute base^100100100100
        let t4 = exp_acc::<BaseElement, STATE_WIDTH, 6>(t3, t3);

        // compute base^100100100100100100100100
        let t5 = exp_acc::<BaseElement, STATE_WIDTH, 12>(t4, t4);

        // compute base^100100100100100100100100100100
        let t6 = exp_acc::<BaseElement, STATE_WIDTH, 6>(t5, t3);

        // compute base^1001001001001001001001001001000100100100100100100100100100100
        let t7 = exp_acc::<BaseElement, STATE_WIDTH, 31>(t6, t6);

        // compute base^1001001001001001001001001001000110110110110110110110110110110111
        for (i, s) in state.iter_mut().enumerate() {
            let a = (t7[i].square() * t6[i]).square().square();
            let b = t1[i] * t2[i] * *s;
            *s = a * b;
        }
    }
}

// MDS
// ================================================================================================
/// RPO MDS matrix
const MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = [
    [
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
    ],
    [
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
    ],
    [
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
    ],
    [
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
    ],
    [
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
    ],
    [
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
    ],
    [
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
    ],
    [
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
    ],
    [
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
    ],
    [
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
        BaseElement::new(8),
    ],
    [
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
        BaseElement::new(23),
    ],
    [
        BaseElement::new(23),
        BaseElement::new(8),
        BaseElement::new(26),
        BaseElement::new(13),
        BaseElement::new(10),
        BaseElement::new(9),
        BaseElement::new(7),
        BaseElement::new(6),
        BaseElement::new(22),
        BaseElement::new(21),
        BaseElement::new(8),
        BaseElement::new(7),
    ],
];

/// RPO Inverse MDS matrix
const INV_MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = [
    [
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
    ],
    [
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
    ],
    [
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
    ],
    [
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
    ],
    [
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
    ],
    [
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
    ],
    [
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
    ],
    [
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
    ],
    [
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
    ],
    [
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
    ],
    [
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
        BaseElement::new(13278298489594233127),
    ],
    [
        BaseElement::new(13278298489594233127),
        BaseElement::new(389999932707070822),
        BaseElement::new(9782021734907796003),
        BaseElement::new(4829905704463175582),
        BaseElement::new(7567822018949214430),
        BaseElement::new(14205019324568680367),
        BaseElement::new(15489674211196160593),
        BaseElement::new(17636013826542227504),
        BaseElement::new(16254215311946436093),
        BaseElement::new(3641486184877122796),
        BaseElement::new(11069068059762973582),
        BaseElement::new(14868391535953158196),
    ],
];

// ROUND CONSTANTS
// ================================================================================================

/// Rescue round constants;
/// computed as in section ?? from [specifications](https://github.org/aszepieniec/rpo/)
///
/// The constants are broken up into two arrays ARK1 and ARK2; ARK1 contains the constants for the
/// first half of RPO round, and ARK2 contains constants for the second half of RPO round.
const ARK1: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = [
    [
        BaseElement::new(5789762306288267392),
        BaseElement::new(6522564764413701783),
        BaseElement::new(17809893479458208203),
        BaseElement::new(107145243989736508),
        BaseElement::new(6388978042437517382),
        BaseElement::new(15844067734406016715),
        BaseElement::new(9975000513555218239),
        BaseElement::new(3344984123768313364),
        BaseElement::new(9959189626657347191),
        BaseElement::new(12960773468763563665),
        BaseElement::new(9602914297752488475),
        BaseElement::new(16657542370200465908),
    ],
    [
        BaseElement::new(12987190162843096997),
        BaseElement::new(653957632802705281),
        BaseElement::new(4441654670647621225),
        BaseElement::new(4038207883745915761),
        BaseElement::new(5613464648874830118),
        BaseElement::new(13222989726778338773),
        BaseElement::new(3037761201230264149),
        BaseElement::new(16683759727265180203),
        BaseElement::new(8337364536491240715),
        BaseElement::new(3227397518293416448),
        BaseElement::new(8110510111539674682),
        BaseElement::new(2872078294163232137),
    ],
    [
        BaseElement::new(18072785500942327487),
        BaseElement::new(6200974112677013481),
        BaseElement::new(17682092219085884187),
        BaseElement::new(10599526828986756440),
        BaseElement::new(975003873302957338),
        BaseElement::new(8264241093196931281),
        BaseElement::new(10065763900435475170),
        BaseElement::new(2181131744534710197),
        BaseElement::new(6317303992309418647),
        BaseElement::new(1401440938888741532),
        BaseElement::new(8884468225181997494),
        BaseElement::new(13066900325715521532),
    ],
    [
        BaseElement::new(5674685213610121970),
        BaseElement::new(5759084860419474071),
        BaseElement::new(13943282657648897737),
        BaseElement::new(1352748651966375394),
        BaseElement::new(17110913224029905221),
        BaseElement::new(1003883795902368422),
        BaseElement::new(4141870621881018291),
        BaseElement::new(8121410972417424656),
        BaseElement::new(14300518605864919529),
        BaseElement::new(13712227150607670181),
        BaseElement::new(17021852944633065291),
        BaseElement::new(6252096473787587650),
    ],
    [
        BaseElement::new(4887609836208846458),
        BaseElement::new(3027115137917284492),
        BaseElement::new(9595098600469470675),
        BaseElement::new(10528569829048484079),
        BaseElement::new(7864689113198939815),
        BaseElement::new(17533723827845969040),
        BaseElement::new(5781638039037710951),
        BaseElement::new(17024078752430719006),
        BaseElement::new(109659393484013511),
        BaseElement::new(7158933660534805869),
        BaseElement::new(2955076958026921730),
        BaseElement::new(7433723648458773977),
    ],
    [
        BaseElement::new(16308865189192447297),
        BaseElement::new(11977192855656444890),
        BaseElement::new(12532242556065780287),
        BaseElement::new(14594890931430968898),
        BaseElement::new(7291784239689209784),
        BaseElement::new(5514718540551361949),
        BaseElement::new(10025733853830934803),
        BaseElement::new(7293794580341021693),
        BaseElement::new(6728552937464861756),
        BaseElement::new(6332385040983343262),
        BaseElement::new(13277683694236792804),
        BaseElement::new(2600778905124452676),
    ],
    [
        BaseElement::new(7123075680859040534),
        BaseElement::new(1034205548717903090),
        BaseElement::new(7717824418247931797),
        BaseElement::new(3019070937878604058),
        BaseElement::new(11403792746066867460),
        BaseElement::new(10280580802233112374),
        BaseElement::new(337153209462421218),
        BaseElement::new(13333398568519923717),
        BaseElement::new(3596153696935337464),
        BaseElement::new(8104208463525993784),
        BaseElement::new(14345062289456085693),
        BaseElement::new(17036731477169661256),
    ],
];

const ARK2: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = [
    [
        BaseElement::new(6077062762357204287),
        BaseElement::new(15277620170502011191),
        BaseElement::new(5358738125714196705),
        BaseElement::new(14233283787297595718),
        BaseElement::new(13792579614346651365),
        BaseElement::new(11614812331536767105),
        BaseElement::new(14871063686742261166),
        BaseElement::new(10148237148793043499),
        BaseElement::new(4457428952329675767),
        BaseElement::new(15590786458219172475),
        BaseElement::new(10063319113072092615),
        BaseElement::new(14200078843431360086),
    ],
    [
        BaseElement::new(6202948458916099932),
        BaseElement::new(17690140365333231091),
        BaseElement::new(3595001575307484651),
        BaseElement::new(373995945117666487),
        BaseElement::new(1235734395091296013),
        BaseElement::new(14172757457833931602),
        BaseElement::new(707573103686350224),
        BaseElement::new(15453217512188187135),
        BaseElement::new(219777875004506018),
        BaseElement::new(17876696346199469008),
        BaseElement::new(17731621626449383378),
        BaseElement::new(2897136237748376248),
    ],
    [
        BaseElement::new(8023374565629191455),
        BaseElement::new(15013690343205953430),
        BaseElement::new(4485500052507912973),
        BaseElement::new(12489737547229155153),
        BaseElement::new(9500452585969030576),
        BaseElement::new(2054001340201038870),
        BaseElement::new(12420704059284934186),
        BaseElement::new(355990932618543755),
        BaseElement::new(9071225051243523860),
        BaseElement::new(12766199826003448536),
        BaseElement::new(9045979173463556963),
        BaseElement::new(12934431667190679898),
    ],
    [
        BaseElement::new(18389244934624494276),
        BaseElement::new(16731736864863925227),
        BaseElement::new(4440209734760478192),
        BaseElement::new(17208448209698888938),
        BaseElement::new(8739495587021565984),
        BaseElement::new(17000774922218161967),
        BaseElement::new(13533282547195532087),
        BaseElement::new(525402848358706231),
        BaseElement::new(16987541523062161972),
        BaseElement::new(5466806524462797102),
        BaseElement::new(14512769585918244983),
        BaseElement::new(10973956031244051118),
    ],
    [
        BaseElement::new(6982293561042362913),
        BaseElement::new(14065426295947720331),
        BaseElement::new(16451845770444974180),
        BaseElement::new(7139138592091306727),
        BaseElement::new(9012006439959783127),
        BaseElement::new(14619614108529063361),
        BaseElement::new(1394813199588124371),
        BaseElement::new(4635111139507788575),
        BaseElement::new(16217473952264203365),
        BaseElement::new(10782018226466330683),
        BaseElement::new(6844229992533662050),
        BaseElement::new(7446486531695178711),
    ],
    [
        BaseElement::new(3736792340494631448),
        BaseElement::new(577852220195055341),
        BaseElement::new(6689998335515779805),
        BaseElement::new(13886063479078013492),
        BaseElement::new(14358505101923202168),
        BaseElement::new(7744142531772274164),
        BaseElement::new(16135070735728404443),
        BaseElement::new(12290902521256031137),
        BaseElement::new(12059913662657709804),
        BaseElement::new(16456018495793751911),
        BaseElement::new(4571485474751953524),
        BaseElement::new(17200392109565783176),
    ],
    [
        BaseElement::new(17130398059294018733),
        BaseElement::new(519782857322261988),
        BaseElement::new(9625384390925085478),
        BaseElement::new(1664893052631119222),
        BaseElement::new(7629576092524553570),
        BaseElement::new(3485239601103661425),
        BaseElement::new(9755891797164033838),
        BaseElement::new(15218148195153269027),
        BaseElement::new(16460604813734957368),
        BaseElement::new(9643968136937729763),
        BaseElement::new(3611348709641382851),
        BaseElement::new(18256379591337759196),
    ],
];