/****************************************************************************
* MeshLab                                                           o o     *
* A versatile mesh processing toolbox                             o     o   *
*                                                                _   O  _   *
* Copyright(C) 2005                                                \/)\/    *
* Visual Computing Lab                                            /\/|      *
* ISTI - Italian National Research Council                           |      *
*                                                                    \      *
* All rights reserved.                                                      *
*                                                                           *
* This program is free software; you can redistribute it and/or modify      *
* it under the terms of the GNU General Public License as published by      *
* the Free Software Foundation; either version 2 of the License, or         *
* (at your option) any later version.                                       *
*                                                                           *
* This program is distributed in the hope that it will be useful,           *
* but WITHOUT ANY WARRANTY; without even the implied warranty of            *
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the             *
* GNU General Public License (http://www.gnu.org/licenses/gpl.txt)          *
* for more details.                                                         *
*                                                                           *
****************************************************************************/
#include <vcg/space/distance3.h>
#include "balltree.h"

namespace GaelMls {

template<typename _Scalar>
BallTree<_Scalar>::BallTree(const vcg::ConstDataWrapper<VectorType>& points, const vcg::ConstDataWrapper<Scalar>& radii)
    : mPoints(points), mRadii(radii), mRadiusScale(1.), mTreeIsUptodate(false)
{
    mRootNode = 0;
    mMaxTreeDepth = 12;
    mTargetCellSize = 24;
}

template<typename _Scalar>
void BallTree<_Scalar>::computeNeighbors(const VectorType& x, Neighborhood<Scalar>* pNei) const
{
    if (!mTreeIsUptodate)
        const_cast<BallTree*>(this)->rebuild();

    pNei->clear();
    mQueryPosition = x;
    queryNode(*mRootNode, pNei);
}

template<typename _Scalar>
void BallTree<_Scalar>::queryNode(Node& node, Neighborhood<Scalar>* pNei) const
{
    if (node.leaf)
    {
        for (unsigned int i=0 ; i<node.size ; ++i)
        {
            int id = node.indices[i];
            Scalar d2 = vcg::SquaredNorm(mQueryPosition - mPoints[id]);
            Scalar r = mRadiusScale * mRadii[id];
            if (d2<r*r)
                pNei->insert(id, d2);
        }
    }
    else
    {
        if (mQueryPosition[node.dim] - node.splitValue < 0)
            queryNode(*node.children[0], pNei);
        else
            queryNode(*node.children[1], pNei);
    }
}

template <typename Scalar>
inline vcg::Point3<Scalar> CwiseAdd(vcg::Point3<Scalar> const & p1, Scalar s)
{
    return vcg::Point3<Scalar>(p1.X() + s, p1.Y() + s, p1.Z() + s);
}

template<typename _Scalar>
void BallTree<_Scalar>::rebuild(void)
{
        delete mRootNode;

        mRootNode = new Node();
        IndexArray indices(mPoints.size());
        AxisAlignedBoxType aabb;
        aabb.Set(mPoints[0]);
        for (unsigned int i=0 ; i<mPoints.size() ; ++i)
        {
                indices[i] = i;
                aabb.Add(mPoints[i],mRadii[i]*mRadiusScale);
//				aabb.min = Min(aabb.min, CwiseAdd(mPoints[i], -mRadii[i]*mRadiusScale));
//				aabb.max = Max(aabb.max, CwiseAdd(mPoints[i],  mRadii[i]*mRadiusScale));
        }
        buildNode(*mRootNode, indices, aabb, 0);

        mTreeIsUptodate = true;
}

template<typename _Scalar>
void BallTree<_Scalar>::split(const IndexArray& indices, const AxisAlignedBoxType& aabbLeft, const AxisAlignedBoxType& aabbRight, IndexArray& iLeft, IndexArray& iRight)
{
    for (std::vector<int>::const_iterator it=indices.begin(), end=indices.end() ; it!=end ; ++it)
    {
        unsigned int i = *it;
        if (vcg::PointFilledBoxDistance(mPoints[i], aabbLeft) < mRadii[i]*mRadiusScale)
            iLeft.push_back(i);

        if (vcg::PointFilledBoxDistance(mPoints[i], aabbRight) < mRadii[i]*mRadiusScale)
            iRight.push_back(i);
    }
}

template<typename _Scalar>
void BallTree<_Scalar>::buildNode(Node& node, std::vector<int>& indices, AxisAlignedBoxType aabb, int level)
{
    Scalar avgradius = 0.;
    for (std::vector<int>::const_iterator it=indices.begin(), end=indices.end() ; it!=end ; ++it)
            avgradius += mRadii[*it];
    avgradius = mRadiusScale * avgradius / Scalar(indices.size());
    VectorType diag = aabb.max - aabb.min;
    if  (int(indices.size())<mTargetCellSize
        || avgradius*0.9 > std::max(std::max(diag.X(), diag.Y()), diag.Z())
        || int(level)>=mMaxTreeDepth)
    {
        node.leaf = true;
        node.size = indices.size();
        node.indices = new unsigned int[node.size];
        for (unsigned int i=0 ; i<node.size ; ++i)
            node.indices[i] = indices[i];
        return;
    }

    unsigned int dim = diag.MaxCoeffId();
    node.dim = dim;
    node.splitValue = Scalar(0.5*(aabb.max[dim] + aabb.min[dim]));
    node.leaf = 0;

    AxisAlignedBoxType aabbLeft=aabb, aabbRight=aabb;
    aabbLeft.max[dim] = node.splitValue;
    aabbRight.min[dim] = node.splitValue;

    std::vector<int> iLeft, iRight;
    split(indices, aabbLeft, aabbRight, iLeft,iRight);

    // we don't need the index list anymore
    indices.clear();

    {
        // left child
        //mNodes.resize(mNodes.size()+1);
        Node* pChild = new Node();
        node.children[0] = pChild;
        buildNode(*pChild, iLeft, aabbLeft, level+1);
    }

    {
        // right child
        //mNodes.resize(mNodes.size()+1);
        Node* pChild = new Node();
        node.children[1] = pChild;
        buildNode(*pChild, iRight, aabbRight, level+1);
    }
}

template class BallTree<float>;
template class BallTree<double>;

}
