//#####################################################################
// Copyright 2004-2007, Zhaosheng Bao, Ron Fedkiw, Geoffrey Irving, Sergey Koltakov, Frank Losasso, Andrew Selle, Jerry Talton.
// This file is part of PhysBAM whose distribution is governed by the license  contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Data_Structures/Compute_Order_Using_KD_Tree.h>
#include <algorithm>

namespace PhysBAM
{
static void Stable_Partition_Helper(ARRAY_VIEW<int> a,ARRAY_VIEW<int> pred,
    ARRAY_VIEW<int> tmp)
{
    assert(tmp.m>=a.m);
    int k=0,n=0;
    for(int i=0;i<a.m;i++)
    {
        if(pred(a(i))) a(k++)=a(i);
        else tmp(n++)=a(i);
    }
    for(int i=0;i<n;i++) a(i+k)=tmp(i);
}

template<int d> void Compute_Helper(VECTOR<ARRAY<int>,d>& w,ARRAY<int>& scratch,
    ARRAY<int>& part_scratch,int s,int e,int axis,int p2,int r)
{
    int n=e-s;
    if(n<=1)
    {
        r-=scratch.m-1;
        assert((unsigned)r<scratch.m);
        scratch(w(axis)(s))=r;
        return;
    }
    int m=(n>=p2*3?s+2*p2:e-p2);
    
    for(int i=s;i<m;i++) scratch(w(axis)(i))=true;
    for(int i=m;i<e;i++) scratch(w(axis)(i))=false;

    for(int k=0;k<d;k++)
        if(k!=axis)
            Stable_Partition_Helper(w(k).Array_View(s,e-s),scratch,part_scratch);
    Compute_Helper(w,scratch,part_scratch,s,m,(axis+1)%d,(n>=p2*3)?p2:p2/2,2*r+1);
    Compute_Helper(w,scratch,part_scratch,m,e,(axis+1)%d,p2/2,2*r+2);
}
//#####################################################################
// Function Compute
//#####################################################################
template<class TV> void
Compute_Order_Using_KD_Tree(ARRAY<int>& leaf_order,const ARRAY<TV>& X)
{
    if(!X.m) return;
    VECTOR<ARRAY<int>,TV::m> w;
    for(int k=0;k<TV::m;k++) Set_Counting(w(k),X.m);
    ARRAY<int> scratch(X.m),part_scratch(X.m);
    leaf_order.Resize(X.m);

    for(int k=0;k<TV::m;k++)
        std::sort(w[k].begin(),w[k].end(),
            [k,&X](int a,int b){return X(a)(k)<X(b)(k);});

    int p2=1;
    while(p2*4<=X.m) p2*=2;
    Compute_Helper(w,scratch,part_scratch,0,X.m,0,p2,0);
    assert(w(0)==w(1));
    if constexpr(TV::m==3) assert(w(0)==w(2));
    for(int i=0;i<X.m;i++)
        leaf_order(scratch(i))=i;
}
template void Compute_Order_Using_KD_Tree<VECTOR<double,1> >(
    ARRAY<int,int>&,ARRAY<VECTOR<double,1>,int> const&);
template void Compute_Order_Using_KD_Tree<VECTOR<double,2> >(
    ARRAY<int,int>&,ARRAY<VECTOR<double,2>,int> const&);
template void Compute_Order_Using_KD_Tree<VECTOR<double,3> >(
    ARRAY<int,int>&,ARRAY<VECTOR<double,3>,int> const&);
template void Compute_Order_Using_KD_Tree<VECTOR<float,1> >(
    ARRAY<int,int>&,ARRAY<VECTOR<float,1>,int> const&);
template void Compute_Order_Using_KD_Tree<VECTOR<float,2> >(
    ARRAY<int,int>&,ARRAY<VECTOR<float,2>,int> const&);
template void Compute_Order_Using_KD_Tree<VECTOR<float,3> >(
    ARRAY<int,int>&,ARRAY<VECTOR<float,3>,int> const&);
}
