#include "Octree.h"



void Octree::buildOctree(const std::vector<int>& tet_ids, const AABB& box, double buffer_distance, int depth, int max_depth) 
{
    // Build the octree recursively
    root_octree_node = buildOctree_function(tet_ids, box, buffer_distance, depth, max_depth);
}


OctreeNode* Octree::buildOctree_function(const std::vector<int>& tet_ids, const AABB& box, double buffer_distance, int depth, int max_depth) 
{

    OctreeNode* node = new OctreeNode();
    node->box = box;
    node->tetra_ids = tet_ids;

    // If small enough or depth limit reached, make leaf
    if (depth >= max_depth)
        return node;

    // Split current bounding box into 8 children        
    double midx = (box.xmin + box.xmax) * 0.5;
    double midy = (box.ymin + box.ymax) * 0.5;
    double midz = (box.zmin + box.zmax) * 0.5;

    AABB child_boxes[8];
    for (int i = 0; i < 8; ++i) {
        child_boxes[i] = {
            (i & 1) ? midx : box.xmin,
            (i & 1) ? box.xmax : midx,
            (i & 2) ? midy : box.ymin,
            (i & 2) ? box.ymax : midy,
            (i & 4) ? midz : box.zmin,
            (i & 4) ? box.zmax : midz
        };
    }

    // Prepare tetra IDs list for each child
    std::vector<int> child_tet_ids[8];

    for (int id : tet_ids) {
        const AABB& tet_box = tetra_boxes[id];

        for (int i = 0; i < 8; ++i) {
            const AABB& child = child_boxes[i];
            bool overlap_x = (tet_box.xmin <= child.xmax + buffer_distance && tet_box.xmax >= child.xmin - buffer_distance);
            bool overlap_y = (tet_box.ymin <= child.ymax + buffer_distance && tet_box.ymax >= child.ymin - buffer_distance);
            bool overlap_z = (tet_box.zmin <= child.zmax + buffer_distance && tet_box.zmax >= child.zmin - buffer_distance);

            if (overlap_x && overlap_y && overlap_z) {
                child_tet_ids[i].push_back(id);
            }
        }
    }

    // Recursively build children
    for (int i = 0; i < 8; ++i) {
        if (!child_tet_ids[i].empty()) {
            node->children[i] = buildOctree_function(child_tet_ids[i], child_boxes[i], buffer_distance, depth + 1, max_depth);
        }
    }

    return node;
}



void Octree::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) 
{
    // Build the octree recursively
    root_octree_node = buildOctree_withNCFLAGS_function(tet_ids, NC_tet_ids, box, buffer_distance, depth, max_depth);
}


OctreeNode* Octree::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) 
{

    OctreeNode* node = new OctreeNode();
    node->box = box;
    node->tetra_ids = tet_ids;   // All the Tetra IDS
    node->NC_tetra_ids = NC_tet_ids; // All the NC Tetra IDS

    // If small enough or depth limit reached, make leaf
    if (depth >= max_depth)
        return node;

    // Split current bounding box into 8 children        
    double midx = (box.xmin + box.xmax) * 0.5;
    double midy = (box.ymin + box.ymax) * 0.5;
    double midz = (box.zmin + box.zmax) * 0.5;

    AABB child_boxes[8];
    for (int i = 0; i < 8; ++i) {
        child_boxes[i] = {
            (i & 1) ? midx : box.xmin,
            (i & 1) ? box.xmax : midx,
            (i & 2) ? midy : box.ymin,
            (i & 2) ? box.ymax : midy,
            (i & 4) ? midz : box.zmin,
            (i & 4) ? box.zmax : midz
        };
    }

    // Prepare tetra IDs list for each child
    std::vector<int> child_tet_ids[8];
    std::vector<int> child_NC_tet_ids[8];

    for (int id : tet_ids) {
        const AABB& tet_box = tetra_boxes[id];

        for (int i = 0; i < 8; ++i) {
            const AABB& child = child_boxes[i];
            bool overlap_x = (tet_box.xmin <= child.xmax + buffer_distance && tet_box.xmax >= child.xmin - buffer_distance);
            bool overlap_y = (tet_box.ymin <= child.ymax + buffer_distance && tet_box.ymax >= child.ymin - buffer_distance);
            bool overlap_z = (tet_box.zmin <= child.zmax + buffer_distance && tet_box.zmax >= child.zmin - buffer_distance);

            if (overlap_x && overlap_y && overlap_z) 
            {
                child_tet_ids[i].push_back(id);

                // Check if id is in NC_tetra_ids
                if (std::find(node->NC_tetra_ids.begin(), node->NC_tetra_ids.end(), id) != node->NC_tetra_ids.end()) 
                {
                    child_NC_tet_ids[i].push_back(id); 
                }
            }
        }
    }

    // Recursively build children
    for (int i = 0; i < 8; ++i) {
        if (!child_tet_ids[i].empty()) {
            node->children[i] = buildOctree_withNCFLAGS_function(child_tet_ids[i], child_NC_tet_ids[i], child_boxes[i], buffer_distance, depth + 1, max_depth);
        }
    }

    return node;
}












