#include "pch.h"
#include "HDBSCAN.h"
#include <queue>
#include <cmath>
#include <algorithm>
#include <limits>
#include <unordered_map>
#include <set>
#include <vector>

namespace
{
    float sqrDistance(const cv::Point2f& a, const cv::Point2f& b)
    {
        float dx = a.x - b.x;
        float dy = a.y - b.y;
        return dx * dx + dy * dy;
    }

    float distance(const cv::Point2f& a, const cv::Point2f& b)
    {
        return std::sqrt(sqrDistance(a, b));
    }

    // Simple KD-Tree node for 2D points (optimized for k-NN search)
    struct KDNode2D {
        cv::Point2f point;
        int index;
        KDNode2D* left;
        KDNode2D* right;
        int axis; // 0 = x, 1 = y

        KDNode2D(const cv::Point2f& p, int idx, int ax = 0)
            : point(p), index(idx), left(nullptr), right(nullptr), axis(ax) {
        }
    };

    // Build KD-Tree for 2D points (simplified, but effective for small-medium datasets)
    KDNode2D* buildKDTree(std::vector<cv::Point2f>& points, std::vector<int>& indices, int depth = 0)
    {
        if (indices.empty()) return nullptr;
        if (indices.size() == 1) {
            return new KDNode2D(points[indices[0]], indices[0], depth % 2);
        }

        int axis = depth % 2;
        int median = static_cast<int>(indices.size() / 2);

        // Partial sort to find median
        std::nth_element(indices.begin(), indices.begin() + median, indices.end(),
            [&](int a, int b) {
                return axis == 0 ? (points[a].x < points[b].x) : (points[a].y < points[b].y);
            });

        int rootIdx = indices[median];
        KDNode2D* node = new KDNode2D(points[rootIdx], rootIdx, axis);

        std::vector<int> leftIndices(indices.begin(), indices.begin() + median);
        std::vector<int> rightIndices(indices.begin() + median + 1, indices.end());

        node->left = buildKDTree(points, leftIndices, depth + 1);
        node->right = buildKDTree(points, rightIndices, depth + 1);

        return node;
    }

    // Find k nearest neighbors using KD-Tree (max-heap to keep k smallest)
    void findKNN(KDNode2D* node, const cv::Point2f& query, int k,
        std::priority_queue<std::pair<float, int>>& heap)
    {
        if (!node) return;

        float dist = sqrDistance(node->point, query);

        // Use max-heap to keep k smallest distances
        if (heap.size() < static_cast<size_t>(k)) {
            heap.push({ dist, node->index });
        }
        else if (dist < heap.top().first) {
            heap.pop();
            heap.push({ dist, node->index });
        }

        int axis = node->axis;
        float diff = (axis == 0) ? (query.x - node->point.x) : (query.y - node->point.y);
        float diffSq = diff * diff;

        KDNode2D* nearNode = nullptr;
        KDNode2D* farNode = nullptr;
        if (diff <= 0) {
            nearNode = node->left;
            farNode = node->right;
        }
        else {
            nearNode = node->right;
            farNode = node->left;
        }

        if (nearNode) findKNN(nearNode, query, k, heap);
        if (farNode && (heap.size() < static_cast<size_t>(k) || diffSq < heap.top().first)) {
            findKNN(farNode, query, k, heap);
        }
    }

    // Delete KD-Tree
    void deleteKDTree(KDNode2D* node)
    {
        if (!node) return;
        deleteKDTree(node->left);
        deleteKDTree(node->right);
        delete node;
    }

    // Union-Find data structure for MST clustering (optimized)
    class UnionFind {
    private:
        std::vector<int> parent;
        std::vector<int> rank;
        std::vector<int> size;

    public:
        UnionFind(int n) : parent(n), rank(n, 0), size(n, 1) {
            for (int i = 0; i < n; ++i) {
                parent[i] = i;
            }
        }

        int find(int x) {
            if (parent[x] != x) {
                parent[x] = find(parent[x]); // Path compression
            }
            return parent[x];
        }

        void unite(int x, int y) {
            int px = find(x);
            int py = find(y);
            if (px == py) return;

            // Union by rank
            if (rank[px] < rank[py]) {
                std::swap(px, py);
            }
            parent[py] = px;
            size[px] += size[py];
            if (rank[px] == rank[py]) {
                rank[px]++;
            }
        }

        int getSize(int x) {
            return size[find(x)];
        }
    };

    // Edge structure for MST
    struct Edge {
        int u, v;
        float weight;

        bool operator<(const Edge& other) const {
            return weight < other.weight;
        }
    };

