#include <stdlib.h>
#include <stdio.h>
#define log(...) printf(__VA_ARGS__)


// Branchy ASM (0)
//int find_branchy(int array[], int high, int element) {
//    int size = (high + 1) >> 1;
//    int current = size - 1;
//    log("Searching for %d in array of size %d\n", element, high + 1);
//    while (1) {
//        log("Checking index %d, size %d\n", current, size);
//        if (size == 0) {
//            return -1; // Element not found
//        }
//        if (array[current] == element) {
//            return current;
//        }
//        size = (size + 1) >> 1;
//        if (array[current] < element) {
//            current += size;
//        } else {
//            current -= size;
//        }
//    }
//}

// Branchy ASM (0)
int find_branchy(int array[], int high, int element) {
    int result = -1;
    asm volatile(
        "mv a2, %[high]\n\t"
        "srli a2, a2, 1\n\t"
        "mv a1, a2\n"
        // Loop start
        "loop%=:\n\t"
        "beq a2, zero, found%= \n\t"
        //Load the element at index a1
        "slli a3, a1, 2\n\t"
        "add a3, a3, %[array]\n\t"
        "lw a3, 0(a3)\n\t"
        // Compare the element with the target
        "beq a3, %[element], found%= \n\t"
        "addi a2, a2, 1\n\t"
        "srli a2, a2, 1\n\t"
        "bge a3, %[element], greater%= \n\t"
        // Less than
        "add a1, a1, a2\n\t"
        "j loop%= \n\t"
        // Greater than
        "greater%=:\n\t"
        "sub a1, a1, a2\n\t"
        "j loop%= \n\t"
        "found%=:\n\t"
        "mv %[result], a1\n\t"
        : [result] "=r" (result)
        : [array] "r" (array), [high] "r" (high), [element] "r" (element)
        : "a1", "a2", "a3", "cc", "memory"
    );
    return (result);
}

// Branch free (1)
int find_branchfree(int array[], int high, int element) {
    int result = -1;
    asm volatile(
        "mv a2, %[high]\n\t"
        "srli a2, a2, 1\n\t"
        "mv a1, a2\n"
        // Loop start
        "loop%=:\n\t"
        "beq a2, zero, found%= \n\t"
        //Load the element at index a1
        "slli a3, a1, 2\n\t"
        "add a3, a3, %[array]\n\t"
        "lw a3, 0(a3)\n\t"
        // Compare the element with the target
        "beq a3, %[element], found%= \n\t"
        "addi a2, a2, 1\n\t"
        "srli a2, a2, 1\n\t"
        "slt a4, %[element], a3\n\t"
        "czero.nez a5, a2, a4\n\t"
        "czero.eqz a6, a2, a4\n\t"
        "add a1, a1, a5\n\t"
        "sub a1, a1, a6\n\t"
        "j loop%= \n\t"
        "found%=:\n\t"
        "mv %[result], a1\n\t"
        : [result] "=r" (result)
        : [array] "r" (array), [high] "r" (high), [element] "r" (element)
        : "a1", "a2", "a3", "a4", "a5", "a6", "cc", "memory"
    );
    return (result);
}

#define STRINGIFY(x) #x
#define ASM_VALUE(x) STRINGIFY(x)