// --------------------------------------------------------------------------------------------------------------------


void Octree::remove_duplicate_tetra_matches(std::vector<std::pair<int, std::array<double, 4>>>& matches) 
{
    // dictionary
    // Uses int as the key (in your case, the tetrahedron ID),
    // Stores a std::array<double, 4> as the value (your barycentric coordinates)
    std::unordered_map<int, std::array<double, 4>> unique_map;

    // Keep the first barycentric result for each unique tet_id
    for (const auto& match : matches) {
        int tet_id = match.first;
        if (unique_map.find(tet_id) == unique_map.end()) {
            unique_map[tet_id] = match.second;
        }
    }

    // Convert map back to vector
    matches.clear();
    for (const auto& entry : unique_map) {
        matches.emplace_back(entry.first, entry.second);
    }
}








bool Octree::findTetraInOctree( const double probe_xyz[3],
    std::vector<std::pair<int, std::array<double, 4>>>& matches,
    double tol)
{
    bool found = findTetraInOctree_function(root_octree_node, probe_xyz, matches, tol);

    // Remove duplicates
    remove_duplicate_tetra_matches(matches);

    return found;
}



bool Octree::findTetraInOctree_function( const OctreeNode* node,
                                const double probe_xyz[3],
                        std::vector<std::pair<int, std::array<double, 4>>>& matches,
                        double tol)
{
    if (!node || !node->box.contains_with_buffer(probe_xyz[0], probe_xyz[1], probe_xyz[2], 0.01))
        return false;

    bool found_any = false;

    // Perform check if the node is a leaf
    if (node->is_leaf()) 
    {
        for (int tet_id : node->tetra_ids) 
        {
            const tetra& tet = tet_ptr[tet_id];

            // Step 1: Use buffered AABB check for early exit
            const AABB& box = tetra_boxes[tet_id];
            if (!box.contains_with_buffer(probe_xyz[0], probe_xyz[1], probe_xyz[2], 0.01))  // 1% buffer
                continue;            
            
                
            // Step 2: Check if the point is near a corner
            // Detect if the probe point is very close to one of the tetrahedron's corner nodes, 
            // and if so, treat it as inside the tetrahedron and optionally snap it to that node 
            // (or tag the corresponding node as the match).
            int nearest_corner_idx = -1;

            // Compute the diagonal of the bounding box
            double dx = box.xmax - box.xmin;
            double dy = box.ymax - box.ymin;
            double dz = box.zmax - box.zmin;
            double diag = std::sqrt(dx * dx + dy * dy + dz * dz);

            double corner_tol = 0.01 * diag;

            // Loop through the tetrahedron's corner nodes
            for (int i = 0; i < 4; ++i) 
            {
                const auto& coord = tet.nd[i]->getCoord();
                double delta_x = probe_xyz[0] - coord.getx();
                double delta_y = probe_xyz[1] - coord.gety();
                double delta_z = probe_xyz[2] - coord.getz();
                double dist = std::sqrt(delta_x * delta_x + delta_y * delta_y + delta_z * delta_z);
                
                if (dist <= corner_tol) 
                {
                    nearest_corner_idx = i;
                    corner_tol = dist;  // Update tolerance to the closest corner
                }
            }

            if (nearest_corner_idx != -1) 
            {
                // Treat it as inside, directly assign barycentric coordinates (1 at corner, 0 elsewhere)
                std::array<double, 4> barycentric_coord = {0, 0, 0, 0};
                barycentric_coord[nearest_corner_idx] = 1.0;
            
                matches.emplace_back(tet_id, barycentric_coord);
                found_any = true;
            }            
            else
            {
                // Step 3: Check if the point is inside the tetrahedron using barycentric coordinates

                // More detailed check using barycentric coordinates
                double n0[3] = {tet.nd[0]->getCoord().getx(), tet.nd[0]->getCoord().gety(), tet.nd[0]->getCoord().getz()};
                double n1[3] = {tet.nd[1]->getCoord().getx(), tet.nd[1]->getCoord().gety(), tet.nd[1]->getCoord().getz()};
                double n2[3] = {tet.nd[2]->getCoord().getx(), tet.nd[2]->getCoord().gety(), tet.nd[2]->getCoord().getz()};
                double n3[3] = {tet.nd[3]->getCoord().getx(), tet.nd[3]->getCoord().gety(), tet.nd[3]->getCoord().getz()};

                std::array<double, 4> barycentric_coord = {0, 0, 0, 0};
                vtkTetra::BarycentricCoords(const_cast<double*>(probe_xyz), n0, n1, n2, n3, barycentric_coord.data());

                // Compute adaptive tolerance from box diagonal
                double adaptive_tol = tol * diag;  // scale base tol 
                if (diag < 1e-3)  // If the diagonal is very small, use a minimum tolerance
                {
                    adaptive_tol = diag*0.1;
                }


                double negative_sum = 0.0;
                for (double lambda : barycentric_coord) 
                {
                    if (lambda < 0.0)
                        negative_sum += -lambda;  // accumulate only the negative part
                }

                bool is_inside = (negative_sum < adaptive_tol);  // tolerance threshold

                if (is_inside) 
                {
                                    
                    // Clamp negative values and normalize the barycentric coordinates to sum to 1
                    double sum = 0.0;
                    for (double& lambda : barycentric_coord) 
                    {
                        if (lambda < 0.0)
                            lambda = 0.0;
                        sum += lambda;
                    }
                    if (sum > 0.0) 
                    {
                        for (double& lambda : barycentric_coord)
                            lambda /= sum;  // Normalize so sum == 1
                    }                    

                    matches.emplace_back(tet_id, barycentric_coord);
                    found_any = true;
                }
            }
        }
    }       
    // If not a leaf, check children
    else 
    {
        for (int i = 0; i < 8; ++i) 
        {
            if (node->children[i]) 
            {
                if (findTetraInOctree_function(node->children[i], probe_xyz, matches, tol))
                    found_any = true;
            }
        }
    }

    return found_any;
}


