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

void printArray(int array[], int size) {
    for (int i = 0; i < size; i++) {
        log("%d ", array[i]);
    }
    log("\n");
}

void swap(int *a, int *b) {
    int t = *a;
    *a = *b;
    *b = t;
}

//void sort_indexes(int array[], int *i1, int *i2) {
//    int x = *i1;
//    int y = *i2;
//    int c = array[*i1] <= array[*i2];
//    *i1 = c ? x : y;
//    *i2 = c ? y : x;
//}

void sort_indexes(int array[], int *i1, int *i2) {
    if (array[*i1] > array[*i2]) {
        swap(&array[*i1], &array[*i2]);
    }
}

void insertionSort(int array[], int n) {
    int i, element, j;
    for (i = 1; i < n; i++) {
        element = array[i];
        j = i - 1;

        while (j >= 0 && array[j] > element) {
            array[j + 1] = array[j];
            j--;
        }
        array[j + 1] = element;
    }
}

int median_of_three(int array[], int low, int high) {
  if (high - low < 2) {
    return array[high];
  }

  int l = low;
  int m = (low + high) / 2;
  int h = high;

  sort_indexes(array, &l, &m);
  sort_indexes(array, &m, &h);
  sort_indexes(array, &l, &m);

  swap(&array[m], &array[high]);

  return array[high];
}

int median_of_three_of_median_of_three(int array[], int low, int high) {
  if (high - low < 100) {
    return median_of_three(array, low, high);
  } else {
    int offset = (high - low) / 8;

    int a = low;
    int b = low + offset;
    int c = low + offset * 2;
    sort_indexes(array, &a, &b);
    sort_indexes(array, &b, &c);
    sort_indexes(array, &a, &b);

    int d = low + offset * 3;
    int e = low + offset * 4;
    int f = low + offset * 5;
    sort_indexes(array, &d, &e);
    sort_indexes(array, &e, &f);
    sort_indexes(array, &d, &e);
    
    int g = low + offset * 6;
    int h = low + offset * 7;
    int i = high;
    sort_indexes(array, &g, &h);
    sort_indexes(array, &h, &i);
    sort_indexes(array, &g, &h);
    
    sort_indexes(array, &b, &e);
    sort_indexes(array, &e, &h);
    sort_indexes(array, &b, &e);

    swap(&array[e], &array[high]);

    return array[high];
  }
}

