#include <stdio.h>
#include <stdlib.h>
#include <math.h>

int max(int a, int b) {
    return a > b ? a : b;
}

class Node {
    public:
    int value;
    bool isRoot = false;
    Node *left;
    Node *right;
    Node *parent;

    Node(int value) {
        this->value = value;
        this->left = NULL;
        this->right = NULL;
        this->parent = NULL;
    }


    __attribute__((noinline))
    Node* findValue(int value) {
        if (this->value == value) {
            return this;
        }
        if (value < this->value) {
            if (left == NULL) {
                return NULL;
            }
            return left->findValue(value);
        } else {
            if (right == NULL) {
                return NULL;
            }
            return right->findValue(value);
        }
    }

    void insert(Node* newChild) {
        if (newChild->value < value) {
            if (left == NULL) {
                left = newChild;
                newChild->parent = this;
                rebalanceBottomUp();
            } else {
                left->insert(newChild);
            }
        } else {
            if (right == NULL) {
                right = newChild;
                newChild->parent = this;
                rebalanceBottomUp();
            } else {
                right->insert(newChild);
            }
        }
    }

    int height() {
        if (left == NULL && right == NULL) {
            return 1;
        }
        if (left != NULL && right != NULL) {
            return max(left->height(), right->height()) + 1;
        }
        if (left != NULL) {
            return left->height() + 1;
        }
        if (right != NULL) {
            return right->height() + 1;
        }
    }

    void rebalanceBottomUp() {
        if (isRoot) {
            return;
        }
        rebalance();
        parent->rebalanceBottomUp();
        
    }

    void rebalance() {
        if (isRoot) {
            return;
        }
        if (left != NULL && right != NULL) {
            int leftHeight = left->height();
            int rightHeight = right->height();
            if (abs(leftHeight - rightHeight) <= 1) {
                return;
            }
            if (rightHeight > leftHeight) {
                rotateLeft();
            } else {
                rotateRight();
            }
        } else if (left != NULL && left->height() > 1) {
            rotateRight();
        } else if (right != NULL && right->height() > 1) {
            rotateLeft();
        }
    }

    void rotateLeft() {
        Node *previousRight = right;
        right = right->left;
        if (right != NULL) {
            right->parent = this;
        }
        if (previousRight->value < parent->value) {
            parent->left = previousRight;
        } else {
            parent->right = previousRight;
        }
        previousRight->parent = parent;
        previousRight->left = this;
        parent = previousRight;
    }

    void rotateRight() {
        Node *previousLeft = left;
        left = left->right;
        if (left != NULL) {
            left->parent = this;
        }
        if (previousLeft->value < parent->value) {
            parent->left = previousLeft;
        } else {
            parent->right = previousLeft;
        }
        previousLeft->parent = parent;
        previousLeft->right = this;
        parent = previousLeft;
    }

    void print(int level = 0) {
        printf("%*s", level, "");
        printf("%d\n", value);
        if (left != NULL) {
            left->print(level + 1);
        }
        if (right != NULL) {
            right->print(level + 1);
        }
    }
};

class RootNode: public Node {
    public:
    RootNode(): Node(__INT_MAX__) {
        isRoot = true;
    }

    void insert(Node* newChild) {
        if (left == NULL) {
            left = newChild;
            newChild->parent = this;
        } else {
            left->insert(newChild);
        }
    }

    void print() {
        if (left != NULL) {
            left->print();
        }
    }

    void rebalance() {
        return;
    }

    void rebalanceBottomUp() {
        return;
    }
    
    void rotateRight() {
        return;
    }

    void rotateLeft() {
        return;
    }
};