// --------------------------------------------------------------------------------------------------------------------

// Function to validate all tetrahedra in the leaves of the octree
void Octree::validateAllTetsInLeaves(OctreeNode* node, std::vector<bool>& tet_found_flags) 
{
    if (!node) return;

    if (node->is_leaf()) {
        
        //std::cout << "[Leaf Node] Tetra Count = " << node->tetra_ids.size() << std::endl;

        for (int id : node->tetra_ids) {
            if (id >= 0 && id < (int)tet_found_flags.size()) {
                tet_found_flags[id] = true;  // Mark as found at least once
            } 
            //else {
            //    std::cerr << "[Warning] Invalid tetrahedron ID " << id << " in leaf node" << std::endl;
            //}
        }
    } 
    else 
    {
        //std::cout << "[Internal Node] Tetra Count = " << node->tetra_ids.size() << std::endl;

        for (int i = 0; i < 8; ++i) {
            if (node->children[i]) {
                validateAllTetsInLeaves(node->children[i], tet_found_flags);
            }
        }
    }
}






void Octree::deleteOctree(OctreeNode* node) {
    if (!node) return;
    for (int i = 0; i < 8; ++i)
        deleteOctree(node->children[i]);
    delete node;
}


// ------------------------------------------------------------------------------------------------------------



void Octree::remove_duplicate_tetra_id_matches(std::vector<int>& matches) 
{
    // Use unordered_set to remove duplicates (order not preserved)
    std::unordered_set<int> unique_set(matches.begin(), matches.end());
    matches.assign(unique_set.begin(), unique_set.end());
}



bool Octree::findNCTetraInOctree(const AABB& face_box,std::vector<int>& matches)
{
    bool found = findNCTetraInOctree_function(root_octree_node, face_box, matches);

    // Remove duplicates
    remove_duplicate_tetra_id_matches(matches);

    return found;

}


