/** @file PenningTrap.cpp
 *
 *  @author Cory Alexander Balaton (coryab)
 *  @author Janita Ovidie Sandtrøen Willumsen (janitaws)
 *
 *  @version 0.1
 *
 *  @brief The implementation of the PenningTrap class.
 *
 *  @bug No known bugs
 *
 *  @todo Implement evolve_RK4
 *  @todo Implement evolve_forward_euler
 * */

#include "PenningTrap.hpp"
#include "constants.hpp"
#include "typedefs.hpp"
#include "utils.hpp"
#include <cstdlib>
#include <functional>
#include <string>
#include <vector>

PenningTrap::PenningTrap(double B_0, double V_0, double d)
{
    this->B_0 = B_0;
    this->V_0 = V_0;
    this->d = d;

    // Create lambda functions for each iteration in RK4
    v_funcs[0] = [this](int i, double dt) { return .5 * dt * this->k_v[0][i]; };
    v_funcs[1] = [this](int i, double dt) { return .5 * dt * this->k_v[1][i]; };
    v_funcs[2] = [this](int i, double dt) { return dt * this->k_v[2][i]; };
    v_funcs[3] = [this](int i, double dt) {
        return (dt / 6) * (this->k_v[0][i].eval() + this->k_v[1][i].eval() +
                           this->k_v[2][i].eval() + this->k_v[3][i].eval());
    };

    r_funcs[0] = [this](int i, double dt) { return .5 * dt * this->k_r[0][i]; };
    r_funcs[1] = [this](int i, double dt) { return .5 * dt * this->k_r[0][i]; };
    r_funcs[2] = [this](int i, double dt) { return dt * this->k_r[2][i]; };
    r_funcs[3] = [this](int i, double dt) {
        return (dt / 6) * (this->k_r[0][i].eval() + this->k_r[1][i].eval() +
                           this->k_r[2][i].eval() + this->k_r[3][i].eval());
    };
}

PenningTrap::PenningTrap(int i, double B_0, double V_0, double d)
    : PenningTrap::PenningTrap(B_0, V_0, d)
{
    vec_3d r, v;
    for (int j = 0; j < i; j++) {
        r = vec_3d().randn() * .1 * this->d;
        v = vec_3d().randn() * .1 * this->d;
        this->add_particle(Particle(1., 40., r, v));
    }
}

PenningTrap::PenningTrap(std::vector<Particle> particles, double B_0,
                         double V_0, double d)
    : PenningTrap::PenningTrap(B_0, V_0, d)
{
    this->particles = particles;
}

void PenningTrap::add_particle(Particle particle)
{
    this->particles.push_back(particle);
}

vec_3d PenningTrap::external_E_field(vec_3d r)
{
    r(2) *= -2.;
    double f = this->V_0 / (this->d * this->d);

    return f * r;
}

vec_3d PenningTrap::external_B_field(vec_3d r)
{
    return vec_3d{0.,0.,this->B_0};
}

vec_3d PenningTrap::force_on_particle(int i, int j)
{
    // Calculate the difference between the particles' position
    vec_3d res = this->particles[i].r_vec - this->particles[j].r_vec;

    // Get the distance between the particles
    double norm = arma::norm(res);

    // Multiply res with p_j's charge divided by the norm cubed
    res *= this->particles[j].q / (norm * norm * norm);

    return res;
}

vec_3d PenningTrap::total_force_external(int i)
{
    Particle p = this->particles[i];

    if (arma::norm(p.r_vec) > this->d) {
        return vec_3d{0., 0., 0.};
    }

    vec_3d B = this->external_B_field(p.r_vec);

    vec_3d v_cross_B{p.v_vec(1) * B(2) - p.v_vec(2) * B(1),
                     p.v_vec(2) * B(0) - p.v_vec(0) * B(2),
                     p.v_vec(0) * B(1) - p.v_vec(1) * B(0)};

    vec_3d force = p.q * (this->external_E_field(p.r_vec) + v_cross_B);

    return force;
}

