#include "raven_spn.h"
#include "spn_regs.h"
#include "init.h"
#include "spn_checker_regs.h"

using spn_1 = spn_regs<0x90000000>;
using spn_2 = spn_regs<0xC0000000>;
using spn_checker = spn_checker_regs<0x10040000>;

void run_xspn1(int in_addr, int out_addr, int num_samples, int in_beats, int out_beats) {
    spn_1::mode_reg() = 0;
    spn_1::input_length_reg() = num_samples;      // each sample consists of 5 uint8 values
    spn_1::input_addr_reg() = in_addr;
    spn_1::output_addr_reg() = out_addr;
    spn_1::num_of_in_beats_reg() = in_beats;  // Number of AXI4 burst beats needed to load all input data
    spn_1::num_of_out_beats_reg() = out_beats;    // Number of AXI4 burst beats needed to store all result data
    spn_1::start_reg() = 1;
    printf("Starting first XSPN instance\n");
}

void run_xspn2(int in_addr, int out_addr, int num_samples, int in_beats, int out_beats) {
    spn_2::mode_reg() = 0;
    spn_2::input_length_reg() = num_samples;      // each sample consists of 5 uint8 values
    spn_2::input_addr_reg() = in_addr;
    spn_2::output_addr_reg() = out_addr;
    spn_2::num_of_in_beats_reg() = in_beats;  // Number of AXI4 burst beats needed to load all input data
    spn_2::num_of_out_beats_reg() = out_beats;    // Number of AXI4 burst beats needed to store all result data
    printf("Starting second XSPN instance\n");
    spn_2::start_reg() = 1;
}

static void spn1_interrupt_handler(){
	spn1_hw_interrupt = false;
    spn_1::interrupt_reg() = 1;
}

static void spn2_interrupt_handler(){
	spn2_hw_interrupt = false;
    spn_2::interrupt_reg() = 1;
}
/*! \brief main function
 *
 */
int main() {


    platform_init();
    configure_irq(2, spn1_interrupt_handler);
    configure_irq(22, spn2_interrupt_handler);


    uint32_t xspn_count = spn_checker::xspn_count_reg();
    uint32_t batch_size = spn_checker::batch_size_reg();
    uint32_t iterations = spn_checker::num_iterations_reg();

    printf("XSPN COUNT: %d\n", xspn_count);
    if (xspn_count < 1 || xspn_count > 2) {
        printf("ERROR: invalid XSPN COUNT");
        return 1;
    }
    printf("BATCH SIZE: %d\n", batch_size);
    printf("ITERATIONS: %d\n", iterations);



    int in_addr   = 0x30000000; // place input samples in the SPI memory
    int out_addr1 = 0x3C000000;
    int out_addr2 = 0x3E000000;

    spn_1::mode_reg() = 1;
    spn_1::start_reg() = 1;
    wait_for_spn1_interrupt();
    uint32_t readout = spn_1::readout_reg();
    printf("READOUT first HW instance:0x%x\n", readout);

    if (xspn_count == 2) {
        spn_2::mode_reg() = 1;
        spn_2::start_reg() = 1;
        wait_for_spn2_interrupt();
        uint32_t readout2 = spn_2::readout_reg();
        printf("READOUT second HW instance:0x%x\n", readout2);
    }

    uint32_t axi_bytes = readout;
    axi_bytes = axi_bytes & 0xff;
    axi_bytes = 1 << axi_bytes;

    printf("AXI Bytes: %d\n", axi_bytes);

    uint32_t sample_bytes = readout;
    sample_bytes = sample_bytes >> 16;
    sample_bytes = sample_bytes / 8;
    printf("Sample Bytes: %d\n", sample_bytes);

    uint32_t result_bytes = 8;
    printf("Result Bytes: %d\n", result_bytes);

    uint32_t in_bytes = batch_size * sample_bytes;
    uint32_t out_bytes = batch_size * result_bytes;

    uint32_t total_in = in_bytes * iterations;

    if (total_in > (out_addr1 - in_addr)) {
        printf("ERROR: input data requires %d bytes, only %d bytes available\n", total_in, out_addr1 - in_addr);
        return 1;
    }
    if (out_bytes > (out_addr2 - out_addr1)) {
        printf("ERROR: output data requires %d bytes, only %d bytes available\n", out_bytes, out_addr2 - out_addr1);
        return 1;
    }

    
    uint32_t in_beats = in_bytes / axi_bytes;
    if (in_beats * axi_bytes < batch_size * sample_bytes) in_beats++;
    uint32_t out_beats = out_bytes / axi_bytes;
    if (out_beats * axi_bytes < batch_size * result_bytes) out_beats++;

    uint32_t current_in_addr = in_addr;
    
    // inject SPN input data
    spn_checker::input_addr_reg() = current_in_addr;
    spn_checker::num_input_samples_reg() = sample_bytes * batch_size * iterations;
    spn_checker::start_data_trans_reg() = 1;
    spn_checker::output_addr_reg() = out_addr1;
    if (xspn_count == 2) {
        spn_checker::output_addr2_reg() = out_addr2;
    }
    for (int k = 0; k < iterations*batch_size; k+=batch_size) {
        run_xspn1(current_in_addr, out_addr1, batch_size, in_beats, out_beats);
        if (xspn_count == 2) {
            run_xspn2(current_in_addr, out_addr2, batch_size, in_beats, out_beats);
            wait_for_spn_interrupts();
        } else {
            wait_for_spn1_interrupt();
        }
        printf("XSPN finished\n");
        spn_checker::offset_reg() = k;
        spn_checker::length_reg() = batch_size;
        spn_checker::start_result_check_reg() = 1;

        current_in_addr += batch_size * sample_bytes; // 5 bytes in each sample (NIPS5)
    }

    return 0;
}