/** @file PenningTrap.hpp
 *
 *  @author Cory Alexander Balaton (coryab)
 *  @author Janita Ovidie Sandtrøen Willumsen (janitaws)
 *
 *  @version 1.0
 *
 *  @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"
#include "typedefs.hpp"
#include "utils.hpp"

#pragma omp declare reduction(+ : vec3 : 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
{
    /** @brief Make PenningTrapTest a friend of PenningTrap.
     * */
    friend class PenningTrapTest;
private:
    double B_0;                                 ///< Magnetic field strength
    double V_0;                                 ///< Applied potential
    std::function<double(double)> perturbation; ///< Time-dependent perturbation
    double d;                                   ///< Characteristic dimension
    double t;                                   ///< Current time
    std::vector<Particle> particles; ///< The particles in the Penning trap

    /** @brief A 2D vector containing all \f$k_{v,i,j}\f$ where \f$j\f$ is the
     *  index of a particle
     * */
    sim_arr k_v;

    /** @brief A 2D vector containing all \f$k_{r,i,j}\f$ where \f$j\f$ is the
     *  index of a particle
     * */
    sim_arr k_r;

    /** @brief Helper for evolve_RK4 when calculating \f$k_{v,i,j}\f$ values
     *
     *  @details Something
     *
     *  @param i Index i for \f$k_{v,i,j}\f$
     *  @param j Index j for \f$k_{v,i,j}\f$
     *  @param dt the step length (delta time)
     *
     *  @return vec3
     * */
    vec3 v_func(uint i, uint j, double dt);

    /** @brief Helper for evolve_RK4 when calculating \f$k_{r,i,j}\f$ values
     *
     *  @details Something
     *
     *  @param i Index i for \f$k_{r,i,j}\f$
     *  @param j Index j for \f$k_{r,i,j}\f$
     *  @param dt The step length (delta time)
     *
     *  @return vec3
     * */
    vec3 r_func(uint i, uint j, double dt);

    /** @brief Calculate E at point r
     *
     *  @param r The position where we want to calculate the E field
     *
     *  @return vec3
     * */
    vec3 external_E_field(vec3 r);

    /** @brief Calculate B at point r
     *
     *  @param r The position where we want to calculate the B field
     *
     *  @return vec3
     * */
    vec3 external_B_field(vec3 r);

    /** @brief Calculate the force between 2 particles.
     *
     *  @details Calculate the force exhibited on particle p_i from
     *  particle p_j.
     *
     *  @param i The index of particle p_i
     *  @param j The index of particle p_j
     *
     *  @return vec3
     * */
    vec3 force_on_particle(uint i, uint 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.
     *
     *  @param i The index of particle p_i
     *
     *  @return vec3
     * */
    vec3 total_force_external(uint i);

    /** @brief Calculate the total force on a particle p_i from other
     * particles.
     *
     *  @param i The index of particle p_i
     *
     *  @return vec3
     * */
    vec3 total_force_particles(uint i);

    /** @brief calculate the total force on a particle p_i.
     *
     *  @param i The index of particle p_i
     *
     *  @return vec3
     * */
    vec3 total_force(uint i);

    /** @brief calculate the total force on a particle p_i without interaction
     *
     *  @param i The index of particle p_i
     *
     *  @return vec3
     * */
    vec3 total_force_no_interaction(uint i);

public:
    /** @brief Constructor for the PenningTrap class
     *
     *  @param B_0 The magnetic field strength
     *  @param V_0 The time dependent applied potential
     *  @param d The characteristic dimension
     *  @param t The starting time
     * */
    PenningTrap(double B_0 = T, double V_0 = (25. * V) / 1000., double d = 500.,
                double t = 0.);

    /** @brief Constructor for the PenningTrap class
     *
     *  @param i The number of particles to generate
     *  @param B_0 The magnetic field strength
     *  @param V_0 The time dependent applied potential
     *  @param d The characteristic dimension
     *  @param t The starting time
     * */
    PenningTrap(uint i, double B_0 = T, double V_0 = (25. * V) / 1000.,
                double d = 500., double t = 0.);

    /** @brief Constructor for the PenningTrap class
     *
     *  @param particles The starting particles
     *  @param B_0 The magnetic field strength
     *  @param V_0 The time dependent applied potential
     *  @param d The characteristic dimension
     *  @param t The starting time
     * */
    PenningTrap(std::vector<Particle> particles, double B_0 = T,
                double V_0 = (25. * V) / 1000., double d = 500., double t = 0.);

    /** @brief Time dependent perturbation to V_0
     *
     *  @param f The amplitude of the perturbation
     *  @parma omega_V the angular frequency of the perturbation
     * */
    void set_pertubation(double f, double omega_V);

    /** @brief Give all particles new positions and velocities, and change t
     *  and V_0.
     *
     *  @param V_0 The tiome dependent applied potential
     *  @param t The starting time
     * */
    void reinitialize(double f, double omega_V, double t = 0.);

    /** @brief Add a particle to the system
     *
     *  @param particle The particle to add to the Penning trap
     * */
    void add_particle(Particle particle);

    /** @brief Go forward one timestep using the RK4 method
     *
     *  @param dt The step length
     *  @param particle_interaction Turn particle interactions on/off
     * */
    void evolve_RK4(double dt, bool particle_interaction = true);

    /** @brief Go forward one timestep using the forward Euler method
     *
     *  @param dt The step length
     *  @param particle_interaction Turn particle interactions on/off
     * */
    void evolve_forward_euler(double dt, bool particle_interaction = true);

    /** @brief Simulate the particle system inside the Penning trap over
     *  a certain amount of time.
     *
     *  @param time The time to simulate in microseconds
     *  @param steps The amount of steps for the whole simulation
     *  @param method The method to use when moving forward a timestep
     *  @param particle_interaction Turn particle interactions on/off
     *
     *  @return simulation_t
     * */
    simulation_t simulate(double time, uint steps, std::string method = "rk4",
                          bool particle_interaction = true);

    /** @brief Simulate and write the displacement of all particles to files.
     *
     *  @param path The directory to save the data
     *  @param time The time to simulate in microseconds
     *  @param steps The amount of steps for the whole simulation
     *  @param method The method to use when moving forward a timestep
     *  @param particle_interaction Turn particle interactions on/off
     * */
    void write_simulation_to_dir(std::string path, double time, uint steps,
                                 std::string method = "rk4",
                                 bool particle_interaction = true);

    /** @brief Simulate and calculate what fraction of particles are still
     *  left inside the Penning trap after the simulation.
     *
     *  @param time The time to simulate in microseconds
     *  @param steps The amount of steps for the whole simulation
     *  @param method The method to use when moving forward a timestep
     *  @param particle_interaction Turn particle interactions on/off
     *
     *  @return double
     * */
    double fraction_of_particles_left(double time, uint steps,
                                      std::string method = "rk4",
                                      bool particle_interaction = true);

};

#endif
