#include "geometry.h"
#include <algorithm>
#include <chrono>
#include <glm/glm.hpp>
#include <iostream>
#include <string>
#include <vector>

class KDTree {
  public:
    KDTree(std::vector<Point *> points) {
        std::cout << "Starting to build tree with " << points.size() << " points took ";
        auto start = std::chrono::high_resolution_clock::now();
        root = build(points, 0);
        auto end = std::chrono::high_resolution_clock::now();

        std::cout << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
                  << " microseconds" << std::endl;
    }

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

    const Triangle *intersect_ray(const Ray ray, Vector &result) {
        auto start = std::chrono::high_resolution_clock::now();

        float nearest = 1000.0; // Initial max distance
        const Triangle *nearest_triangle = nullptr;
        intersect_ray_recurse(nearest_triangle, result, ray, root, 0, nearest);

        auto end = std::chrono::high_resolution_clock::now();

        std::cout << "Intersection took "
                  << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
                  << " microseconds" << std::endl;
        return nearest_triangle;
    }

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

    std::vector<Vector> get_lines() {
        std::vector<Vector> lines;
        float mins[3] = {-1000, -1000, -1000};
        float maxs[3] = {1000, 1000, 1000};

        get_lines_recursive(lines, root, mins, maxs);

        return lines;
    }

  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->pos[axis] < p2->pos[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->pos[it_axis] - min->pos[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));
    }

    void intersect_ray_recurse(const Triangle *&nearest_triangle, Vector &result, Ray ray,
                               Node *node, int depth, float &nearest) {
        // Exit condition: There was no collision
        if (node == nullptr) { return; }

        // Check for a collision here
        for (const Triangle *triangle : node->point->triangles) {
            Vector current_result(0, 0, 0);
            float current_distance;
            if (ray.intersects_triangle(triangle, current_result, current_distance)) {
                if (current_distance < nearest) {
                    nearest = current_distance;
                    nearest_triangle = triangle;
                    result = current_result;
                }
            }
        }

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

        if (ray.direction[node->axis] == 0.0) {
            intersect_ray_recurse(nearest_triangle, result, ray, near, depth + 1, nearest);
        } else {
            float t =
                (node->point->pos[node->axis] - ray.origin[node->axis]) / ray.direction[node->axis];

            // Check this side for intersections up to the distance of the currently best
            // intersection
            intersect_ray_recurse(nearest_triangle, result, ray, near, depth + 1, nearest);

            // If the far side is closer than the distance to the best current intersection, check
            // that side too
            if (t < nearest) {
                intersect_ray_recurse(nearest_triangle, result, ray, far, depth + 1, nearest);
            }
        }
    }

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

        Point *point = node->point;

        str += std::string(depth * 2, ' ') + std::to_string(point->pos[0]) + ", " +
               std::to_string(point->pos[1]) + ", " + std::to_string(point->pos[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);
    }

    void get_lines_recursive(std::vector<Vector> &lines, Node *node, float mins[3], float maxs[3]) {
        if (node == nullptr) return;

        int axis = node->axis;

        // Add lines for this node
        // We add a whole bounding box, so we need 8 points -> 12 lines
        lines.emplace_back(Vector(mins[0], mins[1], mins[2]));
        lines.emplace_back(Vector(maxs[0], mins[1], mins[2]));

        lines.emplace_back(Vector(maxs[0], mins[1], mins[2]));
        lines.emplace_back(Vector(maxs[0], maxs[1], mins[2]));

        lines.emplace_back(Vector(maxs[0], maxs[1], mins[2]));
        lines.emplace_back(Vector(mins[0], maxs[1], mins[2]));

        lines.emplace_back(Vector(mins[0], maxs[1], mins[2]));
        lines.emplace_back(Vector(mins[0], mins[1], mins[2]));

        lines.emplace_back(Vector(mins[0], mins[1], maxs[2]));
        lines.emplace_back(Vector(maxs[0], mins[1], maxs[2]));

        lines.emplace_back(Vector(maxs[0], mins[1], maxs[2]));
        lines.emplace_back(Vector(maxs[0], maxs[1], maxs[2]));

        lines.emplace_back(Vector(maxs[0], maxs[1], maxs[2]));
        lines.emplace_back(Vector(mins[0], maxs[1], maxs[2]));

        lines.emplace_back(Vector(mins[0], maxs[1], maxs[2]));
        lines.emplace_back(Vector(mins[0], mins[1], maxs[2]));

        lines.emplace_back(Vector(mins[0], mins[1], mins[2]));
        lines.emplace_back(Vector(mins[0], mins[1], maxs[2]));

        lines.emplace_back(Vector(maxs[0], mins[1], mins[2]));
        lines.emplace_back(Vector(maxs[0], mins[1], maxs[2]));

        lines.emplace_back(Vector(mins[0], maxs[1], mins[2]));
        lines.emplace_back(Vector(mins[0], maxs[1], maxs[2]));

        lines.emplace_back(Vector(maxs[0], maxs[1], mins[2]));
        lines.emplace_back(Vector(maxs[0], maxs[1], maxs[2]));

        // TODO: We don't actually have to add all as some are always going to overlap with the
        // previous lines...

        // Recurse on the left
        float new_maxs[3] = {maxs[0], maxs[1], maxs[2]}; // Copy initially
        new_maxs[axis] = node->point->pos[axis];
        get_lines_recursive(lines, node->left, mins, new_maxs);

        // Recurse on the right
        float new_mins[3] = {mins[0], mins[1], mins[2]}; // Copy initially
        new_mins[axis] = node->point->pos[axis];
        get_lines_recursive(lines, node->right, new_mins, maxs);
    }
};