vec_3d PenningTrap::total_force_particles(unsigned int i)
{
    Particle p = this->particles[i];

    vec_3d res;
    size_t size = this->particles.size();

    for (size_t j = 0; j < size; j++) {
        if (i == j) {
            continue;
        }

        res += this->force_on_particle(i, j);
    }

    res *= K_E * (p.q / p.m);

    return res;
}

vec_3d PenningTrap::total_force(int i)
{
    return this->total_force_external(i) - this->total_force_particles(i);
}

void PenningTrap::evolve_RK4(double dt, bool particle_interaction)
{

    std::vector<Particle> original_particles = this->particles;
    std::vector<Particle> tmp_particles = this->particles;

    std::function<vec_3d(int)> force;
    if (particle_interaction) {
        force = [this](int i) { return this->total_force(i); };
    }
    else {
        force = [this](int i) { return this->total_force_external(i); };
    }

    int size = this->particles.size();

    this->k_v = sim_arr(4, sim_cols(size));
    this->k_r = sim_arr(4, sim_cols(size));

    for (int i = 0; i < 4; i++) {
#pragma omp parallel for
        for (int j = 0; j < size; j++) {
            this->k_v[i][j] = this->total_force(j) / this->particles[j].m;
            this->k_r[i][j] = this->particles[j].v_vec;

            Particle *p = &tmp_particles[j];

            p->v_vec = original_particles[j].v_vec + this->v_funcs[i](j, dt);
            p->r_vec = original_particles[j].r_vec + this->r_funcs[i](j, dt);
        }
        this->particles = tmp_particles;
    }
}

void PenningTrap::evolve_forward_euler(double dt, bool particle_interaction)
{
    std::vector<Particle> new_state = this->particles;

    Particle *p;

    std::function<vec_3d(int)> force;
    if (particle_interaction) {
        force = [this](int i) { return this->total_force(i); };
    }
    else {
        force = [this](int i) { return this->total_force_external(i); };
    }

#pragma omp parallel for private(p)
    for (size_t i = 0; i < this->particles.size(); i++) {
        p = &new_state[i];
        p->v_vec += dt * force(i) / p->m;
        p->r_vec += dt * this->particles[i].v_vec;
    }

    this->particles = new_state;
}

sim_arr PenningTrap::simulate(double time, unsigned int steps,
                              std::string method, bool particle_interaction)
{
    double dt = time / (double)steps;

    sim_arr res(this->particles.size(), sim_cols(steps));

    std::function<void(double, bool)> func;
    if (method == "rk4") {
        func = [this](double dt, bool particle_interaction) {
            this->evolve_RK4(dt, particle_interaction);
        };
    }
    else if (method == "euler") {
        func = [this](double dt, bool particle_interaction) {
            this->evolve_forward_euler(dt, particle_interaction);
        };
    }
    else {
        std::cout << "Not a valid method!" << std::endl;
        abort();
    }

    for (size_t j = 0; j < steps; j++) {
        for (size_t i = 0; i < this->particles.size(); i++) {
            res[i][j] = this->particles[i].r_vec;
        }
        func(dt, particle_interaction);
    }

    return res;
}

void PenningTrap::write_simulation_to_dir(std::string path, double time,
                                          int steps, std::string method,
                                          bool particle_interaction)
{
    if (path.back() != '/') {
        path += '/';
    }
    if (mkpath(path, 0777) != 0) {
        std::cout << "Hello" << std::endl;
        return;
    }

    sim_arr res = this->simulate(time, steps, method, particle_interaction);

    std::ofstream ofile;

#pragma omp parallel for private(ofile)
    for (size_t i = 0; i < this->particles.size(); i++) {
        ofile.open(path + "particle_" + std::to_string(i) + ".txt");
        for (vec_3d &vec : res[i]) {
            ofile << vec(0) << "," << vec(1) << "," << vec(2) << "\n";
        }
        ofile.close();
    }
}