use.std::crypto::fri::ext2fri
use.std::crypto::stark::random_coin
use.std::crypto::stark::constants

#! Compute the number of FRI layers given log2 of the size of LDE domain. It also computes the
#! LDE domain generator and, from it, the trace generator and store these for later use.
#!
#! Input: [...]
#! Output: [num_fri_layers, ...]
#! Cycles: 52
export.generate_fri_parameters
    # Load FRI verifier data
    padw exec.constants::lde_size_ptr mem_loadw
    #=> [lde_size, log(lde_size), lde_g, 0, ...] (6 cycles)

    # Store in `TMP5` in order to use it for fri layer loading
    exec.constants::tmp5 mem_storew

    # Compute [gz1, gz0, z1, z0] using domain generator
    # TODO: move to somewhere else
    # ---------------------------------------------------------------------------------------------

    # load z from memory
    padw
    exec.constants::z_ptr mem_loadw
    #=> [(z1, z0)^n, z1, z0, lde_size, log2(lde_size), lde_g, 0, ...] (6 cycles)

    # prepare stack
    drop
    drop
    dup.1
    dup.1
    dup.6
    drop
    #=> [z1, z0, z1, z0, lde_size, log2(lde_size), lde_g, 0, ...] (6 cycles)

    # Load `trace_g` from memory
    exec.constants::trace_domain_generator_ptr mem_load
    #=> [trace_g, z1, z0, z1, z0, lde_size, log2(lde_size), lde_g, 0, ...] (2 cycles)

    # Compute `gz0` = `trace_g * z_0`
    dup
    movup.3
    mul
    #=> [gz0, trace_g, z1, z0, z1, z0, lde_size, log2(lde_size), lde_g, 0, ...] (3 cycles)

    # Compute `gz1` = `trace_g * z_1`
    swap.2
    mul
    #=> [gz1, gz0, z1, z0, lde_size, log2(lde_size), lde_g, 0, ...] (2 cycles)

    # Save `[gz1, gz0, z1, z0]` and clean the stack
    exec.constants::tmp1 mem_storew
    dropw
    #=> [lde_size, log2(lde_size), lde_g, 0, ...] (6 cycles)

    # Compute the number of FRI layers
    dup
    dup.2
    dup
    is_odd
    if.true
        push.32
        swap
        sub.5
        div.2
    else
        push.64
        swap
        sub.6
        div.2
    end
    # => [num_fri_layers, remainder_size, lde_size, lde_size, log2(lde_size), domain_gen, 0, ...] (12 cycles)

    # Save `[num_fri_layers, remainder_size, lde_size, lde_size]` in memory
    exec.constants::tmp6 mem_storew
    movdn.6
    dropw
    drop
    drop
    # => [num_fri_layers, ...] (10 cycles)
end

#! Get FRI layer commitments and reseed with them in order to draw folding challenges i.e. alphas.
#!
#! Input: [ptr_layer, num_layers, ...]
#! Output: [...]
#! Cycles: 21 + 83 * num_fri_layers
export.load_fri_layer_commitments
    # Address containing the first layer commitment
    push.0.0
    movup.3
    movup.3
    swap
    push.0.0.0.0
    swapw
    # => [num_layers, ptr_layer, y, y, Y, ...]

    dup
    push.0
    neq
    while.true
        swapw               # [Y, num_layers, ptr_layer, y, y, ...]
        adv_loadw           # [Com, num_layers, ptr_layer, y, y, ...]

        # Save FRI layer commitment
        dup.5
        add.1
        swap.6
        mem_storew
        #=> [Com, num_layers, ptr_layer + 1, y, y, ...]

        # Reseed
        exec.random_coin::reseed
        # => [num_layers, ptr_layer + 1, y, y, ...]

        push.0.0.0.0
        exec.random_coin::get_rate_1
        push.0.0
        exec.constants::tmp5 mem_loadw
        # => [lde_size, log2(lde_size), lde_generator, 0, a1, a0, Y, num_layers, ptr_layer + 1, y, y, ...]

        # Compute and save to memory new lde_size and its new logarithm
        div.4
        swap
        sub.2
        swap
        exec.constants::tmp5 mem_storew

        # Move the pointer higher up the stack
        movup.2 drop
        movup.2 drop
        swapw
        dropw
        # => [lde_size, log2(lde_size), a1, a0, num_layers, ptr_layer + 1, y, y, Y, ...]

        # Save [lde_size, log2(lde_size), a1, a0] in memory next to the layer commitment
        dup.5
        add.1
        swap.6
        mem_storew
        swapw
        # => [num_layers, ptr_layer + 2, y, y, lde_size, log2(lde_size), a1, a0, Y]

        # Decrement the FRI layer counter
        sub.1
        dup
        push.0
        neq
    end
    # => [Y, Y]
    dropw
    dropw
    #=> [...]
