/* Copyright 2022 Remi Lehe
 *
 * This file is part of WarpX.
 *
 * License: BSD-3-Clause-LBNL
 */
#ifndef NamedComponentParticleContainer_H_
#define NamedComponentParticleContainer_H_

#include "Utils/TextMsg.H"

#include <AMReX.H>
#include <AMReX_AmrParGDB.H>
#include <AMReX_Particles.H>

#include <map>
#include <string>

/** Particle Attributes stored in amrex::ParticleContainer's struct of array
 */
struct PIdx
{
    enum {
        w = 0,      ///< weight
        ux, uy, uz,
#ifdef WARPX_DIM_RZ
        theta,      ///< RZ needs all three position components
#endif
        nattribs    ///< number of attributes
    };
};

/** Particle Container class that allows to add/access particle components
 *  with a name (string) instead of doing so with an integer index.
 *  (The "components" are all the particle quantities - except those
 *  that are stored in an AoS by amrex, i.e. the particle positions and ID)
 *
 *  This is done by storing maps that give the index of the component
 *  that corresponds to a given string.
 *
 * @tparam T_Allocator Mainly controls in which type of memory (e.g. device
 * arena, pinned memory arena, etc.) the particle data will be stored
 */
template <template<class> class T_Allocator=amrex::DefaultAllocator>
class NamedComponentParticleContainer :
public amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>
{
public:
    /** Construct an empty NamedComponentParticleContainer **/
    NamedComponentParticleContainer () : amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>() {}

    /** Construct a NamedComponentParticleContainer from an AmrParGDB object
     *
     * In this case, the only components are the default ones:
     * weight, momentum and (in RZ geometry) theta.
     *
     * @param amr_pgdb A pointer to a ParGDBBase, which contains pointers to
     * the Geometry, DistributionMapping, and BoxArray objects that define the
     * AMR hierarchy. Usually, this is generated by an AmrCore or AmrLevel object.
     */
    NamedComponentParticleContainer (amrex::AmrParGDB* amr_pgdb)
    : amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>(amr_pgdb) {
        // build up the map of string names to particle component numbers
        particle_comps["w"]  = PIdx::w;
        particle_comps["ux"] = PIdx::ux;
        particle_comps["uy"] = PIdx::uy;
        particle_comps["uz"] = PIdx::uz;
#ifdef WARPX_DIM_RZ
        particle_comps["theta"] = PIdx::theta;
#endif
    }

    /** Destructor for NamedComponentParticleContainer */
    virtual ~NamedComponentParticleContainer() = default;

    /** Construct a NamedComponentParticleContainer from a regular
     * amrex::ParticleContainer, and additional name-to-index maps
     *
     * @param pc regular particle container, where components are not named (only indexed)
     * @param p_comps name-to-index map for compile-time and run-time real components
     * @param p_icomps name-to-index map for compile-time and run-time integer components
     * @param p_rcomps name-to-index map for run-time real components
     * @param p_ricomps name-to-index map for run-time integer components
     */
    NamedComponentParticleContainer(
        amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator> && pc,
        std::map<std::string, int> p_comps,
        std::map<std::string, int> p_icomps,
        std::map<std::string, int> p_rcomps,
        std::map<std::string, int> p_ricomps)
    : amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>(std::move(pc)),
    particle_comps(p_comps),
    particle_icomps(p_icomps),
    particle_runtime_comps(p_rcomps),
    particle_runtime_icomps(p_ricomps) {}

    /** Move constructor for NamedComponentParticleContainer */
    NamedComponentParticleContainer ( NamedComponentParticleContainer && ) = default;
    /** Move operator for NamedComponentParticleContainer */
    NamedComponentParticleContainer& operator= ( NamedComponentParticleContainer && ) = default;

    /** Create an empty particle container
     *
     * This creates a new NamedComponentParticleContainer with same compile-time
     * and run-time attributes. But it can change its allocator.
     *
     * This function overloads the corresponding function from the parent
     * class (amrex::ParticleContainer)
     */
    template <template<class> class NewAllocator=amrex::DefaultAllocator>
    NamedComponentParticleContainer<NewAllocator>
    make_alike () const {
        auto tmp = NamedComponentParticleContainer<NewAllocator>(
            amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::template make_alike<NewAllocator>(),
            particle_comps,
            particle_icomps,
            particle_runtime_comps,
            particle_runtime_icomps);

        return tmp;
    }

    using amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::NumRealComps;
    using amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::NumIntComps;
    using amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::AddRealComp;
    using amrex::ParticleContainer<0,0,PIdx::nattribs,0,T_Allocator>::AddIntComp;

    /** Allocate a new run-time real component
     *
     * @param name Name of the new component
     * @param comm Whether to communicate this component, in the particle Redistribute
     */
    void AddRealComp (const std::string& name, bool comm=true)
    {
        auto search = particle_comps.find(name);
        if (search == particle_comps.end()) {
            particle_comps[name] = NumRealComps();
            particle_runtime_comps[name] = NumRealComps() - PIdx::nattribs;
            AddRealComp(comm);
        } else {
            amrex::Print() << Utils::TextMsg::Info(
                name + " already exists in particle_comps, not adding.");
        }
    }

    /** Allocate a new run-time integer component
     *
     * @param name Name of the new component
     * @param comm Whether to communicate this component, in the particle Redistribute
     */
    void AddIntComp (const std::string& name, bool comm=true)
    {
        auto search = particle_icomps.find(name);
        if (search == particle_icomps.end()) {
            particle_icomps[name] = NumIntComps();
            particle_runtime_icomps[name] = NumIntComps() - 0;
            AddIntComp(comm);
        } else {
            amrex::Print() << Utils::TextMsg::Info(
                name + " already exists in particle_icomps, not adding.");
        }
    }

    /** Return the name-to-index map for the compile-time and runtime-time real components */
    std::map<std::string, int> getParticleComps () const noexcept { return particle_comps;}
    /** Return the name-to-index map for the compile-time and runtime-time integer components */
    std::map<std::string, int> getParticleiComps () const noexcept { return particle_icomps;}
    /** Return the name-to-index map for the runtime-time real components */
    std::map<std::string, int> getParticleRuntimeComps () const noexcept { return particle_runtime_comps;}
    /** Return the name-to-index map for the runtime-time integer components */
    std::map<std::string, int> getParticleRuntimeiComps () const noexcept { return particle_runtime_icomps;}

protected:
    std::map<std::string, int> particle_comps;
    std::map<std::string, int> particle_icomps;
    std::map<std::string, int> particle_runtime_comps;
    std::map<std::string, int> particle_runtime_icomps;
};

#endif