bool Octree::findNCTetraInOctree_function(
    const OctreeNode* node,
    const AABB& face_box,
    std::vector<int>& matches)
{
    if (!node)
        return false;

    // Step 1: Check if this node's AABB overlaps the face's AABB
    bool overlap_x = (face_box.xmin <= node->box.xmax) && (face_box.xmax >= node->box.xmin);
    bool overlap_y = (face_box.ymin <= node->box.ymax) && (face_box.ymax >= node->box.ymin);
    bool overlap_z = (face_box.zmin <= node->box.zmax) && (face_box.zmax >= node->box.zmin);

    if (!(overlap_x && overlap_y && overlap_z))
        return false;

    bool found = false;

    if (node->is_leaf()) 
    {
        // Step 2: Test each tetra AABB for overlap
        for (int tet_id : node->NC_tetra_ids) 
        {
            const AABB& tet_box = tetra_boxes[tet_id];

            bool tet_overlap_x = (face_box.xmin <= tet_box.xmax) && (face_box.xmax >= tet_box.xmin);
            bool tet_overlap_y = (face_box.ymin <= tet_box.ymax) && (face_box.ymax >= tet_box.ymin);
            bool tet_overlap_z = (face_box.zmin <= tet_box.zmax) && (face_box.zmax >= tet_box.zmin);

            if (tet_overlap_x && tet_overlap_y && tet_overlap_z) 
            {
                matches.push_back(tet_id);
                found = true;
            }
        }
    }
    else 
    {
        // Step 3: Recurse into children
        for (int i = 0; i < 8; ++i) 
        {
            if (node->children[i]) 
            {
                if (findNCTetraInOctree_function(node->children[i], face_box, matches))
                    found = true;
            }
        }
    }

    return found;
}



/*
bool Octree::findNCTetraInOctree(double centroid[3], std::vector<int>& matches, double tol) 
{
    bool found = findNCTetraInOctree_function(root_octree_node, centroid, matches, tol);

    // Remove duplicates
    remove_duplicate_tetra_id_matches(matches);

    return found;
}


bool Octree::findNCTetraInOctree_function(
    const OctreeNode* node,
    double centroid[3],
    std::vector<int>& matches,
    double tol)
{
    // Step 1: Early exit if node is null or doesn't contain the centroid
    if (!node || !node->box.contains_with_buffer(centroid[0], centroid[1], centroid[2], 0.01))
        return false;

    bool found_any = false;

    if (node->is_leaf()) 
    {
        for (int tet_id : node->tetra_ids) 
        {
            const tetra& tet = tet_ptr[tet_id];
            const AABB& box = tetra_boxes[tet_id];
    
            if (!box.contains_with_buffer(centroid[0], centroid[1], centroid[2], 0.01))
                continue;
    
            // Step 1: Compute box diagonal for tolerance
            double dx = box.xmax - box.xmin;
            double dy = box.ymax - box.ymin;
            double dz = box.zmax - box.zmin;
            double diag = std::sqrt(dx * dx + dy * dy + dz * dz);
            double adaptive_tol = tol * diag;
    
            // Step 2: Check for proximity to any corner
            bool skip = false;
            for (int i = 0; i < 4; ++i) 
            {
                const auto& coord = tet.nd[i]->getCoord();
                double dist = std::sqrt(
                    std::pow(centroid[0] - coord.getx(), 2) +
                    std::pow(centroid[1] - coord.gety(), 2) +
                    std::pow(centroid[2] - coord.getz(), 2)
                );
    
                if (dist <= 0.01 * diag) 
                {
                    matches.push_back(tet_id);
                    found_any = true;
                    skip = true;
                    break;  // Stop checking this tet
                }
            }
    
            if (skip)
                continue;
    
            // Step 3: Compute barycentric coordinates
            double n0[3] = {tet.nd[0]->getCoord().getx(), tet.nd[0]->getCoord().gety(), tet.nd[0]->getCoord().getz()};
            double n1[3] = {tet.nd[1]->getCoord().getx(), tet.nd[1]->getCoord().gety(), tet.nd[1]->getCoord().getz()};
            double n2[3] = {tet.nd[2]->getCoord().getx(), tet.nd[2]->getCoord().gety(), tet.nd[2]->getCoord().getz()};
            double n3[3] = {tet.nd[3]->getCoord().getx(), tet.nd[3]->getCoord().gety(), tet.nd[3]->getCoord().getz()};
    
            std::array<double, 4> barycentric_coord;
            vtkTetra::BarycentricCoords(centroid, n0, n1, n2, n3, barycentric_coord.data());
    
            double negative_sum = 0.0;
            for (double lambda : barycentric_coord)
                if (lambda < 0.0)
                    negative_sum += -lambda;
    
            if (negative_sum < adaptive_tol) 
            {
                matches.push_back(tet_id);
                found_any = true;
            }
        }
    }
    

    else 
    {
        // Step 3: Recurse into child nodes
        for (int i = 0; i < 8; ++i) 
        {
            if (node->children[i]) 
            {
                if (findNCTetraInOctree_function(node->children[i], centroid, matches, tol))
                    found_any = true;
            }
        }
    }

    return found_any;
}
*/