#if defined(VARIANT) && VARIANT == 0
// Branchy
int partition(int array[], int low, int high) {
    int pivot = median_of_three_of_median_of_three(array, low, high);

    int i = (low - 1);
    for (int j = low; j < high; j++) {
        if (array[j] <= pivot) {
            i++;
            swap(&array[i], &array[j]);
        }
    }
    swap(&array[i + 1], &array[high]);
    return (i + 1);
}
#elif defined(VARIANT) && VARIANT == 1
// Branch free
int partition(int array[], int low, int high) {
    int pivot = median_of_three_of_median_of_three(array, low, high);

    int i = low;
    for (int j = low; j < high; j++) {
        int c = pivot > array[j];
        int y = array[i];
        int z = array[j];
        array[i] = c ? z : y;
        array[j] = c ? y : z;
        i += c;
    }
    swap(&array[i], &array[high]);
    return (i);
}
#elif defined(VARIANT) && VARIANT == 2
// BMS
int partition(int array[], int low, int high) {
    int pivot = median_of_three_of_median_of_three(array, low, high);

    int i = low;
    asm volatile(
        // a1 == j == low
        "mv a1, %[low]\n\t"
        // Loop start
        "bge a1, %[high], loop_end%=\n"
        "loop%=:\n\t"
        // Get the array element at index a1
        "slli a2, a1, 2\n\t"
        "add a2, a2, %[array]\n\t"
        "lw a3, 0(a2)\n"
        "bms problematicbranch%=,branchfree%=,19\n\t"
        // Branchy variant
        "problematicbranch%=:\n\t"
        // Compare the array element with the pivot
        "bgt a3, %[pivot], branchy_loop_increment%=\n\t"
        // Swap the array element j with the element i
        "slli a6, %[i], 2\n\t"
        "add a6, a6, %[array]\n\t"
        "lw a7, 0(a6)\n\t"
        "sw a7, 0(a2)\n\t"
        "sw a3, 0(a6)\n"
        // Increase i
        "addiw %[i], %[i], 1\n\t"
        // Increase a1 for the next iteration
        "branchy_loop_increment%=:\n\t"
        "addi a1, a1, 1\n\t"
        "blt a1, %[high], loop%=\n\t"
        "jal loop_end%=\n"
        // Branch free variant
        "branchfree%=:\n\t"
        // Load the array element at index i
        "slli a6, %[i], 2\n\t"
        "add a6, a6, %[array]\n\t"
        "lw a7, 0(a6)\n\t"
        "sltu a5, %[pivot], a3\n\t"
        "bjr a5, problematicbranch%=, branchy_loop_increment%=\n\t"
        // Write from index a1 to index i when a3 < pivot
        "czero.eqz a4, a3, a5\n\t"
        "czero.nez a0, a7, a5\n\t"
        "or a4, a4, a0\n\t"
        "sw a4, 0(a2)\n\t"
        // Write from index i to index a1 when a3 < pivot
        "czero.eqz a4, a7, a5\n\t"
        "czero.nez a0, a3, a5\n\t"
        "or a4, a4, a0\n\t"
        "sw a4, 0(a6)\n\t"
        // Increase i
        "addi %[i], %[i], 1\n\t"
        "sub %[i], %[i], a5\n\t"
        // Increase a1 for the next iteration
        "loop_increment%=:\n\t"
        "addi a1, a1, 1\n\t"
        "blt a1, %[high], loop%=\n\t"
        "loop_end%=:\n\t"
        : [i] "+r" (i)
        : [array] "r" (array), [pivot] "r" (pivot), [high] "r" (high), [low] "r" (low)
        : "a1", "a2", "a3", "a4", "a5", "a6", "a7", "a0", "s2", "cc", "memory"
    );
    swap(&array[i], &array[high]);
    return (i);
}
#elif defined(VARIANT) && VARIANT == 3
// Branchy ASM
int partition(int array[], int low, int high) {
    int pivot = median_of_three_of_median_of_three(array, low, high);

    int i = low;
    asm volatile(
        "mv a1, %[low]\n\t"
        // Loop start
        "bge a1, %[high], loop_end%=\n"
        "loop%=:\n\t"
        // Get the array element at index a1
        "slli a2, a1, 2\n\t"
        "add a2, a2, %[array]\n\t"
        "lw a3, 0(a2)\n"
        // Branchy variant
        "problematicbranch%=:\n\t"
        // Compare the array element with the pivot
        "bgt a3, %[pivot], loop_increment%=\n\t"
        // Swap the array element j with the element i
        "slli a6, %[i], 2\n\t"
        "add a6, a6, %[array]\n\t"
        "lw a7, 0(a6)\n\t"
        "sw a7, 0(a2)\n\t"
        "sw a3, 0(a6)\n"
        // Increase i
        "addiw %[i], %[i], 1\n\t"
        // Increase a1 for the next iteration
        "loop_increment%=:\n\t"
        "addiw a1, a1, 1\n\t"
        "blt a1, %[high], loop%=\n\t"
        "loop_end%=:\n\t"
        : [i] "+r" (i)
        : [array] "r" (array), [pivot] "r" (pivot), [high] "r" (high), [low] "r" (low)
        : "a1", "a2", "a3", "a6", "a7", "cc", "memory"
    );
    swap(&array[i], &array[high]);
    return (i);
}
#elif defined(VARIANT) && VARIANT == 4
// Branch free ASM
int partition(int array[], int low, int high) {
    int pivot = median_of_three_of_median_of_three(array, low, high);

    int i = low;
    asm volatile(
        "mv a1, %[low]\n\t"
        // Loop start
        "bge a1, %[high], loop_end%=\n"
        "loop%=:\n\t"
        // Get the array element at index a1
        "slli a2, a1, 2\n\t"
        "add a2, a2, %[array]\n\t"
        "lw a3, 0(a2)\n"
        // Branch free variant
        "branchfree%=:\n\t"
        // Load the array element at index i
        "slli a6, %[i], 2\n\t"
        "add a6, a6, %[array]\n\t"
        "lw a7, 0(a6)\n\t"
        // Set bit based on comparison with pivot
        "sltu a5, a3, %[pivot]\n\t"
        // Swap indexes based on bit
        "sub a4, a2, a6\n\t"
        "mul a4, a4, a5\n\t"
        "add a6, a6, a4\n\t"
        "sub a2, a2, a4\n\t"
        // Write back to array
        "sw a3, 0(a2)\n\t"
        "sw a7, 0(a6)\n\t"
        // Increase i
        "add %[i], %[i], a5\n"
        // Increase a1 for the next iteration
        "loop_increment%=:\n\t"
        "addiw a1, a1, 1\n\t"
        "blt a1, %[high], loop%=\n\t"
        "loop_end%=:\n\t"
        : [i] "+r" (i)
        : [array] "r" (array), [pivot] "r" (pivot), [high] "r" (high), [low] "r" (low)
        : "a1", "a2", "a3", "a4", "a5", "a6", "a7", "cc", "memory"
    );
    swap(&array[i], &array[high]);
    return (i);
}
#elif defined(VARIANT) && VARIANT == 5
// Branch free ASM
int partition(int array[], int low, int high) {
    int pivot = median_of_three_of_median_of_three(array, low, high);

    int i = low;
    asm volatile(
        // a1 == j == low
        "mv a1, %[low]\n\t"
        // Loop start
        "bge a1, %[high], loop_end%=\n"
        "loop%=:\n\t"
        // Get the array element at index a1
        "slli a2, a1, 2\n\t"
        "add a2, a2, %[array]\n\t"
        "lw a3, 0(a2)\n"
        // Branch free variant
        "branchfree%=:\n\t"
        // Load the array element at index i
        "slli a6, %[i], 2\n\t"
        "add a6, a6, %[array]\n\t"
        "lw a7, 0(a6)\n\t"
        "sltu a5, %[pivot], a3\n\t"
        // Write from index a1 to index i when a3 < pivot
        "czero.eqz a4, a3, a5\n\t"
        "czero.nez a0, a7, a5\n\t"
        "or a4, a4, a0\n\t"
        "sw a4, 0(a2)\n\t"
        // Write from index i to index a1 when a3 < pivot
        "czero.eqz a4, a7, a5\n\t"
        "czero.nez a0, a3, a5\n\t"
        "or a4, a4, a0\n\t"
        "sw a4, 0(a6)\n\t"
        // Increase i
        "addi %[i], %[i], 1\n\t"
        "sub %[i], %[i], a5\n\t"
        // Increase a1 for the next iteration
        "loop_increment%=:\n\t"
        "addiw a1, a1, 1\n\t"
        "blt a1, %[high], loop%=\n\t"
        "loop_end%=:\n\t"
        : [i] "+r" (i)
        : [array] "r" (array), [pivot] "r" (pivot), [high] "r" (high), [low] "r" (low)
        : "a1", "a2", "a3", "a4", "a5", "a6", "a7", "a0", "s2", "cc", "memory"
    );
    swap(&array[i], &array[high]);
    return (i);
}
#elif defined(VARIANT) && VARIANT == 6
// BMS with VARIANT 4
int partition(int array[], int low, int high) {
    int pivot = median_of_three_of_median_of_three(array, low, high);

    int i = low;
    asm volatile(
        // a1 == j == low
        "mv a1, %[low]\n\t"
        // Loop start
        "bge a1, %[high], loop_end%=\n"
        "loop%=:\n\t"
        // Get the array element at index a1
        "slli a2, a1, 2\n\t"
        "add a2, a2, %[array]\n\t"
        "lw a3, 0(a2)\n"
        "bms problematicbranch%=,branchfree%=,19\n\t"
        // Branchy variant
        "problematicbranch%=:\n\t"
        // Compare the array element with the pivot
        "bgt a3, %[pivot], branchy_loop_increment%=\n\t"
        // Swap the array element j with the element i
        "slli a6, %[i], 2\n\t"
        "add a6, a6, %[array]\n\t"
        "lw a7, 0(a6)\n\t"
        "sw a7, 0(a2)\n\t"
        "sw a3, 0(a6)\n"
        // Increase i
        "addiw %[i], %[i], 1\n\t"
        // Increase a1 for the next iteration
        "branchy_loop_increment%=:\n\t"
        "addi a1, a1, 1\n\t"
        "blt a1, %[high], loop%=\n\t"
        "jal loop_end%=\n"
        // Branch free variant
        "branchfree%=:\n\t"
        // Load the array element at index i
        "slli a6, %[i], 2\n\t"
        "add a6, a6, %[array]\n\t"
        "lw a7, 0(a6)\n\t"
        "sltu a5, a3, %[pivot]\n\t"
        "bjr a5, problematicbranch%=, branchy_loop_increment%=\n\t"
        // Swap indexes based on bit
        "sub a4, a2, a6\n\t"
        "mul a4, a4, a5\n\t"
        "add a6, a6, a4\n\t"
        "sub a2, a2, a4\n\t"
        // Write back to array
        "sw a3, 0(a2)\n\t"
        "sw a7, 0(a6)\n\t"
        // Increase i
        "add %[i], %[i], a5\n"
        // Increase a1 for the next iteration
        "loop_increment%=:\n\t"
        "addi a1, a1, 1\n\t"
        "blt a1, %[high], loop%=\n\t"
        "loop_end%=:\n\t"
        : [i] "+r" (i)
        : [array] "r" (array), [pivot] "r" (pivot), [high] "r" (high), [low] "r" (low)
        : "a1", "a2", "a3", "a4", "a5", "a6", "a7", "a0", "s2", "cc", "memory"
    );
    swap(&array[i], &array[high]);
    return (i);
}
#endif

