//#####################################################################
// Copyright 2015, Andre Pradhana, Chenfanfu Jiang.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
// Class POISSON_DISK
//#####################################################################
#include <Core/Seeding/POISSON_DISK.h>
#include <Core/Grids/GRID.h>
#include <Core/Implicit_Objects/IMPLICIT_OBJECT.h>
#include <Core/Math_Tools/pow.h>
#include <Core/Parallel_Computation/Parallel_Partition.h>
#include <Core/Random_Numbers/RANDOM_NUMBERS.h>
#include <Core/Viewer/VIEWER_OUTPUT.h>
#ifdef USE_OPENMP
#include <omp.h>
#endif
namespace PhysBAM{
//#####################################################################
// Constructor
//#####################################################################
template<class TV> POISSON_DISK<TV>::
POISSON_DISK(const typename TV::SCALAR min_distance,const int max_attempts)
    :min_distance(min_distance),h(min_distance/sqrt((T)(TV::m))),max_attempts(max_attempts)
{
}
//#####################################################################
// Destructor
//#####################################################################
template<class TV> POISSON_DISK<TV>::
~POISSON_DISK()
{
}
//#####################################################################
// Function Sample
//#####################################################################
template<class TV> void POISSON_DISK<TV>::
Sample(RANDOM_NUMBERS<T>& random,IMPLICIT_OBJECT<TV>& object,ARRAY<TV>& X,T contour,bool threaded)
{
    int num_seeds=X.m;
    object.Update_Box();
    if(threaded) Sample_Raw_Threaded(random,object.box,X);
    else Sample_Raw(random,object.box,X);
    for(int i=X.m-1;i>=num_seeds;i--)
        if(object.Extended_Phi(X(i))>contour)
            X.Remove_Index_Lazy(i);
    Make_Samples_Periodic(X);
}
//#####################################################################
// Function Insert_In_Array
//#####################################################################
template<class TV> void POISSON_DISK<TV>::
Insert_In_Array(ARRAY<int,IV>& a,const GRID<TV>& grid,const TV& X,int index)
{
    if(!is_periodic.Contains(true))
    {
        a(grid.Cell_I(X))=index;
        return;
    }
    TV Y[pow<TV::m>(3)]={X},e=periodic_boundary.Edge_Lengths();
    int num=1;
    for(int i=0;i<TV::m;i++)
    {
        if(!need_periodic(i)) continue;
        for(int p=0,n=num;p<n;p++)
        {
            if(Y[p](i)<wrap_zone.min_corner(i))
                Y[num++]=Y[p].Add_Axis(i,e(i));
            if(Y[p](i)>wrap_zone.max_corner(i))
                Y[num++]=Y[p].Add_Axis(i,-e(i));
        }
    }
    for(int i=0;i<num;i++)
        a(grid.Cell_I(Y[i]))=index;
}
//#####################################################################
// Function Make_Samples_Periodic
//#####################################################################
template<class TV> void POISSON_DISK<TV>::
Make_Samples_Periodic(ARRAY<TV>& X)
{
    TV periodic_size=periodic_boundary.Edge_Lengths();
    for(auto&x:X)
        for(int i=0;i<TV::m;i++)
            if(is_periodic(i))
            {
                if(x(i)<periodic_boundary.min_corner(i)) x(i)+=periodic_size(i);
                else if(x(i)>periodic_boundary.max_corner(i)) x(i)-=periodic_size(i);
            }
}
//#####################################################################
// Function Sample
//#####################################################################
template<class TV> void POISSON_DISK<TV>::
Sample(RANDOM_NUMBERS<T>& random,const RANGE<TV>& box,ARRAY<TV>& X,bool threaded)
{
    if(threaded) Sample_Raw_Threaded(random,box,X);
    else Sample_Raw(random,box,X);
    Make_Samples_Periodic(X);
}
//#####################################################################
// Function Sample
//#####################################################################
template<class TV> void POISSON_DISK<TV>::
Sample_Raw(RANDOM_NUMBERS<T>& random,const RANGE<TV>& box,ARRAY<TV>& X)
{
    GRID<TV> grid=GRID<TV>::Create_Grid_Given_Cell_Size(box,h);
    ARRAY<int,IV> grid_array(grid.Cell_Indices(2),init_new,-1);
    RANGE<TV> safe_domain=grid.domain.Thickened(min_distance);
    TV safe_size=safe_domain.Edge_Lengths();
    TV periodic_size=periodic_boundary.Edge_Lengths();
    need_periodic.Fill(false);
    for(int i=0;i<TV::m;i++)
        if(is_periodic(i) && safe_size(i)>=periodic_size(i))
        {
            need_periodic(i)=true;
            wrap_zone.min_corner(i)=safe_domain.max_corner(i)-periodic_size(i);
            wrap_zone.max_corner(i)=safe_domain.min_corner(i)+periodic_size(i);
        }

    bool has_seed=false;
    auto ok_seed_box=box.Thickened(min_distance);
    for(int i=0;i<X.m;i++)
    {
        if(!grid_array.domain.Lazy_Inside_Half_Open(grid.Cell_I(X(i)))) continue;
        Insert_In_Array(grid_array,grid,X(i),i);
        if(ok_seed_box.Lazy_Inside(X(i)))
            has_seed=true;
    }
    if(!has_seed)
    {
        TV first_point;
        random.Fill_Uniform(first_point,box);
        Insert_In_Array(grid_array,grid,first_point,0);
        X.Append(first_point);
    }
    ARRAY<int> active(X.m);
    Set_Counting(active);
    while(active.m){
        int random_index=random.Get_Uniform_Integer(0,active.m-1);
        TV point=X(active(random_index));
        bool found_at_least_one=false;
        for(int i=0;i<max_attempts;i++){
            TV new_point=Generate_Random_Point_Around_Annulus(random,point);
            if(box.Lazy_Inside(new_point) && Check_Distance(grid,grid_array,new_point,X)){
                found_at_least_one=true;
                int index=X.Append(new_point);
                active.Append(index);
                Insert_In_Array(grid_array,grid,new_point,index);
                break;}}
        if(!found_at_least_one) active.Remove_Index_Lazy(random_index);}
}
//#####################################################################
// Function Sample
//#####################################################################
template<class TV> void POISSON_DISK<TV>::
Sample_Raw_Threaded(RANDOM_NUMBERS<T>& random,const RANGE<TV>& box,ARRAY<TV>& X)
{
    // TODO: fix periodic

#ifdef USE_OPENMP
    int threads=1;
#pragma omp parallel
#pragma omp master
    threads=omp_get_num_threads();

    T width=box.Edge_Lengths().x;
    int num_threads=std::min((int)(width/(10*min_distance)),threads);
    int num_strips=num_threads*2;
    LOG::printf("Perform Poisson seeding on %P of %P threads.\n",num_threads,threads);
    if(num_strips<=2) Sample_Raw(random,box,X);
    T strip_width=width/num_strips;

    ARRAY<ARRAY<TV>> Xt(num_strips),Yt(num_strips);
    Parallel_Partition(Xt,X,
        [z=box.min_corner.x,strip_width](const TV& X)
        {
            int strip_id=(X.x-z)/strip_width;
            return strip_id/2;
        });

    ARRAY<int> seeds(2*num_strips);
    for(auto&i:seeds) seeds(i)=random.Get_Uniform_Integer(0,INT_MAX);

    for(int p=0;p<2;p++)
    {
#pragma omp parallel for
        for(int t=0;t<threads;t++)
        {
            int s=2*t+p;
            if(s>=num_strips) continue;
            RANDOM_NUMBERS<T> rn;
            rn.Set_Seed(seeds(s));
            POISSON_DISK<TV> pd(min_distance,max_attempts);
            ARRAY<TV> Y(Xt(s));
            Y.Append(Yt(s));
            if(s<num_strips-1)
            {
                Y.Append(Xt(s+1));
                Y.Append(Yt(s+1));
            }
            else if(s==num_strips-1 && is_periodic(0))
            {
                Y.Append(Xt(0));
                Y.Append(Yt(0));
            }
            if(s>0)
            {
                Y.Append(Xt(s-1));
                Y.Append(Yt(s-1));
            }
            else if(s==0 && is_periodic(0))
            {
                Y.Append(Xt.Last());
                Y.Append(Yt.Last());
            }
            int num_seeds=Y.m;
            auto box2=box;
            box2.min_corner.x=box.min_corner.x+width*(s)/num_strips;
            box2.max_corner.x=box.min_corner.x+width*(s+1)/num_strips;
            pd.Sample_Raw(rn,box2,Y);
            Yt(s)=Y.Array_View(num_seeds,Y.m-num_seeds);
        }
        Flush_Frame();
    }
    ARRAY<int> start_offset(num_strips);
    int start=X.m;
    for(int s=0;s<num_strips;s++)
    {
        start_offset(s)=start;
        start+=Yt(s).m;
    }
    X.Resize(start,no_init);
#pragma omp parallel for
    for(int s=0;s<num_strips;s++)
        X.Array_View(start_offset(s),Yt(s).m)=Yt(s);
#else
    Sample_Raw(random,box,X);
#endif
}
//#####################################################################
// Function Generate_Random_Point_Around_Annulus
//#####################################################################
template<class TV> TV POISSON_DISK<TV>::
Generate_Random_Point_Around_Annulus(RANDOM_NUMBERS<T>& random,const TV& center) const
{
    while(1){
        TV v;
        random.Fill_Uniform(v,-1,1);
        T mag2=v.Magnitude_Squared();
        if(mag2>=0.25 && mag2<=1) return v*min_distance*2+center;}
    return TV();
}
//#####################################################################
// Function Check_Distance
//#####################################################################
template<class TV> bool POISSON_DISK<TV>::
Check_Distance(const GRID<TV>& grid,ARRAY<int,IV>& grid_array,
    const TV& point,ARRAY<TV>& X) const
{
    IV cell=grid.Cell_I(point);
    RANGE<IV> candidate_range(cell-2,cell+3);
    TV half_edge=periodic_boundary.Edge_Lengths()/2;
    for(RANGE_ITERATOR<TV::m> it(candidate_range);it.Valid();it.Next()){
        if(grid_array(it.index)==-1) continue;
        TV Y=point-X(grid_array(it.index));
        for(int a=0;a<TV::m;a++)
            if(is_periodic(a))
                Y(a)=wrap(Y(a),-half_edge(a),half_edge(a));
        if(Y.Magnitude_Squared()<sqr(min_distance)) return false;}
    return true;
}
//#####################################################################
// Function Set_Distance_By_Volume
//#####################################################################
template<class TV> void POISSON_DISK<TV>::
Set_Distance_By_Volume(T volume_per_sample)
{
    if(TV::m==1) min_distance=volume_per_sample*((T)2/3);
    else if(TV::m==2) min_distance=sqrt(volume_per_sample*((T)2/3));
    else min_distance=cbrt(volume_per_sample*((T)13/18));
    h=min_distance/sqrt((T)(TV::m));
}
//#####################################################################
// Function Set_Distance_By_Volume
//#####################################################################
template<class TV> void POISSON_DISK<TV>::
Set_Min_Distance(T min_dist)
{
    min_distance=min_dist;
    h=min_distance/sqrt((T)(TV::m));
}
template class POISSON_DISK<VECTOR<float,1> >;
template class POISSON_DISK<VECTOR<float,2> >;
template class POISSON_DISK<VECTOR<float,3> >;
template class POISSON_DISK<VECTOR<double,1> >;
template class POISSON_DISK<VECTOR<double,2> >;
template class POISSON_DISK<VECTOR<double,3> >;
}
