#include "avl/avl.h"
#include "bst/bst.h"
#include "list/list.h"

#include <algorithm>
#include <chrono>
#include <cstdio>
#include <iostream>
#include <ostream>
#include <random>
#include <set>

void measureList(std::vector<int> *sequence, FILE *file) {
  float buildTime = 0, searchTime = 0, deleteTime = 0;

  for (int i = 0; i < 10; i++) {
    std::cout << "- Running list " << i + 1 << "/10\r";
    fflush(stdout);

    // Measure build time
    auto start = std::chrono::high_resolution_clock::now();
    List *head = nullptr;
    for (int value : *sequence) {
      head = insert(head, value);
    }
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed_seconds = end - start;
    buildTime += elapsed_seconds.count() * 1000; // Convert to milliseconds

    // Now we run the search
    start = std::chrono::high_resolution_clock::now();
    for (int value : *sequence) {
      search(head, value);
    }
    end = std::chrono::high_resolution_clock::now();
    elapsed_seconds = end - start;
    searchTime += elapsed_seconds.count() * 1000; // Convert to milliseconds

    // Measure deletion time for list
    start = std::chrono::high_resolution_clock::now();
    while (head != nullptr) {
      head = remove(head);
    }
    end = std::chrono::high_resolution_clock::now();
    elapsed_seconds = end - start;
    deleteTime += elapsed_seconds.count() * 1000; // Convert to milliseconds
  }

  buildTime /= 10;
  searchTime /= 10;
  deleteTime /= 10;

  std::cout << "- List built in " << buildTime << "ms | searched in "
            << searchTime << "ms | deleted in " << deleteTime << "ms"
            << std::endl;
  fprintf(file, "List,%f,%f,%f\n", buildTime, searchTime, deleteTime);
}

void measureBST(std::vector<int> *sequence, FILE *file) {
  float buildTime = 0, searchTime = 0, deleteTime = 0;

  for (int i = 0; i < 10; i++) {
    std::cout << "- Running bst " << i + 1 << "/10\r";
    fflush(stdout);

    // Measure build time
    auto start = std::chrono::high_resolution_clock::now();
    Tree *root = nullptr;
    for (int value : *sequence) {
      root = insert(root, value);
    }
    // root = balance(root);
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed_seconds = end - start;
    buildTime += elapsed_seconds.count() * 1000; // Convert to milliseconds

    // Now we run the search
    start = std::chrono::high_resolution_clock::now();
    for (int value : *sequence) {
      search(root, value);
    }
    end = std::chrono::high_resolution_clock::now();
    elapsed_seconds = end - start;
    searchTime += elapsed_seconds.count() * 1000; // Convert to milliseconds

    // Measure deletion time for tree
    start = std::chrono::high_resolution_clock::now();
    deleteTree(root);
    end = std::chrono::high_resolution_clock::now();
    elapsed_seconds = end - start;
    deleteTime += elapsed_seconds.count() * 1000; // Convert to milliseconds
  }

  buildTime /= 10;
  searchTime /= 10;
  deleteTime /= 10;

  std::cout << "- Tree built in " << buildTime << "ms | searched in "
            << searchTime << "ms | deleted in " << deleteTime << "ms"
            << std::endl;
  fprintf(file, "BST,%f,%f,%f\n", buildTime, searchTime, deleteTime);
}

void benchmarkAVL(std::vector<int> *sequence, FILE *file) {
  Tree *bst = nullptr;
  for (int value : *sequence) {
    bst = insert(bst, value);
  }

  int heightBST = getHeight(bst, 0);
  deleteTree(bst);

  AVL_tree *avl = nullptr;
  for (int value : *sequence) {
    avl = insertAVL(avl, value);
  }

  int heightAVL = getHeight(avl);
  deleteAVL(avl);

  fprintf(file, "%d,%d\n", heightBST, heightAVL);
}

int main() {
  // Init the random number generator
  std::random_device rd;
  std::mt19937 gen(rd());
  std::uniform_int_distribution<> dis(1, 1000000);

  int mode = 1;
  std::cout << "Select benchmark to run:\n 1 - Benchmark List vs BST\n 2 - "
               "Benchmark BST vs AVL"
            << std::endl;
  std::cin >> mode;

  if (mode < 1 && mode > 2) {
    mode = 1;
  }

  FILE *file;
  if (mode == 1) {
    // Open a file for writing results
    file = fopen("results.csv", "w");
    fprintf(file, "Structure,BuildTime,SearchTime,DeleteTime\n");
  } else {
    file = fopen("avl.csv", "w");
    std::fprintf(file, "Structure,Height\n");
  }

  for (int n = 1; n < 26; n++) {
    // Using a set here ensures that there are no duplicates
    std::set<int> sequence;
    while (sequence.size() < n * 1000) {
      sequence.insert(dis(gen));
    }

    std::vector<int> random_sequence_vec;
    for (int val : sequence) {
        random_sequence_vec.push_back(val);
    }

    // Then randomize the sequence so the bst isn't a list (Set keeps the elements in order)
    std::shuffle(random_sequence_vec.begin(), random_sequence_vec.end(), gen);

    // Display the times like a cascade
    std::cout << "Running tests for " << n * 1000 << " elements..."
              << std::endl;
    if (mode == 1) {
      measureList(&random_sequence_vec, file);
      measureBST(&random_sequence_vec, file);
    } else {
      benchmarkAVL(&random_sequence_vec, file);
    }
  }

  fclose(file);
  return 0;
}