use std::machines::hash::poseidon_bb::PoseidonBB;
use std::machines::range::Byte2;
use std::machines::range::Bit12;
use std::machines::small_field::memory::Memory;
use std::machines::split::ByteCompare;
use std::machines::split::split_bb::SplitBB;

let main_degree: int = 2**10;
let memory_degree: int = 2**12;
let poseidon2_degree: int = 2**12;
let split_bb_degree: int = 2**12;

machine Main with degree: main_degree {
    reg pc[@pc];
    reg X1[<=];
    reg X2[<=];
    reg ADDR1_LOW[<=];
    reg ADDR1_HIGH[<=];
    reg ADDR2_LOW[<=];
    reg ADDR2_HIGH[<=];

    ByteCompare byte_compare;
    SplitBB split(byte_compare, split_bb_degree, split_bb_degree);

    // Increase the time step by 2 in each row, so that the poseidon machine
    // can read in the given time step and write in the next time step.
    col fixed STEP(i) { 2 * i };
    Byte2 byte2;
    Bit12 bit12;
    Memory memory(bit12, byte2, memory_degree, memory_degree);
    instr mstore_le ADDR1_HIGH, ADDR1_LOW, X1, X2 ->
        link ~> memory.mstore(ADDR1_HIGH, ADDR1_LOW, STEP, X1, X2);

    PoseidonBB poseidon(memory, split, poseidon2_degree, poseidon2_degree);
    instr poseidon ADDR1_HIGH, ADDR1_LOW, ADDR2_HIGH, ADDR2_LOW ->
        link ~> poseidon.poseidon_permutation(
            ADDR1_HIGH, ADDR1_LOW,
            ADDR2_HIGH, ADDR2_LOW,
            STEP
        );

    col witness val_low, val_high;
    instr assert_eq ADDR1_HIGH, ADDR1_LOW, X1 ->
        link ~> (val_high, val_low) = memory.mload(ADDR1_HIGH, ADDR1_LOW, STEP)
    {
        val_low + 2**16 * val_high = X1
    }

    function main {

        // Test vectors generated by gen_poseidon_bb_consts

        // All zeros:
        mstore_le 0, 0, 0, 0;
        mstore_le 0, 4, 0, 0;
        mstore_le 0, 8, 0, 0;
        mstore_le 0, 12, 0, 0;
        mstore_le 0, 16, 0, 0;
        mstore_le 0, 20, 0, 0;
        mstore_le 0, 24, 0, 0;
        mstore_le 0, 28, 0, 0;
        mstore_le 0, 32, 0, 0;
        mstore_le 0, 36, 0, 0;
        mstore_le 0, 40, 0, 0;
        mstore_le 0, 44, 0, 0;
        mstore_le 0, 48, 0, 0;
        mstore_le 0, 52, 0, 0;
        mstore_le 0, 56, 0, 0;
        mstore_le 0, 60, 0, 0;

        poseidon 0, 0, 0, 0;

        assert_eq 0, 0, 958909011;
        assert_eq 0, 4, 1204189760;
        assert_eq 0, 8, 1051708848;
        assert_eq 0, 12, 538282630;
        assert_eq 0, 16, 129608908;
        assert_eq 0, 20, 219931460;
        assert_eq 0, 24, 1794729453;
        assert_eq 0, 28, 1386481683;

        // All ones:
        mstore_le 0, 0, 0, 1;
        mstore_le 0, 4, 0, 1;
        mstore_le 0, 8, 0, 1;
        mstore_le 0, 12, 0, 1;
        mstore_le 0, 16, 0, 1;
        mstore_le 0, 20, 0, 1;
        mstore_le 0, 24, 0, 1;
        mstore_le 0, 28, 0, 1;
        mstore_le 0, 32, 0, 1;
        mstore_le 0, 36, 0, 1;
        mstore_le 0, 40, 0, 1;
        mstore_le 0, 44, 0, 1;
        mstore_le 0, 48, 0, 1;
        mstore_le 0, 52, 0, 1;
        mstore_le 0, 56, 0, 1;
        mstore_le 0, 60, 0, 1;

        poseidon 0, 0, 0, 0;

        assert_eq 0, 0, 605188316;
        assert_eq 0, 4, 179227070;
        assert_eq 0, 8, 987768411;
        assert_eq 0, 12, 743435232;
        assert_eq 0, 16, 1037675772;
        assert_eq 0, 20, 926267312;
        assert_eq 0, 24, 842954441;
        assert_eq 0, 28, 1311415732;

        // All elements are -1 (in BabyBear, 0x78000000)
        mstore_le 0, 0, 30720, 0;
        mstore_le 0, 4, 30720, 0;
        mstore_le 0, 8, 30720, 0;
        mstore_le 0, 12, 30720, 0;
        mstore_le 0, 16, 30720, 0;
        mstore_le 0, 20, 30720, 0;
        mstore_le 0, 24, 30720, 0;
        mstore_le 0, 28, 30720, 0;
        mstore_le 0, 32, 30720, 0;
        mstore_le 0, 36, 30720, 0;
        mstore_le 0, 40, 30720, 0;
        mstore_le 0, 44, 30720, 0;
        mstore_le 0, 48, 30720, 0;
        mstore_le 0, 52, 30720, 0;
        mstore_le 0, 56, 30720, 0;
        mstore_le 0, 60, 30720, 0;

        poseidon 0, 0, 0, 0;

        assert_eq 0, 0, 1100602725;
        assert_eq 0, 4, 287073314;
        assert_eq 0, 8, 1141334375;
        assert_eq 0, 12, 212177781;
        assert_eq 0, 16, 1688593373;
        assert_eq 0, 20, 1515905513;
        assert_eq 0, 24, 1635387847;
        assert_eq 0, 28, 693019811;

        // Some other values (ported from poseidon_gl test):
        mstore_le 0, 0, 14, 6474;
        mstore_le 0, 4, 3225, 31229;
        mstore_le 0, 8, 13499, 22633;
        mstore_le 0, 12, 1, 47334;
        mstore_le 0, 16, 26147, 51257;
        mstore_le 0, 20, 14081, 12102;
        mstore_le 0, 24, 25381, 5145;
        mstore_le 0, 28, 0, 2087;
        mstore_le 0, 32, 0, 0;
        mstore_le 0, 36, 0, 0;
        mstore_le 0, 40, 0, 0;
        mstore_le 0, 44, 0, 0;
        mstore_le 0, 48, 0, 0;
        mstore_le 0, 52, 0, 0;
        mstore_le 0, 56, 0, 0;
        mstore_le 0, 60, 0, 0;

        poseidon 0, 0, 0, 0;

        assert_eq 0, 0, 731906575;
        assert_eq 0, 4, 1685915147;
        assert_eq 0, 8, 1200872284;
        assert_eq 0, 12, 1493840138;
        assert_eq 0, 16, 404883058;
        assert_eq 0, 20, 1229250173;
        assert_eq 0, 24, 1012016786;
        assert_eq 0, 28, 481888550;

        // Repeat the second test, but be fancy with the memory pointers being passed:
        mstore_le 42, 65520, 0, 1;
        mstore_le 42, 65524, 0, 1;
        mstore_le 42, 65528, 0, 1;
        mstore_le 42, 65532, 0, 1;
        mstore_le 43, 0, 0, 1;
        mstore_le 43, 4, 0, 1;
        mstore_le 43, 8, 0, 1;
        mstore_le 43, 12, 0, 1;
        mstore_le 43, 16, 0, 1;
        mstore_le 43, 20, 0, 1;
        mstore_le 43, 24, 0, 1;
        mstore_le 43, 28, 0, 1;
        mstore_le 43, 32, 0, 1;
        mstore_le 43, 36, 0, 1;
        mstore_le 43, 40, 0, 1;
        mstore_le 43, 44, 0, 1;

        // This will read 64 bytes starting at address 0x002afff0 and
        // write the result to bytes starting at address 0x002afff4.
        // Both operations should overflow the lower 16 bits of the address.
        poseidon 42, 65520, 42, 65524;

        assert_eq 42, 65524, 605188316;
        assert_eq 42, 65528, 179227070;
        assert_eq 42, 65532, 987768411;
        assert_eq 43, 0, 743435232;
        assert_eq 43, 4, 1037675772;
        assert_eq 43, 8, 926267312;
        assert_eq 43, 12, 842954441;
        assert_eq 43, 16, 1311415732;

        return;
    }
}
