/* ============================================================
 *
 * This file is a part of digiKam
 *
 * Date        : 02-02-2012
 * Description : Face database interface for spatial storage of face embedding.
 *
 * SPDX-FileCopyrightText: 2012-2013 by Marcel Wiesweg <marcel dot wiesweg at gmx dot de>
 * SPDX-FileCopyrightText: 2010-2024 by Gilles Caulier <caulier dot gilles at gmail dot com>
 * SPDX-FileCopyrightText: 2019      by Thanh Trung Dinh <dinhthanhtrung1996 at gmail dot com>
 * SPDX-FileCopyrightText: 2020      by Nghia Duong <minhnghiaduong997 at gmail dot com>
 * SPDX-FileCopyrightText: 2024      by Michael Miller <michael underscore miller at msn dot com>
 *
 * SPDX-License-Identifier: GPL-2.0-or-later
 *
 * ============================================================ */

#include "facedb_p.h"

// Local includes

#include "kd_nodebase.h"

namespace Digikam
{

class FaceDb::DataNode
{
public:

    DataNode() = default;

    bool isNull() const
    {
        return (nodeID == 0);
    }

public:

    int     nodeID      = 0;
    int     label       = 0;
    int     splitAxis   = 0;
    int     left        = -1;
    int     right       = -1;
    cv::Mat position;
    cv::Mat minRange;
    cv::Mat maxRange;
};

bool FaceDb::insertToTreeDb(const int nodeID, const cv::Mat& faceEmbedding) const
{
    bool isLeftChild    = false;
    int parentSplitAxis = 0;
    int parentID        = findParentTreeDb(faceEmbedding, isLeftChild, parentSplitAxis);

    if (parentID < 0)
    {
        qCWarning(DIGIKAM_FACEDB_LOG) << "fail to find parent node";
        return false;
    }

    QVariantList bindingValues;

    bindingValues << (parentSplitAxis + 1) % 128;
    bindingValues << nodeID;
    bindingValues << QByteArray::fromRawData((char*)faceEmbedding.ptr<float>(), (sizeof(float) * 128));
    bindingValues << QByteArray::fromRawData((char*)faceEmbedding.ptr<float>(), (sizeof(float) * 128));
    bindingValues << parentID;

    // Insert node to database.

    DbEngineSqlQuery query = d->db->execQuery(QLatin1String("INSERT INTO KDTree "
                                                            "(split_axis, position, max_range, min_range, parent, `left`, `right`) "
                                                            "VALUES (?, ?, ?, ?, ?, NULL, NULL)"),
                                              bindingValues);

    int newNode            = query.lastInsertId().toInt();

    if (newNode <= 0)
    {
        qCWarning(DIGIKAM_FACEDB_LOG) << "error insert into treedb" << query.lastError();
    }

    if (parentID > 0)
    {
        bindingValues.clear();
        bindingValues << newNode;
        bindingValues << parentID;

        // Not root -> update parent.

        if (isLeftChild)
        {
            query = d->db->execQuery(QLatin1String("UPDATE KDTree SET left = ? WHERE id = ?;"), bindingValues);
        }
        else
        {
            query = d->db->execQuery(QLatin1String("UPDATE KDTree SET right = ? WHERE id = ?;"), bindingValues);
        }
    }

    return true;
}

QMap<double, QVector<int> > FaceDb::getClosestNeighborsTreeDb(const cv::Mat& position,
                                                              float sqRange,
                                                              float cosThreshold,
                                                              int maxNbNeighbors) const
{
    QMap<double, QVector<int> > closestNeighbors;

    DataNode root;

    DbEngineSqlQuery query = d->db->execQuery(QLatin1String("SELECT position, max_range, min_range, `left`, `right` "
                                                            "FROM KDTree WHERE id = 1"));
    if (query.next())
    {
        // Encapsulate data node.

        root.nodeID     = 1;
        int embeddingID = query.value(0).toInt();
        root.maxRange   = cv::Mat(1, 128, CV_32F, query.value(1).toByteArray().data()).clone();
        root.minRange   = cv::Mat(1, 128, CV_32F, query.value(2).toByteArray().data()).clone();
        root.left       = query.value(3).toInt();
        root.right      = query.value(4).toInt();

        QVariantList bindingValues;
        bindingValues << embeddingID;

        query = d->db->execQuery(QLatin1String("SELECT identity, embedding FROM FaceMatrices WHERE id = ?"),
                                 bindingValues);

        if (query.next())
        {
            root.label    = query.value(0).toInt();
            root.position = cv::Mat(1, 128, CV_32F, query.value(1).toByteArray().data()).clone();
        }

        getClosestNeighborsTreeDb(root, closestNeighbors, position, sqRange, cosThreshold, maxNbNeighbors);
    }

    return closestNeighbors;
}

void FaceDb::clearTreeDb() const
{
    d->db->execSql(QLatin1String("DELETE FROM KDTree;"));
}

void FaceDb::updateRangeTreeDb(int nodeId, cv::Mat& minRange, cv::Mat& maxRange, const cv::Mat& position) const
{
    float* const min = minRange.ptr<float>();
    float* const max = maxRange.ptr<float>();
    const float* pos = position.ptr<float>();

    for (int i = 0 ; i < position.cols ; ++i)
    {
        max[i] = std::max(max[i], pos[i]);
        min[i] = std::min(min[i], pos[i]);
    }

    QVariantList bindingValues;

    bindingValues << QByteArray::fromRawData((char*)max, (sizeof(float) * 128));
    bindingValues << QByteArray::fromRawData((char*)min, (sizeof(float) * 128));
    bindingValues << nodeId;

    DbEngineSqlQuery query = d->db->execQuery(QLatin1String("UPDATE KDTree SET max_range = ?, min_range = ? WHERE id = ?;"),
                                              bindingValues);
}

int FaceDb::findParentTreeDb(const cv::Mat& nodePos, bool& leftChild, int& parentSplitAxis) const
{
    int parent           = 1;
    QVariant currentNode = parent;

    while (currentNode.isValid() && !currentNode.isNull())
    {
        parent = currentNode.toInt();

        QVariantList bindingValues;
        bindingValues << parent;

        DbEngineSqlQuery query = d->db->execQuery(QLatin1String("SELECT split_axis, position, max_range, min_range, `left`, `right` "
                                                                "FROM KDTree WHERE id = ?"),
                                                  bindingValues);

        if (! query.next())
        {
            if (parent == 1)
            {
/*
                qCDebug(DIGIKAM_FACEDB_LOG) << "add root";
*/
                return 0;
            }

            qCWarning(DIGIKAM_FACEDB_LOG) << "Error query parent =" << parent << query.lastError();

            return -1;
        }

/*
        qCDebug(DIGIKAM_FACEDB_LOG) << "split axis" << query.value(0).toInt()
                                    << "left"       << query.value(4)
                                    << "right"      << query.value(5);
*/

        int split        = query.value(0).toInt();
        parentSplitAxis  = split;

        int embeddingId  = query.value(1).toInt();
        cv::Mat maxRange = cv::Mat(1, 128, CV_32F, query.value(2).toByteArray().data()).clone();
        cv::Mat minRange = cv::Mat(1, 128, CV_32F, query.value(3).toByteArray().data()).clone();
        QVariant left    = query.value(4);
        QVariant right   = query.value(5);

        bindingValues.clear();
        bindingValues << embeddingId;

        query            = d->db->execQuery(QLatin1String("SELECT embedding FROM FaceMatrices WHERE id = ?"),
                                            bindingValues);

        if (! query.next())
        {
            qCWarning(DIGIKAM_FACEDB_LOG) << "fail to find parent face embedding" << query.lastError();

            return -1;
        }

        cv::Mat position = cv::Mat(1, 128, CV_32F, query.value(0).toByteArray().data()).clone();

        if (nodePos.at<float>(0, split) >= position.at<float>(0, split))
        {
            currentNode = right;
            leftChild   = false;
        }
        else
        {
            currentNode = left;
            leftChild   = true;
        }

         updateRangeTreeDb(parent, minRange, maxRange, nodePos);
    }

    return parent;
}

double FaceDb::getClosestNeighborsTreeDb(const DataNode& subTree,
                                         QMap<double, QVector<int> >& neighborList,
                                         const cv::Mat& position,
                                         float sqRange,
                                         float cosThreshold,
                                         int maxNbNeighbors) const
{
    if (subTree.isNull())
    {
        return sqRange;
    }

    // Try to add current node to the list.

    const float sqrdistanceToCurrentNode = KDNodeBase::sqrDistance(position.ptr<float>(), subTree.position.ptr<float>(), 128);

    if (
        (sqrdistanceToCurrentNode < sqRange) &&
        (KDNodeBase::cosDistance(position.ptr<float>(), subTree.position.ptr<float>(), 128) > cosThreshold)
       )
    {
        neighborList[sqrdistanceToCurrentNode].append(subTree.label);

        // Limit the size of the Map to maxNbNeighbors.

        int size = 0;

        for (QMap<double, QVector<int> >::const_iterator iter  = neighborList.cbegin();
                                                         iter != neighborList.cend();
                                                         ++iter)
        {
            size += iter.value().size();
        }

        if (size > maxNbNeighbors)
        {
            // Eliminate the farthest neighbor.

            QMap<double, QVector<int> >::iterator farthestNodes = std::prev(neighborList.end(), 1);

            if (farthestNodes.value().size() == 1)
            {
                neighborList.erase(farthestNodes);
            }
            else
            {
                farthestNodes.value().pop_back();
            }

            // Update the searching range.

            sqRange = neighborList.lastKey();
        }
    }

    // Sub-trees Traversal.

    double sqrDistanceLeftTree  = 0.0;
    double sqrDistanceRightTree = 0.0;
    DataNode leftNode;
    DataNode rightNode;

    if (subTree.left <= 0)
    {
        sqrDistanceLeftTree = DBL_MAX;
    }
    else
    {
        QVariantList bindingValues;
        bindingValues << subTree.left;

        DbEngineSqlQuery query = d->db->execQuery(QLatin1String("SELECT position, max_range, min_range, `left`, `right` "
                                                                "FROM KDTree WHERE id = ?"),
                                                  bindingValues);

        if (query.next())
        {
            // Encapsulate data node.

            leftNode.nodeID   = subTree.left;
            int embeddingID   = query.value(0).toInt();
            leftNode.maxRange = cv::Mat(1, 128, CV_32F, query.value(1).toByteArray().data()).clone();
            leftNode.minRange = cv::Mat(1, 128, CV_32F, query.value(2).toByteArray().data()).clone();
            leftNode.left     = query.value(3).toInt();
            leftNode.right    = query.value(4).toInt();

            bindingValues.clear();
            bindingValues << embeddingID;

            query             = d->db->execQuery(QLatin1String("SELECT identity, embedding FROM FaceMatrices WHERE id = ?"),
                                                 bindingValues);

            if (query.next())
            {
                leftNode.label        = query.value(0).toInt();
                leftNode.position     = cv::Mat(1, 128, CV_32F, query.value(1).toByteArray().data()).clone();

                const float* minRange = leftNode.minRange.ptr<float>();
                const float* maxRange = leftNode.maxRange.ptr<float>();
                const float* pos      = leftNode.position.ptr<float>();

                for (int i = 0 ; i < 128 ; ++i)
                {
                    sqrDistanceLeftTree += (pow(qMax((minRange[i] - pos[i]),      0.0f), 2) +
                                            pow(qMax((pos[i]      - maxRange[i]), 0.0f), 2));
                }
            }
        }
    }

    if (subTree.right <= 0)
    {
        sqrDistanceRightTree = DBL_MAX;
    }
    else
    {
        QVariantList bindingValues;
        bindingValues << subTree.right;

        DbEngineSqlQuery query = d->db->execQuery(QLatin1String("SELECT position, max_range, min_range, `left`, `right` "
                                                                "FROM KDTree WHERE id = ?"),
                                                  bindingValues);

        if (query.next())
        {
            // Encapsulate data node.

            rightNode.nodeID   = subTree.right;
            int embeddingID    = query.value(0).toInt();
            rightNode.maxRange = cv::Mat(1, 128, CV_32F, query.value(1).toByteArray().data()).clone();
            rightNode.minRange = cv::Mat(1, 128, CV_32F, query.value(2).toByteArray().data()).clone();
            rightNode.left     = query.value(3).toInt();
            rightNode.right    = query.value(4).toInt();

            bindingValues.clear();
            bindingValues << embeddingID;

            query              = d->db->execQuery(QLatin1String("SELECT identity, embedding FROM FaceMatrices WHERE id = ?"),
                                                  bindingValues);

            if (query.next())
            {
                rightNode.label             = query.value(0).toInt();
                rightNode.position          = cv::Mat(1, 128, CV_32F, query.value(1).toByteArray().data()).clone();

                const float* const minRange = rightNode.minRange.ptr<float>();
                const float* const maxRange = rightNode.maxRange.ptr<float>();
                const float* const pos      = rightNode.position.ptr<float>();

                for (int i = 0 ; i < 128 ; ++i)
                {
                    sqrDistanceRightTree += (pow(qMax((minRange[i] - pos[i]),      0.0f), 2) +
                                             pow(qMax((pos[i]      - maxRange[i]), 0.0f), 2));
                }
            }
        }
    }

    // Traverse the closest area.

    if (sqrDistanceLeftTree < sqrDistanceRightTree)
    {
        if (sqrDistanceLeftTree < sqRange)
        {
            // Traverse left Tree.

            sqRange = getClosestNeighborsTreeDb(leftNode, neighborList, position, sqRange, cosThreshold, maxNbNeighbors);

            if (sqrDistanceRightTree < sqRange)
            {
                // Traverse right Tree.

                sqRange = getClosestNeighborsTreeDb(rightNode, neighborList, position, sqRange, cosThreshold, maxNbNeighbors);
            }
        }
    }
    else
    {
        if (sqrDistanceRightTree < sqRange)
        {
            // Traverse right Tree.

            sqRange = getClosestNeighborsTreeDb(rightNode, neighborList, position, sqRange, cosThreshold, maxNbNeighbors);

            if (sqrDistanceLeftTree < sqRange)
            {
                // Traverse left Tree.

                sqRange = getClosestNeighborsTreeDb(leftNode, neighborList, position, sqRange, cosThreshold, maxNbNeighbors);
            }
        }
    }

    return sqRange;
}

} // namespace Digikam
