//#####################################################################
// Copyright 2020, Jason Goulding, Craig Schroeder, Tamar Shinar.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Spatial_Acceleration/Nearest_Neighbors.h>
#include <Core/Arrays/ARRAY.h>
#include <Core/Data_Structures/HASHTABLE.h>
#include <Core/Math_Tools/RANGE_ITERATOR.h>
#include <Core/Vectors/VECTOR.h>
namespace PhysBAM{
//#####################################################################
// Function Nearest_Neighbors_Helper
//#####################################################################
template<class TV,class T,class F> void
Nearest_Neighbors_Helper(ARRAY_VIEW<const TV> X,F func,T dist)
{
    typedef VECTOR<int,TV::m> IV;
    HASHTABLE<IV,ARRAY<int> > hash;
    for(int i=0;i<X.m;i++)
        hash.Get_Or_Insert(IV(floor(X(i)/dist))).Append(i);

    ARRAY<IV> offsets;
    for(RANGE_ITERATOR<TV::m> it(RANGE<IV>(IV()-1,IV()+2));it.Valid();it.Next())
        offsets.Append(it.index);
    offsets.Resize(offsets.m/2);
    T dist2=sqr(dist);
    
    for(const auto& h:hash)
    {
        for(int i=0;i<h.second.m;i++)
            for(int j=i+1;j<h.second.m;j++)
            {
                T d2=(X(h.second(i))-X(h.second(j))).Magnitude_Squared();
                if(d2<=dist2) func(h.second(i),h.second(j),sqrt(d2));
            }
        for(auto o:offsets)
            if(auto*p=hash.Get_Pointer(h.first+o))
                for(auto i:h.second)
                    for(auto j:*p)
                    {
                        T d2=(X(i)-X(j)).Magnitude_Squared();
                        if(d2<=dist2) func(i,j,sqrt(d2));
                    }
    }
}
//#####################################################################
// Function Nearest_Neighbors
//#####################################################################
template<class TV,class T> void
Nearest_Neighbors(ARRAY_VIEW<const TV> X,ARRAY<pair<VECTOR<int,2>,T> >& pairs,T dist)
{
    pairs.Remove_All();
    Nearest_Neighbors_Helper(X,
        [&pairs](int i,int j,T d){pairs.Append({VECTOR<int,2>(i,j).Sorted(),d});},
        dist);
}
//#####################################################################
// Function Nearest_Neighbors
//#####################################################################
template<class TV,class T> void
Nearest_Neighbors(ARRAY_VIEW<const TV> X,std::function<void(int,int,T)> func,T dist)
{
    Nearest_Neighbors_Helper(X,func,dist);
}
//#####################################################################
// Function Nearest_Neighbors_Helper
//#####################################################################
template<class TV,class T,class F> void
Nearest_Neighbors_Helper(ARRAY_VIEW<const TV> X,ARRAY_VIEW<const TV> Y,F func,T dist)
{
    typedef VECTOR<int,TV::m> IV;
    HASHTABLE<IV,ARRAY<int> > hash;
    for(int i=0;i<X.m;i++)
        hash.Get_Or_Insert(IV(floor(X(i)/dist))).Append(i);

    ARRAY<IV> offsets;
    for(RANGE_ITERATOR<TV::m> it(RANGE<IV>(IV()-1,IV()+2));it.Valid();it.Next())
        offsets.Append(it.index);
    T dist2=sqr(dist);

    for(int j=0;j<Y.m;j++)
    {
        TV y=Y(j);
        IV index(floor(y/dist));
        for(const auto& o:offsets)
            if(auto*p=hash.Get_Pointer(index+o))
                for(auto i:*p)
                {
                    T d2=(y-X(i)).Magnitude_Squared();
                    if(d2<=dist2) func(i,j,sqrt(d2));
                }
    }
}
//#####################################################################
// Function Nearest_Neighbors
//#####################################################################
template<class TV,class T> void
Nearest_Neighbors(ARRAY_VIEW<const TV> X,ARRAY_VIEW<const TV> Y,ARRAY<pair<VECTOR<int,2>,T> >& pairs,T dist)
{
    pairs.Remove_All();
    Nearest_Neighbors_Helper(X,Y,
        [&pairs](int i,int j,T d){pairs.Append({VECTOR<int,2>(i,j).Sorted(),d});},
        dist);
}
//#####################################################################
// Function Nearest_Neighbors
//#####################################################################
template<class TV,class T> void
Nearest_Neighbors(ARRAY_VIEW<const TV> X,ARRAY_VIEW<const TV> Y,std::function<void(int,int,T)> func,T dist)
{
    Nearest_Neighbors_Helper(X,Y,func,dist);
}
//#####################################################################
// Function Nearest_Neighbors
//#####################################################################
template<class TV,class T> void
Nearest_Neighbors(ARRAY_VIEW<const TV> X,ARRAY<pair<VECTOR<int,2>,T> >& pairs,T dist,const RANGE<TV>& domain,const VECTOR<bool,TV::m>& periodic)
{
    pairs.Remove_All();
    ARRAY<TV> Y(X);
    ARRAY<int> mapping;
    for(int a=0;a<TV::m;a++)
    {
        if(!periodic(a)) continue;
        PHYSBAM_ASSERT(dist<domain.Edge_Lengths()(a));
        for(int i=0,n=Y.m;i<n;i++)
        {
            if(Y(i)(a)<=domain.min_corner(a)+dist)
            {
                Y.Append(Y(i));
                if(i<X.m) mapping.Append(i);
                else mapping.Append(mapping(i-X.m));
                
                Y.Last()(a)+=domain.Edge_Lengths()(a);
            }
            if(Y(i)(a)>=domain.max_corner(a)-dist)
            {
                Y.Append(Y(i));
                if(i<X.m) mapping.Append(i);
                else mapping.Append(mapping(i-X.m));
                
                Y.Last()(a)-=domain.Edge_Lengths()(a);
            }            
        }
    }
    Nearest_Neighbors<TV,T>(Y,pairs,dist);
    int k=0;
    for(int i=0;i<pairs.m;i++)
    {
        if(pairs(i).first.x>=X.m) continue;
        if(pairs(i).first.y>=X.m) pairs(i).first.y=mapping(pairs(i).first.y-X.m);
        if(pairs(i).first.x>=pairs(i).first.y) continue;
        pairs(k++)=pairs(i);
    }
    pairs.Resize(k);
}
template<class TV,class T> void
Nearest_Neighbors(ARRAY_VIEW<const TV> X,std::function<void(int,int,T)> func,T dist,const RANGE<TV>& domain,const VECTOR<bool,TV::m>& periodic);
template void Nearest_Neighbors<VECTOR<double,2>,double>(ARRAY_VIEW<VECTOR<double,2> const,int>,ARRAY<pair<VECTOR<int,2>,double>,int>&,double);
template void Nearest_Neighbors<VECTOR<double,2>,double>(ARRAY_VIEW<VECTOR<double,2> const,int>,ARRAY<pair<VECTOR<int,2>,double>,int>&,double,RANGE<VECTOR<double,2> > const&,VECTOR<bool,VECTOR<double,2>::m> const&);
template void Nearest_Neighbors<VECTOR<double,3>,double>(ARRAY_VIEW<VECTOR<double,3> const,int>,ARRAY<pair<VECTOR<int,2>,double>,int>&,double);
template void Nearest_Neighbors<VECTOR<double,3>,double>(ARRAY_VIEW<VECTOR<double,3> const,int>,ARRAY<pair<VECTOR<int,2>,double>,int>&,double,RANGE<VECTOR<double,3> > const&,VECTOR<bool,VECTOR<double,3>::m> const&);
template void Nearest_Neighbors<VECTOR<float,2>,float>(ARRAY_VIEW<VECTOR<float,2> const,int>,ARRAY<pair<VECTOR<int,2>,float>,int>&,float);
template void Nearest_Neighbors<VECTOR<float,2>,float>(ARRAY_VIEW<VECTOR<float,2> const,int>,ARRAY<pair<VECTOR<int,2>,float>,int>&,float,RANGE<VECTOR<float,2> > const&,VECTOR<bool,VECTOR<float,2>::m> const&);
template void Nearest_Neighbors<VECTOR<float,3>,float>(ARRAY_VIEW<VECTOR<float,3> const,int>,ARRAY<pair<VECTOR<int,2>,float>,int>&,float);
template void Nearest_Neighbors<VECTOR<float,3>,float>(ARRAY_VIEW<VECTOR<float,3> const,int>,ARRAY<pair<VECTOR<int,2>,float>,int>&,float,RANGE<VECTOR<float,3> > const&,VECTOR<bool,VECTOR<float,3>::m> const&);
template void Nearest_Neighbors<VECTOR<double,2>,double>(ARRAY_VIEW<VECTOR<double,2> const,int>,ARRAY_VIEW<VECTOR<double,2> const,int>,ARRAY<std::pair<VECTOR<int,2>,double>,int>&,double);
template void Nearest_Neighbors<VECTOR<double,3>,double>(ARRAY_VIEW<VECTOR<double,3> const,int>,ARRAY_VIEW<VECTOR<double,3> const,int>,ARRAY<std::pair<VECTOR<int,2>,double>,int>&,double);
template void Nearest_Neighbors<VECTOR<float,2>,float>(ARRAY_VIEW<VECTOR<float,2> const,int>,ARRAY_VIEW<VECTOR<float,2> const,int>,ARRAY<std::pair<VECTOR<int,2>,float>,int>&,float);
template void Nearest_Neighbors<VECTOR<float,3>,float>(ARRAY_VIEW<VECTOR<float,3> const,int>,ARRAY_VIEW<VECTOR<float,3> const,int>,ARRAY<std::pair<VECTOR<int,2>,float>,int>&,float);
template void Nearest_Neighbors<VECTOR<double,3>,double>(ARRAY_VIEW<VECTOR<double,3> const,int>,std::function<void (int,int,double)>,double);
template void Nearest_Neighbors<VECTOR<float,3>,float>(ARRAY_VIEW<VECTOR<float,3> const,int>,std::function<void (int,int,float)>,float);
template void Nearest_Neighbors<VECTOR<double,2>,double>(ARRAY_VIEW<VECTOR<double,2> const,int>,std::function<void (int,int,double)>,double);
template void Nearest_Neighbors<VECTOR<float,2>,float>(ARRAY_VIEW<VECTOR<float,2> const,int>,std::function<void (int,int,float)>,float);
}
