// This code is to be copy pasted in you vivtis application component
// alongside data generated by the python generator
// PLEASE INCREASE YOUR STACK AND HEAP SIZE ! to avoid program stall

#include "xparameters.h"
#include "xaxidma.h"
#include <stdio.h>
#include <xstatus.h>
#include "mnist_samples.h"

XAxiDma AxiDma;

int init_dma(XAxiDma *AxiDma) {
    XAxiDma_Config* CfgPtr;
    int status;

    CfgPtr = XAxiDma_LookupConfig(XPAR_AXI_DMA_0_BASEADDR);
    if (!CfgPtr) {
        xil_printf("No configuration found for %d\n", XPAR_AXI_DMA_0_BASEADDR);
        return XST_FAILURE;
    }

    status = XAxiDma_CfgInitialize(AxiDma, CfgPtr);
    if (status != XST_SUCCESS) {
        xil_printf("Initialization failed\n");
        return XST_FAILURE;
    }

    if (XAxiDma_HasSg(AxiDma)) {
        xil_printf("Device configured as SG mode\n");
        return XST_FAILURE;
    }

    return XST_SUCCESS;
}

static inline void enable_pmu_cycle_counter(void) {
    asm volatile("mcr p15, 0, %0, c9, c12, 1" :: "r"(1 << 31));  // Enable cycle counter
    asm volatile("mcr p15, 0, %0, c9, c12, 0" :: "r"(1));        // Enable all counters
}

static inline uint32_t read_pmu_cycle_counter(void) {
    uint32_t value;
    asm volatile("mrc p15, 0, %0, c9, c13, 0" : "=r"(value));
    return value;
}

int main(void) {
    enable_pmu_cycle_counter();
    uint32_t start, end;

    int status = init_dma(&AxiDma);
    if(status != XST_SUCCESS) {
        xil_printf("Error while initializing the DMA\n");
        return 1;
    }

    xil_printf("DMA initialized successfully\n");

    volatile char TxBuffer[IMAGE_SIZE*NUM_SAMPLES] __attribute__ ((aligned (32)));
    volatile int RxBuffer[NUM_SAMPLES] __attribute__ ((aligned (32)));

    xil_printf("Memory init OKAY\n");

    for(int j = 0; j < NUM_SAMPLES; j++) {
        for(int i = 0; i < IMAGE_SIZE; i++) {
            // xil_printf("I : %d   ///   J : %d\n", i, j); // debug purpose
            TxBuffer[j * IMAGE_SIZE + i] = (char)mnist_samples[j][i]; // fill with variable placeholder data
        }
    }

    xil_printf("Memory allocation OKAY\n");    

    Xil_DCacheFlushRange((UINTPTR)TxBuffer, NUM_SAMPLES * IMAGE_SIZE * sizeof(char));
    Xil_DCacheFlushRange((UINTPTR)RxBuffer, NUM_SAMPLES * sizeof(char));

    xil_printf("Cach flush OKAY, Strating transfers...\n");    

    start = read_pmu_cycle_counter();
    for(int k = 0; k < NUM_SAMPLES; k++) {

        status = XAxiDma_SimpleTransfer(&AxiDma, (UINTPTR)&TxBuffer[k*IMAGE_SIZE], IMAGE_SIZE * sizeof(char), XAXIDMA_DMA_TO_DEVICE);
        //printf("%i TO_DEVICE status code\n", status);
        if (status != XST_SUCCESS) {
            xil_printf("Error: DMA transfer to device failed\n");
            return XST_FAILURE;
        }

        status = XAxiDma_SimpleTransfer(&AxiDma, (UINTPTR)&RxBuffer[k], sizeof(int), XAXIDMA_DEVICE_TO_DMA);
        //printf("%i FROM_DEVICE status code\n", status);
        if (status != XST_SUCCESS) {
            xil_printf("Error: DMA transfer from device failed\n");
            return XST_FAILURE;
        }

        while (XAxiDma_Busy(&AxiDma, XAXIDMA_DMA_TO_DEVICE) || 
               XAxiDma_Busy(&AxiDma, XAXIDMA_DEVICE_TO_DMA)) {
            ;
        }
        xil_printf("#%i iteration done\n", k);
    }
    end = read_pmu_cycle_counter();

    // Output classifier's results & compute the accuracy

    int valid = 0;
    int accuracy_percentage;

    for(int i = 0; i < NUM_SAMPLES; i++) {
        xil_printf("FPGA value RxBuffer[%d] = %d\n", i, RxBuffer[i]);
        if(RxBuffer[i] == mnist_labels[i]){
            valid++;
        }
    }
    // Calculate accuracy as a percentage, multiplied by 100 to preserve precision
    accuracy_percentage = (valid * 100) / NUM_SAMPLES;
    xil_printf("\n\nMODEL ACCURACY = %d%%\n", accuracy_percentage);

    uint32_t cycles = end - start;
    double time_ms = (double)cycles / 667000.0;
    printf("Execution time: %f milliseconds\n", time_ms);

    return 0;
}