/** @file PenningTrap.hpp
 *
 *  @author Cory Alexander Balaton (coryab)
 *  @author Janita Ovidie Sandtrøen Willumsen (janitaws)
 *
 *  @version 0.1
 *
 *  @brief A class for simulating a Penning trap.
 *
 *  @bug No known bugs
 * */
#ifndef __PENNING_TRAP__
#define __PENNING_TRAP__

#include <armadillo>
#include <omp.h>

#include "Particle.hpp"
#include "constants.hpp"

typedef std::vector<arma::vec::fixed<3>> sim_cols;
typedef std::vector<arma::vec::fixed<3>> sim_rows;
typedef std::vector<sim_cols> sim_arr;

#pragma omp declare reduction(+ : arma::vec : omp_out += omp_in)               \
    initializer(omp_priv = omp_orig)

/** @brief A class that simulates a Penning trap.
 *
 *  This class simulates a Penning trap. It can take in a number of particles
 *  and simulate how they would behave inside a Penning trap.
 * */
class PenningTrap {
private:
    double B_0;                      ///< Magnetic field strength
    double V_0;                      ///< Applied potential
    double d;                        ///< Characteristic dimension
    std::vector<Particle> particles; ///< The particles in the Penning trap
    arma::vec::fixed<3> *k_v;
    arma::vec::fixed<3> *k_r;
    std::function<arma::vec(int, double)> v_funcs[4];
    std::function<arma::vec(int, double)> r_funcs[4];

public:
    /** @brief Set B_0, V_0 and d.
     * */
    PenningTrap(double B_0 = T, double V_0 = 25. * V / 1000., double d = 500.);

    PenningTrap(int i, double B_0 = T, double V_0 = 25. * V / 1000.,
                double d = 500.);

    PenningTrap(std::vector<Particle> particles, double B_0 = T,
                double V_0 = 25. * V / 1000., double d = 500.);

    /** @brief Add a particle to the system
     * */
    void add_particle(Particle particle);

    /** @brief Calculate E at point r
     * */
    arma::vec external_E_field(arma::vec r);

    /** @brief Calculate B at point r
     * */
    arma::vec external_B_field(arma::vec r);

    /** @brief Calculate the force between 2 particles.
     *
     *  @details Calculate the force exhibited on particle p_i from
     *  particle p_j.
     * */
    arma::vec force_on_particle(int i, int j);

    /** @brief Calculate the total external force on a particle.
     *
     *  @details Calculate the total amount of force that E and B exhibits
     *  on particle p_i.
     * */
    arma::vec total_force_external(int i);

    /** @brief Calculate the total force on a particle from other particles.
     * */
    arma::vec total_force_particles(int i);

    /** @brief calculate the total force on a particle.
     * */
    arma::vec total_force(int i);

    /** @brief Go forward one timestep using the RK4 method
     * */
    void evolve_RK4(double dt, bool particle_interaction = true);

    /** @brief Go forward one timestep using the forward Euler method
     * */
    void evolve_forward_euler(double dt, bool particle_interaction = true);

    sim_arr simulate(double time, int steps, std::string method = "rk4",
                     bool particle_interaction = true);

    void write_simulation_to_dir(std::string path, double time, int steps,
                                 std::string method = "rk4",
                                 bool particle_interaction = true);
};

#endif