end

#! Load the remainder polynomial from the advice provider and check that its hash corresponds
#! to its commitment and reseed with the latter.
#! Load the remainder code word, i.e. the NTT of the remainder polynomial, and use its hash, together,
#! with the hash of the remainder polynomial in order to generate the Fiat-Shamir challenge `tau` for
#! the `verify_remainder_xx` procedure.
#!
#! Input: [...]
#! Output: [...]
#! Cycles:
#!  1- Remainder of size 32: 1633
#!  2- Remainder of size 64: 3109
export.load_and_verify_remainder
    # Load remainder commitment and save it at `TMP7`
    push.0.0.0.0
    adv_loadw
    exec.constants::tmp7 mem_storew

    # Reseed with remainder commitment
    exec.random_coin::reseed

    # adv_pipe the remainder codeword
    ## Get the length of remainder
    exec.constants::tmp6 mem_loadw
    ## Compute the correct remainder pointer using length of remainder
    exec.constants::fri_com_ptr
    swap
    mul.2
    add
    ## Store for later use
    exec.constants::tmp8 mem_storew
    #=> [ptr_remainder, remainder_size, y, y]

    dup.1
    push.32
    eq
    if.true
        # Remainder length equal to 32
        push.0.0.0.0
        push.0.0.0.0
        push.0.0.0.0
        # => [Y, Y, 0, 0, 0, 0 ptr_remainder, remainder_size, y, y]

        # adv_load remainder polynomial
        # TODO: This is a workaround since the FRI verifier expects the memory layout to be
        # [query_ptr ... layer_ptr ... rem_ptr ...] which leaves only one option for laying out
        # the polynomial coefficients i.e. starting at remainder_ptr + remainder_codeword_length/2.
        # On the other hand, we need to check that the hash of the polynomial coefficients agrees with
        # the commitment already received by the prover. Thus the need for hashing the polynomial
        # coefficients first.
        adv_loadw
        dup.12
        add.16
        mem_storew
        swapw
        adv_loadw
        dup.12
        add.17
        mem_storew
        hperm
        # => [Y, Remainder_poly_com, Y, ptr_remainder, remainder_size, y, y]

        # Compare Remainder_poly_com with the read commitment
        exec.constants::tmp7 mem_loadw
        movup.4
        assert_eq
        movup.3
        assert_eq
        movup.2
        assert_eq
        assert_eq
        # => [Y, ptr_remainder, remainder_size, y, y]
        push.0.0.0.0
        push.0.0.0.0
        repeat.8
            adv_pipe hperm
        end
    else
        # Remainder length equal to 64
        push.0.0.0.0
        push.0.0.0.0
        push.0.0.0.0
        # => [Y, Y, 0, 0, 0, 0 ptr_remainder, remainder_size, y, y]

        # adv_load remainder polynomial
        # TODO: This is a workaround since the FRI verifier expects the memory layout to be
        # [query_ptr ... layer_ptr ... rem_ptr ...] which leaves only one option for laying out
        # the polynomial coefficients i.e. starting at remainder_ptr + remainder_codeword_length/2.
        # On the other hand, we need to check that the hash of the polynomial coefficients agrees with
        # the commitment already received by the prover. Thus the need for hashing the polynomial
        # coefficients first.
        adv_loadw
        dup.12
        add.32
        mem_storew
        swapw
        adv_loadw
        dup.12
        add.33
        mem_storew
        hperm

        adv_loadw
        dup.12
        add.34
        mem_storew
        swapw
        adv_loadw
        dup.12
        add.35
        mem_storew
        hperm
        # => [Y, Remainder_poly_com, Y, ptr_remainder, remainder_size, y, y]

        # Compare Remainder_poly_com with the read commitment
        exec.constants::tmp7 mem_loadw
        movup.4
        assert_eq
        movup.3
        assert_eq
        movup.2
        assert_eq
        assert_eq
        # => [Y, ptr_remainder, remainder_size, y, y]
        push.0.0.0.0
        push.0.0.0.0
        repeat.16
            adv_pipe hperm
        end
    end
    # => [Y, R, Y, Y] where R = [y, y, tau1, tau0]

    dropw
    swapw.2
    dropw
    drop
    drop
    #=> [Y, tau1, tau0] where tau is the challenge of ext2fri::verify_remainder_xx

    # Prepare for remainder verification procedure
    exec.constants::tmp8 mem_loadw
    movup.2 drop
    movup.2 drop
    # => [ptr_remainder, remainder_size, tau1, tau0]

    # Call the correct remainder verification procedure
    movdn.3
    push.32
    eq
    if.true
        exec.ext2fri::verify_remainder_32
    else
        exec.ext2fri::verify_remainder_64
    end
    #=> [...]
end