use crate::types::Z;

use crate::Q;

use sha3::digest::XofReader;

/// Algorithm 7 `SampleNTT(B)` on page 23.

/// Takes a 32-byte seed and two indices as input and outputs a pseudorandom element of `𝑇_𝑞`.

/// This implementation takes the `XofReader` directly.

///

/// Input: byte stream `B ∈ B^{34}`     ▷ a 32-byte seed along with two indices <br>

/// Output: array `a_hat ∈ Z^{256}_q`    ▷ the coefficients of the NTT of a polynomial

pub(crate) fn sample_ntt(mut xof_reader: impl XofReader) -> [Z; 256] {

    //

    let mut array_a_hat = [Z::default(); 256];

    let mut c = [0u8; 3]; // Space for 3 random (byte) draws

    // Not needed as XofReader is passed into function.

    // 1: ctx ← XOF.Init()

    // 2: ctx ← XOF.Absorb(ctx, 𝐵)    ▷ input the given byte array into XOF

    // 3: j ← 0

    let mut j = 0usize;

    // This rejection sampling loop is solely dependent upon rho which crosses a trust boundary

    // in the clear. Thus, it does not need to be constant time.

    // 4: while j < 256 do

    while j < 256 {

        //

        // 5: (ctx, 𝐶) ← XOF.Squeeze(ctx, 3)    ▷ get a fresh 3-byte array 𝐶 from XOF

        xof_reader.read(&mut c); // Draw 3 bytes

        // 6: 𝑑1 ← 𝐶[0] + 256 ⋅ (𝐶[1] mod 16)    ▷ 0 ≤ 𝑑1 < 2^{12}

        let d1 = u16::from(c[0]) + 256 * (u16::from(c[1]) & 0x0F);

        // 7: 𝑑2 ← ⌊𝐶[1]/16⌋ + 16 ⋅ 𝐶[2]    ▷ 0 ≤ 𝑑2 < 2^{12}

        let d2 = (u16::from(c[1]) >> 4) + 16 * u16::from(c[2]);

        // 8: if d1 < q then

        if d1 < Q {

            //

            // 9: a_hat[j] ← d1    ▷ a_hat ∈ Z^{256}_q

            array_a_hat[j].set_u16(d1);

            // 10: j ← j + 1

            j += 1;

            // 11: end if

        }

        // 12: if d2 < q and j < 256 then

        // Bitwise AND; `&&` would do as well, as this loop need not be constant time (see above)

        #[allow(clippy::needless_bitwise_bool)]

        if (d2 < Q) & (j < 256) {

            //

            // 13: a_hat[j] ← d2

            array_a_hat[j].set_u16(d2);

            // 14: j ← j + 1

            j += 1;

            // 15: end if

        }

        // 16: end while

    }

    // 17: return a_hat

    array_a_hat

}

/// Algorithm 8 `SamplePolyCBD_η(B)` on page 23.

/// Takes a seed as input and outputs a pseudorandom sample from the distribution `D_𝜂(𝑅_𝑞)`. <br>

/// This function is an optimized version that avoids the `BytesToBits` function (algorithm 4).

///

/// Input: byte array `B ∈ B^{64·η}` <br>

/// Output: array `f ∈ Z^{256}_q`

///

/// The output is written into `array_f` rather than returned by value, so the secret sample leaves

/// no copy outside the caller's buffer, which the caller wipes (FIPS 203 §3.3).

pub(crate) fn sample_poly_cbd(byte_array_b: &[u8], array_f: &mut [Z; 256]) {

    let eta = u32::try_from(byte_array_b.len()).unwrap() >> 6;

    debug_assert_eq!(byte_array_b.len(), 64 * eta as usize, "Alg 8: byte array not 64 * eta");

    let mut temp = 0;

    let mut int_index = 0;

    let mut bit_index = 0;

    for byte in byte_array_b {

        temp |= u32::from(*byte) << bit_index;

        bit_index += 8;

        while bit_index >= 2 * (eta as usize) {

            let tmask_x = temp & ((1 << eta) - 1);

            let x = count_ones(tmask_x);

            let tmask_y = (temp >> eta) & ((1 << eta) - 1);

            let y = count_ones(tmask_y);

            let (mut xx, mut yy) = (Z::default(), Z::default());

            xx.set_u16(x);

            yy.set_u16(y);

            array_f[int_index] = xx.sub(yy);

            bit_index -= 2 * (eta as usize);

            temp >>= 2 * (eta as usize);

            int_index += 1;

        }

    }

}

// Count u8 ones in constant time (u32 helps perf)

#[allow(clippy::cast_possible_truncation)] // return x as u16

fn count_ones(x: u32) -> u16 {

    let x = (x & 0x5555_5555) + ((x >> 1) & 0x5555_5555);

    let x = (x & 0x3333_3333) + ((x >> 2) & 0x3333_3333);

    let x = (x & 0x0F0F_0F0F) + ((x >> 4) & 0x0F0F_0F0F);

    x as u16

}

// The original pseudocode for Algorithm 8 follows...

// Algorithm 8 `SamplePolyCBDη(B)` on page 23.

// Takes a seed as input and outputs a pseudorandom sample from the distribution `D_𝜂(𝑅_𝑞)`.

//

// Input: byte array B ∈ B^{64·η}

// Output: array f ∈ Z^{256}_q

// 1: b ← BytesToBits(B)

// 2: for (i ← 0; i < 256; i ++)

// 3:   x ← ∑_{j=0}^{η-1} b[2iη + j] //

// 4:   y ← ∑_{j=0}^{η-1} b[2iη + η + j]

// 5:   f [i] ← x − y mod q

// 6: end for

// 7: return f

// }