    // Calculate mutual reachability distance (HDBSCAN* core concept)
    // This is the key difference from DBSCAN - uses max(distance, coreDist1, coreDist2)
    inline float mutualReachabilityDist(
        const cv::Point2f& p1, const cv::Point2f& p2,
        float coreDist1, float coreDist2)
    {
        float d = distance(p1, p2);
        return std::max({ d, coreDist1, coreDist2 });
    }

    inline float bottomY(const cv::Rect& b)
    {
        return static_cast<float>(b.y + b.height);
    }

    inline bool isBottomAligned(const cv::Rect& a, const cv::Rect& b, float bottomDiffRatio)
    {
        const float maxH = static_cast<float>(std::max(a.height, b.height));
        if (maxH <= 0.0f) return false;
        const float diff = std::abs(bottomY(a) - bottomY(b));
        return diff <= maxH * bottomDiffRatio;
    }

    inline bool isCenterClose(const cv::Rect& a, const cv::Rect& b,
        const cv::Point2f& ca, const cv::Point2f& cb)
    {
        const float maxW = static_cast<float>(std::max(a.width, b.width));
        if (maxW <= 0.0f) return false;
        return distance(ca, cb) <= (2.0f * maxW);
    }
}

// HDBSCAN* implementation - Full algorithm matching the original paper:
// 1. Calculate core distance (k-th nearest neighbor) using KD-Tree for speed
// 2. Build MST using mutual reachability distance (not Euclidean distance)
// 3. Extract clusters from MST using stability-based selection (hierarchical clustering)
// 4. Filter clusters by minPts and stability
std::vector<std::vector<int>> HDBSCAN::cluster(const std::vector<cv::Rect>& boxes,
    float eps,
    int minPts,
    float bottomDiffRatio)
{
    std::vector<std::vector<int>> clusters;
    const int n = static_cast<int>(boxes.size());
    if (n == 0 || minPts < 2) return clusters;
    if (n < minPts) return clusters;

    std::vector<cv::Point2f> points;
    points.reserve(n);
    for (const auto& b : boxes) {
        points.emplace_back(b.x + b.width * 0.5f, b.y + b.height * 0.5f);
    }

    // ===== 1. Calculate core distance using KD-Tree (optimized k-NN search) =====
    std::vector<float> coreDist(n, 0.0f);

    std::vector<cv::Point2f> pointsCopy = points;
    std::vector<int> indices(n);
    for (int i = 0; i < n; ++i) indices[i] = i;

    KDNode2D* kdTree = buildKDTree(pointsCopy, indices);

    for (int i = 0; i < n; ++i) {
        std::priority_queue<std::pair<float, int>> heap;

        findKNN(kdTree, points[i], minPts, heap);

        if (heap.size() >= static_cast<size_t>(minPts)) {
            coreDist[i] = std::sqrt(heap.top().first);
        }
        else if (heap.size() > 0) {
            float maxDist = 0.0f;
            while (!heap.empty()) {
                maxDist = std::max(maxDist, std::sqrt(heap.top().first));
                heap.pop();
            }
            coreDist[i] = maxDist;
        }
        else {
            coreDist[i] = 0.0f;
        }
    }

    deleteKDTree(kdTree);

    // ===== 2. Build MST using mutual reachability distance =====
    std::vector<Edge> edges;
    const size_t maxEdges = (n < 100) ? (n * (n - 1) / 2) : (n * 10);
    edges.reserve(std::min(maxEdges, static_cast<size_t>(n * (n - 1) / 2)));

    for (int i = 0; i < n; ++i) {
        for (int j = i + 1; j < n; ++j) {
            float mrDist = mutualReachabilityDist(
                points[i], points[j],
                coreDist[i], coreDist[j]);

            if (eps > 0 && mrDist > eps) continue;

            edges.push_back({ i, j, mrDist });

            if (edges.size() >= maxEdges && n > 100) {
                goto build_mst;
            }
        }
    }

build_mst:
    std::sort(edges.begin(), edges.end());

    UnionFind uf(n);
    std::vector<Edge> mstEdges;
    mstEdges.reserve(n - 1);

    for (const auto& e : edges) {
        if (uf.find(e.u) != uf.find(e.v)) {
            uf.unite(e.u, e.v);
            mstEdges.push_back(e);
            if (static_cast<int>(mstEdges.size()) >= n - 1) {
                break;
            }
        }
    }

    // ===== 3. Extract clusters from MST =====
    std::sort(mstEdges.begin(), mstEdges.end());

    struct ClusterInfo {
        int root;
        int size;
        float birthWeight;
        float deathWeight;
        float stability;
        bool isActive;
    };

    std::vector<ClusterInfo> clusterInfo(n);
    for (int i = 0; i < n; ++i) {
        clusterInfo[i] = { i, 1, 0.0f, std::numeric_limits<float>::max(), 0.0f, true };
    }

    UnionFind clusterUF(n);
    std::vector<int> clusterSizes(n, 1);

    for (const auto& e : mstEdges) {
        int rootU = clusterUF.find(e.u);
        int rootV = clusterUF.find(e.v);

        if (rootU != rootV) {
            int sizeU = clusterSizes[rootU];
            int sizeV = clusterSizes[rootV];

            if (sizeU >= minPts && clusterInfo[rootU].isActive) {
                float lambdaDeath = e.weight;
                float lambdaBirth = clusterInfo[rootU].birthWeight;
                clusterInfo[rootU].stability = (lambdaDeath - lambdaBirth) * sizeU;
                clusterInfo[rootU].deathWeight = lambdaDeath;
                clusterInfo[rootU].isActive = false;
            }

            if (sizeV >= minPts && clusterInfo[rootV].isActive) {
                float lambdaDeath = e.weight;
                float lambdaBirth = clusterInfo[rootV].birthWeight;
                clusterInfo[rootV].stability = (lambdaDeath - lambdaBirth) * sizeV;
                clusterInfo[rootV].deathWeight = lambdaDeath;
                clusterInfo[rootV].isActive = false;
            }

            clusterUF.unite(e.u, e.v);
            int newRoot = clusterUF.find(e.u);
            int newSize = sizeU + sizeV;
            clusterSizes[newRoot] = newSize;

            float minBirth = std::min(clusterInfo[rootU].birthWeight, clusterInfo[rootV].birthWeight);
            if (minBirth == 0.0f) minBirth = e.weight;

            clusterInfo[newRoot] = {
                newRoot, newSize, minBirth,
                std::numeric_limits<float>::max(), 0.0f, true
            };
        }
    }

    for (int i = 0; i < n; ++i) {
        int root = clusterUF.find(i);
        if (clusterInfo[root].isActive && clusterSizes[root] >= minPts) {
            float maxWeight = mstEdges.empty() ? 0.0f : mstEdges.back().weight;
            clusterInfo[root].stability = (maxWeight - clusterInfo[root].birthWeight) * clusterSizes[root];
        }
    }

    // ===== 4. Extract clusters with positive stability =====
    std::unordered_map<int, std::vector<int>> clusterMap;
    std::set<int> processedRoots;

    for (int i = 0; i < n; ++i) {
        int root = clusterUF.find(i);
        int clusterSize = clusterSizes[root];

        if (clusterSize >= minPts && clusterInfo[root].stability >= -0.001f) {
            if (processedRoots.find(root) == processedRoots.end()) {
                clusterMap[root] = std::vector<int>();
                processedRoots.insert(root);
            }
            clusterMap[root].push_back(i);
        }
    }

    // Convert to output format + filter by bottom alignment
    for (auto& pair : clusterMap) {
        std::vector<int> filtered;
        filtered.reserve(pair.second.size());

        for (int idx : pair.second) {
            bool hasNeighbor = false;
            for (int other : pair.second) {
                if (idx == other) continue;
                if (isBottomAligned(boxes[idx], boxes[other], bottomDiffRatio)) {
                    hasNeighbor = true;
                    break;
                }
            }
            if (hasNeighbor) {
                filtered.push_back(idx);
            }
        }

        // Extra rule: if cluster size is 2 or 3, enforce center distance <= max width
        std::vector<int> proximityFiltered;
        if (filtered.size() == 2) {
            int a = filtered[0];
            int b = filtered[1];
            if (isCenterClose(boxes[a], boxes[b], points[a], points[b])) {
                proximityFiltered = filtered;
            }
        }
        else if (filtered.size() == 3) {
            std::vector<bool> keep(3, true);
            for (int i = 0; i < 3; ++i) {
                for (int j = i + 1; j < 3; ++j) {
                    int a = filtered[i];
                    int b = filtered[j];
                    if (!isCenterClose(boxes[a], boxes[b], points[a], points[b])) {
                        keep[i] = false;
                        keep[j] = false;
                    }
                }
            }
            for (int i = 0; i < 3; ++i) {
                if (keep[i]) {
                    proximityFiltered.push_back(filtered[i]);
                }
            }
        }
        else {
            proximityFiltered = std::move(filtered);
        }

        if (static_cast<int>(proximityFiltered.size()) >= minPts) {
            clusters.push_back(std::move(proximityFiltered));
        }
    }
    return clusters;
}
