//#####################################################################
// Copyright 2015, Craig Schroeder, Andre Pradhana, and Gergely Klar.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#ifndef __GATHER_SCATTER__
#define __GATHER_SCATTER__
#include <Core/Geometry_Particles/DEBUG_PARTICLES.h>
#include <Core/Grids/FACE_INDEX.h>
#include <Core/Grids/GRID.h>
#include <Core/Log/SCOPE.h>
#include <Core/Utilities/TIMER.h>
#include <Core/Vectors/VECTOR.h>
#include <Hybrid_Methods/Examples_And_Drivers/MPM_PARTICLES.h>
#include <Hybrid_Methods/Iterators/PARTICLE_GRID_FACE_ITERATOR.h>
#include <Hybrid_Methods/Iterators/PARTICLE_GRID_ITERATOR.h>
#include <utility>
#ifdef USE_OPENMP
#include <omp.h>
#endif
namespace PhysBAM{

template<class TV> class MPM_EXAMPLE;
template<class TV> class MPM_OBJECTIVE;
template<class TV> class PARTICLE_GRID_WEIGHTS;

// TODO: probably need to leave an extra cell width to avoid collisions for MAC transfers.

template<class TV>
class GATHER_SCATTER
{
    typedef typename TV::SCALAR T;
    typedef VECTOR<int,TV::m> IV;
public:
    typedef typename PARTICLE_GRID_ITERATOR<TV>::SCRATCH T_SCRATCH;
    typedef typename PARTICLE_GRID_FACE_ITERATOR<TV>::SCRATCH T_FACE_SCRATCH;
    const ARRAY<int>& simulated_particles;
    const PARTICLE_GRID_WEIGHTS<TV>* weights;
    VECTOR<PARTICLE_GRID_WEIGHTS<TV>*,TV::m> face_weights;
    const GRID<TV>& grid;
    ARRAY<ARRAY<int> > bins;
    int threads;
    int partitions;

    GATHER_SCATTER(const GRID<TV>& grid,const ARRAY<int>& simulated_particles);
    ~GATHER_SCATTER();

    void Prepare_Scatter(const MPM_PARTICLES<TV>& particles);
    int Compute_Optimal_Bins(ARRAY<int>& bin_ends,ARRAY<int>& counts,
        ARRAY<int>& sum_counts,int min_width,int max_bins);
    
    // Safe to write to grid data
    template<class DATA,class REDUCE,class F> void
    Scatter(bool want_gradient,bool want_pseudo_gradient,
        const DATA& data,const REDUCE& reduce,F f)
    {
        if(threads>=2){
            Scatter_Parallel(want_gradient,want_pseudo_gradient,data,reduce,f);}
        else Serial(want_gradient,want_pseudo_gradient,data,reduce,f);
    }

    // Safe to write to grid data
    template<class DATA,class REDUCE,class F> void
    Gather(bool want_gradient,bool want_pseudo_gradient,
        const DATA& data,const REDUCE& reduce,F f)
    {
        if(threads>=2){
            Gather_Parallel(want_gradient,want_pseudo_gradient,data,reduce,f);}
        else Serial(want_gradient,want_pseudo_gradient,data,reduce,f);
    }

private:
    template<class DATA,class REDUCE,class F> void
    Serial(bool want_gradient,bool want_pseudo_gradient,
        const DATA& in_data,const REDUCE& reduce,F f)
    {
        T_FACE_SCRATCH face_scratch;
        DATA data(in_data);
        for(int k=0;k<simulated_particles.m;k++){
            int p=simulated_particles(k);
            Helper(face_scratch,want_gradient,want_pseudo_gradient,p,data,f);}
        if constexpr(!is_same<REDUCE,int>::value) reduce(data);
    }

    template<class DATA,class REDUCE,class F> void
    Scatter_Parallel(bool want_gradient,bool want_pseudo_gradient,
        const DATA& in_data,const REDUCE& reduce,F f)
    {
#pragma omp parallel
        {
            DATA data(in_data);
            for(int pass=0;pass<partitions;++pass)
            {
#pragma omp for
                for(int tid=0;tid<threads;tid++)
                {
                    int bin_id=tid*partitions+pass;
                    T_FACE_SCRATCH face_scratch;
                    for(int k=0;k<bins(bin_id).m;k++){
                        int p=bins(bin_id)(k);
                        Helper(face_scratch,want_gradient,want_pseudo_gradient,p,data,f);}
                }
            }
            if constexpr(!is_same<REDUCE,int>::value)
            {
#pragma omp critical
                reduce(data);
            }
        }
    }

    template<class DATA,class REDUCE,class F> void
    Gather_Parallel(bool want_gradient,bool want_pseudo_gradient,
        const DATA& in_data,const REDUCE& reduce,F f)
    {
#pragma omp parallel for
        for(int tid=0;tid<threads;tid++)
        {
            int a=tid*simulated_particles.m/threads;
            int b=(tid+1)*simulated_particles.m/threads;
            DATA data(in_data);
            T_FACE_SCRATCH face_scratch;
            for(int k=a;k<b;k++){
                int p=simulated_particles(k);
                Helper(face_scratch,want_gradient,want_pseudo_gradient,p,data,f);}
            if constexpr(!is_same<REDUCE,int>::value)
            {
#pragma omp critical
                reduce(data);
            }
        }
    }

    template<class DATA,class F> auto
    Helper(T_FACE_SCRATCH& face_scratch,bool want_gradient,
        bool want_pseudo_gradient,int p,DATA& data,F f)
        -> enable_if_t<sizeof((f(0,std::declval<PARTICLE_GRID_ITERATOR<TV>&>(),data),true))!=0>

    {
        PARTICLE_GRID_ITERATOR<TV> it(weights,p,want_gradient,want_pseudo_gradient,face_scratch(0));
        f(p,it,data);
    }

    template<class DATA,class F> auto
    Helper(T_FACE_SCRATCH& face_scratch,bool want_gradient,
        bool want_pseudo_gradient,int p,DATA& data,F f)
        -> enable_if_t<sizeof((f(0,std::declval<PARTICLE_GRID_FACE_ITERATOR<TV>&>(),data),true))!=0>

    {
        PARTICLE_GRID_FACE_ITERATOR<TV> it(face_weights,p,want_gradient,
            want_pseudo_gradient,face_scratch);
        f(p,it,data);
    }
};
}
#endif
