#pragma once
#ifndef OCTREE_H
#define OCTREE_H


#include <cstdlib>
#include <cmath>
#include <iostream>
#include <algorithm>
#include "matconv.h"
#include "vtkwriter.h"

#ifdef _OPENMP
#include <omp.h>
#endif

#include <map>
#include <vector>
#include <array>
#include <fstream>
#include <string>
#include <sstream>
#include <stdexcept>
#include <iostream>
#include "tetra.h"
#include "vtk-5.0/vtkTetra.h"

#include <unordered_map>
#include <unordered_set>

// Axis-Aligned Bounding Box (AABB) structure
struct AABB 
{
    double xmin, xmax, ymin, ymax, zmin, zmax; // Min and max bounds in x, y, z

    // Check if a point (x, y, z) lies inside the box
    bool contains(double x, double y, double z) const  
    {
        return (x >= xmin && x <= xmax &&
                y >= ymin && y <= ymax &&
                z >= zmin && z <= zmax);
    }

    bool contains_with_buffer(double x, double y, double z, double buffer_ratio = 0.01) const  
    {
        double dx = (xmax - xmin) * buffer_ratio;
        double dy = (ymax - ymin) * buffer_ratio;
        double dz = (zmax - zmin) * buffer_ratio;

        return (x >= xmin - dx && x <= xmax + dx &&
                y >= ymin - dy && y <= ymax + dy &&
                z >= zmin - dz && z <= zmax + dz);
    }


};

struct OctreeNode 
{
    AABB box;                       // Bounding box of this node
    std::vector<int> tetra_ids;     // List of tetrahedron indices (only if leaf)
    std::vector<int> NC_tetra_ids;  // List of NC tetrahedron indices (only if leaf)
    OctreeNode* children[8] = {nullptr}; // Child pointers (nullptr if not exist)

    // Check if current node is a leaf node
    bool is_leaf() const 
    {
        return (children[0] == nullptr);
    }
};


class Octree 
{
    // This class implements an octree data structure for spatial partitioning
    // of tetrahedral elements in 3D space. It is used to efficiently query
    // which tetrahedra contain a given point in space.
    //
    // The octree is built from a set of tetrahedra, and each node in the
    // octree represents a bounding box that contains a subset of the tetrahedra.
    //
    // The class provides methods to build the octree, find tetrahedra
    // containing a point, and validate the contents of the leaves of the octree.
    public:

        friend class FemGrp;  // Allow FemGrp to access private members

        // Public members
        const tetra*  tet_ptr;;  // Pointer to the tetrahedron array in FEMgrp    
        OctreeNode* root_octree_node;

        // For tetrahedra with conformal faces
        int tet_count;
        std::vector<AABB> tetra_boxes;  // Allocate vector to hold all bounding boxes

        // For tetrahedra with non-conformal faces
        int NC_tet_count;
        std::vector<AABB> NC_tetra_boxes; 


        Octree() : root_octree_node(nullptr), tet_ptr(nullptr) {};


        // Octree building for tetrahedra with conformal faces only
        void buildOctree(const std::vector<int>& tet_ids, const AABB& box, double buffer_distance, int depth, int max_depth);
        OctreeNode* buildOctree_function(const std::vector<int>& tet_ids, const AABB& box, double buffer_distance, int depth, int max_depth);

        // Octree building for tetrahedra with non-conformal faces too
        void buildOctree_withNCFLAGS(const std::vector<int>& tet_ids, const std::vector<int>& NC_tet_ids, const AABB& box, double buffer_distance, int depth, int max_depth); 
        OctreeNode* buildOctree_withNCFLAGS_function(const std::vector<int>& tet_ids, const std::vector<int>& NC_tet_ids, const AABB& box, double buffer_distance, int depth, int max_depth); 


        // Finds all tetrahedra in the octree that the point could lie in 
        // (within a tolerance), even if it lies on shared faces, edges, or vertices.
        void remove_duplicate_tetra_matches(std::vector<std::pair<int, std::array<double, 4>>>& matches);
        bool findTetraInOctree( const double probe_xyz[3],
                                std::vector<std::pair<int, std::array<double, 4>>>& matches,
                                double tol = 1e-6);
        bool findTetraInOctree_function( const OctreeNode* node,
                                    const double probe_xyz[3],
                                   std::vector<std::pair<int, std::array<double, 4>>>& matches,
                                    double tol = 1e-6);


        // Need function to find all the tetrahedral that contains 3 probes / nodes
        void remove_duplicate_tetra_id_matches(std::vector<int>& matches);     
        bool findNCTetraInOctree(const AABB& face_box, std::vector<int>& matches);
        bool findNCTetraInOctree_function(const OctreeNode* node, const AABB& face_box, std::vector<int>& matches);
        
                                            


        void validateAllTetsInLeaves(OctreeNode* node, std::vector<bool>& tet_found_flags); 

    private:
        void deleteOctree(OctreeNode* node);

};

#endif // OCTREE_H