#if defined(VARIANT) && VARIANT == 0
__attribute__((noinline))
Node* findValueLoop(Node* node, int value) {
    Node* currentNode = node;

    while (true) {
        if (currentNode == NULL) {
            return NULL;
        }
        if (currentNode->value == value) {
            return currentNode;
        }
        if (currentNode->value > value) {
            currentNode = currentNode->left;
        } else {
            currentNode = currentNode->right;
        }
    }
    return NULL;
}
#elif defined(VARIANT) && VARIANT == 1
__attribute__((noinline))
Node* findValueLoop(Node* node, int value) {
    Node* currentNode = node;

    asm volatile(
        "beqz %[currentNode], end%=\n\t"
        "loop%=:"
        "lw a4, 0x0(%[currentNode])\n\t"
        "beq a4, %[value], end%=\n\t"
        "blt %[value], a4, less_than%=\n\t"
        "ld %[currentNode], 0x10(%[currentNode])\n\t"
        "bnez %[currentNode], loop%=\n\t"
        "j end%=\n\t"
        "less_than%=:\n\t"
        "ld %[currentNode], 0x8(%[currentNode])\n\t"
        "bnez %[currentNode], loop%=\n\t"
        "end%=:"
        : [currentNode] "+r" (currentNode)
        : [value] "r" (value)
        : "a4", "cc", "memory"
    );
    return currentNode;
}
#elif defined(VARIANT) && VARIANT == 2
__attribute__((noinline))
Node* findValueLoop(Node* node, int value) {
    Node* currentNode = node;
    asm volatile(
        "beqz	%[currentNode], end%=\n\t"
        "li	a3, 0x8\n"
        "loop%=:"
        "lw a4, 0x0(%[currentNode])\n\t"
        "beq a4, %[value], end%=\n\t"
        "slt a4, %[value], a4\n\t"
        "czero.nez a4, a3, a4\n\t"
        "add %[currentNode], %[currentNode], a4\n\t"
        "ld %[currentNode], 0x8(%[currentNode])\n\t"
        "bnez %[currentNode], loop%=\n\t"
        "end%=:"
        : [currentNode] "+r" (currentNode)
        : [value] "r" (value)
        : "a3", "a4", "cc", "memory"
    );
    return currentNode;
}
#elif defined(VARIANT) && VARIANT == 3
__attribute__((noinline))
Node* findValueLoop(Node* node, int value) {
    Node* currentNode = node;
    asm volatile(
        "beqz	%[currentNode], end%=\n\t"
        "li	a3, 0x8\n"
        "loop%=:"
        "lw a4, 0x0(%[currentNode])\n\t"
        "beq a4, %[value], end%=\n\t"

        "bms problematicbranch%=,branchfree%=,12\n"

        "problematicbranch%=:\n\t"
        "blt %[value], a4, less_than%=\n\t"
        "ld %[currentNode], 0x10(%[currentNode])\n\t"
        "bnez %[currentNode], loop%=\n\t"
        "j end%=\n\t"
        "less_than%=:\n\t"
        "ld %[currentNode], 0x8(%[currentNode])\n\t"
        "bnez %[currentNode], loop%=\n\t"
        "j end%=\n"

        "branchfree%=:\n\t"
        "slt a4, %[value], a4\n\t"
        "bjr a4, problematicbranch%=, less_than%=\n\t"
        "czero.nez a4, a3, a4\n\t"
        "add %[currentNode], %[currentNode], a4\n\t"
        "ld %[currentNode], 0x8(%[currentNode])\n\t"
        "bnez %[currentNode], loop%=\n"

        "end%=:"
        : [currentNode] "+r" (currentNode)
        : [value] "r" (value)
        : "a3", "a4", "cc", "memory"
    );
    return currentNode;
}
#endif

#define ELEMENTS_IN_TREE 64

int main() {
    RootNode* root = new RootNode();

    for (int i = 0; i < ELEMENTS_IN_TREE; i++) {
        root->insert(new Node(i));
    }

    srand(0);
    int randomNumbers[ELEMENTS];
    for (int i = 0; i < ELEMENTS; i++) {
        randomNumbers[i] = rand() % ELEMENTS_IN_TREE;
    }
    randomNumbers[3] = 0;

    // This guarantees that there is a 0 in the first 1024 numbers
    int smallest = __INT_MAX__;

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

    for (int i = 0; i < ELEMENTS; i++) {
        #if defined(PREDICTABLE) && PREDICTABLE == 0
        int randomNumber = randomNumbers[i];
        #else
        int randomNumber = randomNumbers[(i % 4)];
        #endif
        
        #if defined(SETUP) && SETUP == 1
        asm volatile("nop":::"memory");
        Node* node = root;
        #else
        Node* node = findValueLoop(root, randomNumber);
        #endif

        if (node == NULL) {
            printf("Did not find node %d\n", randomNumber);
            return 1;
        }
        if (node->value < smallest) {
            smallest = node->value;
        }
    }

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

    #if defined(SETUP) && SETUP == 1
    if (smallest != root->value) {
    #else
    if (smallest != 0) {
    #endif
        printf("Did not get smallest\n");
        return 1;
    }

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

    return 0;
}