// BMS (2)
int find_bms(int array[], int high, int element) {
    int result = -1;
    asm volatile(
        "mv a2, %[high]\n\t"
        "srli a2, a2, 1\n\t"
        "mv a1, a2\n"
        // Loop start
        "loop%=:\n\t"
        "beq a2, zero, found%= \n\t"
        //Load the element at index a1
        "slli a3, a1, 2\n\t"
        "add a3, a3, %[array]\n\t"
        "lw a3, 0(a3)\n\t"
        // Compare the element with the target
        "beq a3, %[element], found%= \n\t"
        "addi a2, a2, 1\n\t"
        "srli a2, a2, 1\n\t"
        "bms problematicbranch%=,branchfree%=," ASM_VALUE(BMS_MODIFIER) "\n\t"
        // Branchy variant
        "problematicbranch%=:\n\t"
        "bge a3, %[element], greater%= \n\t"
        // Less than
        "add a1, a1, a2\n\t"
        "j loop%= \n\t"
        // Greater than
        "greater%=:\n\t"
        "sub a1, a1, a2\n\t"
        "j loop%= \n\t"
        // Branch free variant
        "branchfree%=:\n\t"
        "slt a4, a3, %[element]\n\t"
        "bjr a4, problematicbranch%=, greater%=\n\t"
        "czero.eqz a5, a2, a4\n\t"
        "czero.nez a6, a2, a4\n\t"
        "add a1, a1, a5\n\t"
        "sub a1, a1, a6\n\t"
        "j loop%= \n\t"
        "found%=:\n\t"
        "mv %[result], a1\n\t"
        : [result] "=r" (result)
        : [array] "r" (array), [high] "r" (high), [element] "r" (element)
        : "a1", "a2", "a3", "a4", "a5", "a6", "cc", "memory"
    );
    return (result);
}

int main(void) {
    int array[ELEMENTS];

    printf("Elements,Elements to find,Random Seed,Random placement,Bms Modifier,Variant,Percentage,Sum,Cycles\n");

    for (int i = 0; i < ELEMENTS; i++) {
        array[i] = i;
    }

    int elements_to_find[ELEMENTS_TO_FIND];

    for (int random_seed = 0; random_seed < 6; random_seed++) {
        for (int VARIANT = 0; VARIANT < 3; VARIANT++) {
            for (int RANDOM_PERCENTAGE = 0; RANDOM_PERCENTAGE <= 100; RANDOM_PERCENTAGE += 5) {
                srand(random_seed);

                // Fill elements_to_find with predictable values
                int random = rand() % ELEMENTS;
                for (int i = 0; i < ELEMENTS_TO_FIND; i++) {
                    elements_to_find[i] = random;
                }

                if (RANDOM_PLACEMENT == 0) {
                    // Replace random percentage values with unpredictable values
                    for (int i = 0; i < ELEMENTS_TO_FIND; i++) {
                        if ((rand() % ELEMENTS_TO_FIND) * 100 / ELEMENTS_TO_FIND < RANDOM_PERCENTAGE) {
                            elements_to_find[i] = rand() % ELEMENTS;
                        }
                    }
                } else if (RANDOM_PLACEMENT == 1) {
                    // Replace continous block at start of random percentage length with unpredictable values
                    int length = (ELEMENTS_TO_FIND * RANDOM_PERCENTAGE) / 100;
                    for (int i = 0; i < length; i++) {
                        elements_to_find[i] = rand() % ELEMENTS;
                    }
                }


                int (*find)(int array[], int high, int element) = NULL;
                switch (VARIANT) {
                    case 0:
                        find = find_branchy;
                        break;
                    case 1:
                        find = find_branchfree;
                        break;
                    case 2:
                        find = find_bms;
                        break;
                }

                long long cyclesStart;
                asm volatile("csrr %0, cycle" : "=r" (cyclesStart));


                int sum = 0;
                for (int i = 0; i < ELEMENTS_TO_FIND; i++) {
                    int element_to_find = elements_to_find[i];
                    int found = find(array, ELEMENTS - 1, element_to_find);
                    sum += found;
                    if (found != element_to_find) {
                        log("Error: Element %d not found, expected %d\n", found, element_to_find);
                    }
                }

                long long cyclesEnd;
                asm volatile("csrr %0, cycle" : "=r" (cyclesEnd));

                log("RESULT %d,%d,%d,%d,%d,%d,%d%%,%d,%lld\n", ELEMENTS, ELEMENTS_TO_FIND, random_seed, RANDOM_PLACEMENT, BMS_MODIFIER, VARIANT, RANDOM_PERCENTAGE, sum, cyclesEnd - cyclesStart);
            }
        }
    }
    return 0;
}