/* Copyright 2019 Luca Fedeli
 *
 * This file is part of WarpX.
 *
 * License: BSD-3-Clause-LBNL
 */
#ifndef WARPX_quantum_sync_engine_wrapper_h_
#define WARPX_quantum_sync_engine_wrapper_h_

#include "QedWrapperCommons.H"
#include "QuantumSyncEngineInnards.H"

#include <AMReX_Array.H>
#include <AMReX_Vector.H>
#include <AMReX_Gpu.H>

//#define PXRMP_CORE_ONLY allows importing only the 'core functions' of the
//Quantum Synchrotron engine of the QED PICSAR library.
#define PXRMP_CORE_ONLY
#include <quantum_sync_engine.hpp>

//Lookup table building function is in a dedicated (optional) class to
//avoid including heavy dependencies if they are not needed.
#ifdef WARPX_QED_TABLE_GEN
#    include "QuantumSyncEngineTableBuilder.H"
#endif

#include <string>

//Some handy aliases

// The engine has two templated arguments: the numerical type
// and a random number generator. However, random numbers are not
// used to generate the lookup tables and the static member
// functions which are called in the functors do not use
// random numbers as well. Therefore, an empty "DummyStruct"
// can be passed as a template parameter.
using PicsarQuantumSynchrotronEngine = picsar::multi_physics::
    quantum_synchrotron_engine<amrex::Real, QedUtils::DummyStruct>;

using PicsarQuantumSynchrotronCtrl =
    picsar::multi_physics::quantum_synchrotron_engine_ctrl<amrex::Real>;
//__________

// Functors ==================================

// These functors allow using the core elementary functions of the library.
// They are generated by a factory class (QuantumSynchrotronEngine, see below).
// They can be included in GPU kernels.

/**
*  Functor to initialize the optical depth of leptons for the
*  Quantum Synchrotron process
*/
class QuantumSynchrotronGetOpticalDepth
{
public:
    /**
     * Constructor does nothing because optical depth initialization
     * does not require control parameters or lookup tables.
     */
    QuantumSynchrotronGetOpticalDepth ()
    {};

    /**
     * () operator is just a thin wrapper around a very simple function to
     * generate the optical depth. It can be used on GPU.
     */
    AMREX_GPU_HOST_DEVICE
    AMREX_FORCE_INLINE
    amrex::Real operator() () const noexcept
    {
        //A random number in [0,1) should be provided as an argument.
        return PicsarQuantumSynchrotronEngine::
            internal_get_optical_depth(amrex::Random());
    }
};
//____________________________________________

/**
 * Functor to evolve the optical depth of leptons due to the
 * Quantum Synchrotron process
 */
class QuantumSynchrotronEvolveOpticalDepth
{
public:

    QuantumSynchrotronEvolveOpticalDepth () = default;

    /**
     * Constructor acquires pointers to control parameters and
     * lookup tables data.
     * lookup_table uses non-owning vectors under the hood. So no new data
     * allocations should be triggered on GPU
     */
    QuantumSynchrotronEvolveOpticalDepth(
        QuantumSynchrotronEngineInnards& r_innards):
        m_ctrl{r_innards.ctrl},
        m_KKfunc_size{r_innards.KKfunc_coords.size()},
        m_p_KKfunc_coords{r_innards.KKfunc_coords.dataPtr()},
        m_p_KKfunc_data{r_innards.KKfunc_data.dataPtr()}
        {};

    /**
     * Evolves the optical depth. It can be used on GPU.
     * @param[in] px,py,pz momentum components of the lepton (SI units)
     * @param[in] ex,ey,ez electric field components (SI units)
     * @param[in] bx,by,bz magnetic field components (SI units)
     * @param[in] dt timestep (SI units)
     * @param[in,out] opt_depth optical depth of the lepton. It is modified by the method.
     * @return a flag which is 1 if optical depth becomes negative (i.e. a photon has to be generated).
     */
    AMREX_GPU_HOST_DEVICE
    AMREX_FORCE_INLINE
    int operator()(
        amrex::Real px, amrex::Real py, amrex::Real pz,
        amrex::Real ex, amrex::Real ey, amrex::Real ez,
        amrex::Real bx, amrex::Real by, amrex::Real bz,
        amrex::Real dt, amrex::Real& opt_depth) const noexcept
    {
        bool has_event_happened{false};

        //the library provides the time (< dt) at which the event occurs, but this
        //feature won't be used in WarpX for now.
        amrex::Real unused_event_time{0.0};

        PicsarQuantumSynchrotronEngine::
        internal_evolve_opt_depth_and_determine_event(
            px, py, pz,
            ex, ey, ez,
            bx, by, bz,
            dt, opt_depth,
            has_event_happened, unused_event_time,
            m_dummy_lambda,
            picsar::multi_physics::lookup_1d<amrex::Real>{
                m_KKfunc_size,
                m_p_KKfunc_coords,
                m_p_KKfunc_data},
            m_ctrl);

        return has_event_happened;
    }

private:
    //laser wavelength is not used with SI units
    amrex::Real m_dummy_lambda{1.0};

    PicsarQuantumSynchrotronCtrl m_ctrl;

    //lookup table data
    size_t m_KKfunc_size;
    amrex::Real* m_p_KKfunc_coords;
    amrex::Real* m_p_KKfunc_data;
};

/**
 * Functor to generate a photon via the Quantum Synchrotron process
 * and to update momentum accordingly
 */
