#include <algorithm>
#include <glm/glm.hpp>
#include <vector>

// Forward declarations
struct Node;
struct Point;
struct Triangle;

struct Node {
  Node(int axis, Point *point, Node *left, Node *right)
      : axis(axis), point(point), left(left), right(right) {}

  int axis;

  Point *point;

  Node *left;
  Node *right;
};

struct Triangle {
  Triangle(Point *p1, Point *p2, Point *p3) : p1(p1), p2(p2), p3(p3) {}

  Point *p1;
  Point *p2;
  Point *p3;
};

struct Point {
  Point(float coordinates[3], Triangle *triangle)
      : coordinates(coordinates), triangle(triangle) {}

  float *coordinates;

  Triangle *triangle;
};

class KDTree {
public:
  KDTree(std::vector<Point *> points) { root = build(points, 0); }

  ~KDTree() = default;

private:
  Node *root;

  int MAX_DEPTH = 500;

  Node *build(std::vector<Point *> points, int depth) {
    // Exit conditions
    if (points.empty() || depth > MAX_DEPTH) {
      return nullptr;
    }

    // Select axis by choosing the one with maximal extent
    float max_extent = 0;
    int axis = 0;

    for (int it_axis = 0; it_axis < 3; it_axis++) {
      // Get extent along this axis
      auto comparator = [it_axis](Point *p1, Point *p2) {
        return p1->coordinates[it_axis] < p2->coordinates[it_axis];
      };

      Point *min = *std::max_element(points.begin(), points.end(), comparator);
      Point *max = *std::max_element(points.begin(), points.end(), comparator);

      float extent = max->coordinates[it_axis] - min->coordinates[it_axis];

      // Is it greater than max_extent?
      if (extent > max_extent) {
        // If so, make this the splitting axis
        max_extent = extent;
        axis = it_axis;
      }
    }

    // Choose the median as the pivot and sort the points into
    // left-of-median and right-of-median using nth_element
    int middle = points.size() / 2;

    // TODO: Code duplication from the comparator in the earlier axis assessment
    // loop
    std::nth_element(points.begin(), points.begin() + middle, points.end(),
                     [axis](Point *p1, Point *p2) {
                       return p1->coordinates[axis] < p2->coordinates[axis];
                     });

    Point *median;
    float pivot;

    // TODO: This copies. Can we split the vector into two without copying?
    std::vector<Point *> left_of_median(points.begin(),
                                        points.begin() + middle);
    std::vector<Point *> right_of_median(points.begin() + middle, points.end());

    // Create node, recursively call to construct subtree
    return new Node(axis, median, build(left_of_median, depth + 1),
                    build(right_of_median, depth + 1));
  }
};