#include <algorithm>
#include <glm/glm.hpp>
#include <iostream>
#include <string>
#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) {}

    Point operator+(const Point &other) const {
        return Point(new float[3]{coordinates[0] + other.coordinates[0],
                                  coordinates[1] + other.coordinates[1],
                                  coordinates[2] + other.coordinates[2]},
                     nullptr);
    }

    Point operator*(float scalar) const {
        return Point(
            new float[3]{coordinates[0] * scalar, coordinates[1] * scalar, coordinates[2] * scalar},
            nullptr);
    }

    float *coordinates;

    Triangle *triangle;
};

struct Ray {
    Ray(float origin[3], float direction[3]) : origin(origin), direction(direction) {}

    float *origin;

    float *direction;
};

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

    ~KDTree() = default; // TODO: Delete all allocated Nodes

    Point *intersect_ray(Ray ray) { return intersect_ray_recurse(ray, root); }

    std::string to_string() {
        std::string str = "";
        to_string_recurse(str, root, 0);
        return str;
    }

  private:
    Node *root;

    int MAX_DEPTH = 500;

    // Returns a comparator lambda for assessing which of the two points has a
    // greater coordinate in the given axis.
    auto get_point_comparator(int axis) {
        return [axis](Point *p1, Point *p2) {
            return p1->coordinates[axis] < p2->coordinates[axis];
        };
    }

    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 = get_point_comparator(it_axis);

            Point *min = *std::min_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;

        std::nth_element(points.begin(), points.begin() + middle, points.end(),
                         get_point_comparator(axis));

        Point *median = points[middle];

        // 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 + 1, 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));
    }

    Point *intersect_ray_recurse(Ray ray, Node *node) {
        // Exit condition: There was no collision
        if (node == nullptr) { return nullptr; }

        // Is the left or right child node closer to this point?
        Node *near = ray.origin[node->axis] > node->point->coordinates[node->axis] ? node->right
                                                                                   : node->left;
        Node *far = near == node->right ? node->left : node->right;

        std::cout << "Checking " << node->point->coordinates[0] << ", "
                  << node->point->coordinates[1] << ", " << node->point->coordinates[2] << ", "
                  << std::endl;

        // Intersect ray with the point's splitting plane

        // Are they parallel? If so, recurse only to the nearer side
        if (ray.direction[node->axis] == 0.0) {
            return intersect_ray_recurse(ray, near);
        } else {
            // They are not parallel, so check where the intersection occurs
            float t = (node->point->coordinates[node->axis] - ray.origin[node->axis]) /
                      ray.direction[node->axis];

            if (t >= 0.0) {
                // t is positive, so we need to recurse to both children if the nearer one does not
                // result in a collision
                Point *first_attempt = intersect_ray_recurse(ray, near);

                if (first_attempt != nullptr) {
                    return first_attempt;
                } else {
                    // The first attempt did not work, so recurse to the other side too.
                    // For this, calculate a new ray origin ...
                    float new_origin[3]{ray.origin[0] + t * ray.direction[0],
                                        ray.origin[1] + t * ray.direction[1],
                                        ray.origin[2] + t * ray.direction[2]};

                    // ... and continue towards that direction, but with the new origin (we can
                    // leave behind what we already checked)
                    return intersect_ray_recurse(Ray(new_origin, ray.direction), far);
                }
            } else {
                // We only have to check the nearer one, as the other side can't be reached by the
                // ray
                return intersect_ray_recurse(ray, near);
            }
        }
    }

    void to_string_recurse(std::string &str, Node *node, int depth) {
        if (node == nullptr) { return; }

        Point *point = node->point;

        str += std::string(depth, '-') + std::to_string(point->coordinates[0]) + ", " +
               std::to_string(point->coordinates[1]) + ", " +
               std::to_string(point->coordinates[2]) + " with axis " + std::to_string(node->axis) +
               "\n";

        to_string_recurse(str, node->left, depth + 1);
        to_string_recurse(str, node->right, depth + 1);
    }
};