void quicksort(int array[], int low, int high) {

    if (low < high) {
        if (high - low > 2) {
            int pivot;
            pivot = partition(array, low, high);

            quicksort(array, low, pivot - 1);
            quicksort(array, pivot + 1, high);

        } else {
            insertionSort(array + low, high - low + 1);
        }
    }
}

int main(void) {
    int arrayToSort[ELEMENTS];
    // Initialize array of crandom bits
    #if defined(PREDICTABLE) && PREDICTABLE == 0
    unsigned int seed = 0;
    srand(seed);
    for (int i = 0; i < ELEMENTS; i++) {
        //seed = (seed * 1103515245 + 12345) & 0xffffffff;
        arrayToSort[i] = rand() % 10000;
    }
    #elif defined(PREDICTABLE) && PREDICTABLE == 1
    for (int i = 0; i < ELEMENTS; i++) {
        arrayToSort[i] = i;
    }
    #endif

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

    quicksort(arrayToSort, 0, ELEMENTS - 1);

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

    // Check if the array is sorted
    for (int i = 0; i < ELEMENTS - 1; i++) {
        if (arrayToSort[i] > arrayToSort[i + 1]) {
            log("Array is not sorted\n");
            return 1;
        }
    }

    log("Cycles: %lld\n", cyclesEnd - cyclesStart);

    return 0;
}