class QuantumSynchrotronGeneratePhotonAndUpdateMomentum
{
public:
    /**
     * Constructor acquires pointers to control parameters and
     * lookup tables data.
     * lookup_table uses non-owning vectors under the hood. So no new data
     * allocations should be triggered on GPU
     */
    QuantumSynchrotronGeneratePhotonAndUpdateMomentum(
        QuantumSynchrotronEngineInnards& r_innards):
        m_ctrl{r_innards.ctrl},
        m_cum_distrib_coords_1_size{r_innards.cum_distrib_coords_1.size()},
        m_cum_distrib_coords_2_size{r_innards.cum_distrib_coords_2.size()},
        m_p_distrib_coords_1{r_innards.cum_distrib_coords_1.data()},
        m_p_distrib_coords_2{r_innards.cum_distrib_coords_2.data()},
        m_p_cum_distrib_data{r_innards.cum_distrib_data.data()}
        {};

    /**
     * Generates sampling (template parameter) photons according to Quantum Synchrotron process.
     * It can be used on GPU.
     * @param[in,out] px,py,pz momentum components of the lepton. They are modified (SI units)
     * @param[in] ex,ey,ez electric field components (SI units)
     * @param[in] bx,by,bz magnetic field components (SI units)
     * @param[in] weight of the lepton (code units)
     * @param[out] g_px,g_py,g_pz momenta of generated photons. Each array should have size=sampling  (SI units)
     * @param[out] g_weight weight of the generated photons. Array should have size=sampling (code units)
     */
    template <size_t sampling>
    AMREX_GPU_HOST_DEVICE
    AMREX_FORCE_INLINE
    void operator()(
    amrex::Real* px, amrex::Real* py, amrex::Real* pz,
    amrex::Real ex, amrex::Real ey, amrex::Real ez,
    amrex::Real bx, amrex::Real by, amrex::Real bz,
    amrex::Real weight,
    amrex::Real* g_px, amrex::Real* g_py, amrex::Real* g_pz,
    amrex::Real* g_weight) const noexcept
    {
        //[sampling] random numbers are needed
        amrex::GpuArray<amrex::Real, sampling>
            rand_zero_one_minus_epsi;
        for(auto& el : rand_zero_one_minus_epsi) el = amrex::Random();

        PicsarQuantumSynchrotronEngine::
        internal_generate_photons_and_update_momentum(
            *px, *py, *pz,
            ex, ey, ez,
            bx, by, bz,
            weight, sampling,
            g_px, g_py, g_pz,
            g_weight,
            m_dummy_lambda,
            picsar::multi_physics::lookup_2d<amrex::Real>{
                m_cum_distrib_coords_1_size,
                m_p_distrib_coords_1,
                m_cum_distrib_coords_2_size,
                m_p_distrib_coords_2,
                m_p_cum_distrib_data},
            m_ctrl,
            rand_zero_one_minus_epsi.data());
    }

private:
    //laser wavelenght is not used with SI units
    const amrex::Real m_dummy_lambda{1.0};

    const PicsarQuantumSynchrotronCtrl m_ctrl;

    //lookup table data
    size_t m_cum_distrib_coords_1_size;
    size_t m_cum_distrib_coords_2_size;
    amrex::Real* m_p_distrib_coords_1;
    amrex::Real* m_p_distrib_coords_2;
    amrex::Real* m_p_cum_distrib_data;
};

// Factory class =============================

/**
 * Wrapper for the Quantum Synchrotron engine of the PICSAR library
 */
class QuantumSynchrotronEngine
{
public:
    /**
     * Constructor requires no arguments.
     */
    QuantumSynchrotronEngine ();

    /**
     * Builds the functor to initialize the optical depth
     */
    QuantumSynchrotronGetOpticalDepth build_optical_depth_functor ();

    /**
     * Builds the functor to evolve the optical depth
     */
    QuantumSynchrotronEvolveOpticalDepth build_evolve_functor ();

    /**
     * Builds the functor to generate photons
     */
    QuantumSynchrotronGeneratePhotonAndUpdateMomentum build_phot_em_functor ();

    /**
     * Checks if the optical tables are properly initialized
     */
    bool are_lookup_tables_initialized () const;

    /**
     * Init lookup tables from raw binary data.
     * @param[in] raw_data a Vector of char
     * @return true if it succeeds, false if it cannot parse raw_data
     */
    bool init_lookup_tables_from_raw_data (const amrex::Vector<char>& raw_data);

    /**
     * Init lookup tables using built-in dummy tables
     * for test purposes.
     */
    void init_dummy_tables();

    /**
     * Export lookup tables data into a raw binary Vector
     * @return the data in binary format. The Vector is empty if tables were
     * not previously initialized.
     */
    amrex::Vector<char> export_lookup_tables_data () const;

    /**
     * Computes the lookup tables. It does nothing unless WarpX is compiled with QED_TABLE_GEN=TRUE
     * @param[in] ctrl control params to generate the tables
     */
    void compute_lookup_tables (PicsarQuantumSynchrotronCtrl ctrl);

    /**
     * gets default (reasonable) values for the control parameters
     * @return default control params to generate the tables
     */
    PicsarQuantumSynchrotronCtrl get_default_ctrl() const;

    /**
     * returns a constant reference to the control parameters
     * @return const reference to control parameters
     */
    const PicsarQuantumSynchrotronCtrl& get_ref_ctrl() const;

private:
    bool m_lookup_tables_initialized = false;

    QuantumSynchrotronEngineInnards m_innards;

//Table builing is available only if the libray is compiled with QED_TABLE_GEN=TRUE
#ifdef WARPX_QED_TABLE_GEN
    QuantumSynchrotronEngineTableBuilder m_table_builder;
#endif

};

//============================================

#endif //WARPX_quantum_sync_engine_wrapper_h_