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

#include <algorithm>
#include <functional>
#include <sys/types.h>
#include <vector>

#include "PenningTrap.hpp"
#include "typedefs.hpp"

vec3 PenningTrap::v_func(uint i, uint j, double dt)
{
    switch (i) {
    case 0:
        return .5 * dt * this->k_v[0][j];
    case 1:
        return .5 * dt * this->k_v[1][j];
    case 2:
        return dt * this->k_v[2][j];
    case 3:
        return vec3((dt / 6.)
                    * (this->k_v[0][j] + 2. * this->k_v[1][j]
                       + 2. * this->k_v[2][j] + this->k_v[3][j]));
    default:
        std::cout << "Not valid!" << std::endl;
        abort();
    }
}

vec3 PenningTrap::r_func(uint i, uint j, double dt)
{
    switch (i) {
    case 0:
        return .5 * dt * this->k_r[0][j];
    case 1:
        return .5 * dt * this->k_r[1][j];
    case 2:
        return dt * this->k_r[2][j];
    case 3:
        return vec3((dt / 6.)
                    * (this->k_r[0][j] + 2. * this->k_r[1][j]
                       + 2. * this->k_r[2][j] + this->k_r[3][j]));
    default:
        std::cout << "Not valid!" << std::endl;
        abort();
    }
}

vec3 PenningTrap::external_E_field(vec3 r)
{
    r(2) *= -2.;

    return vec3(
        ((this->V_0 * this->perturbation(this->t)) / (this->d * this->d)) * r);
}

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

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

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

    return vec3((this->particles[j].q / (norm * norm * norm)) * res);
}

vec3 PenningTrap::total_force_external(uint i)
{
    Particle *p = &this->particles[i];

    return vec3(p->q
                * (this->external_E_field(p->r_vec)
                   + arma::cross(p->v_vec, this->external_B_field(p->r_vec))));
}

vec3 PenningTrap::total_force_particles(uint i)
{
    vec3 res;

    for (size_t j = 0; j < this->particles.size(); j++) {
        if (i != j)
            res += this->force_on_particle(i, j);
    }

    return vec3(res * (K_E * this->particles[i].q));
}

vec3 PenningTrap::total_force(uint i)
{
    if (arma::norm(this->particles[i].r_vec) > this->d) {
        return vec3{0., 0., 0.};
    }
    return vec3(this->total_force_external(i) - this->total_force_particles(i));
}

vec3 PenningTrap::total_force_no_interaction(uint i)
{
    if (arma::norm(this->particles[i].r_vec) > this->d) {
        return vec3{0., 0., 0.};
    }
    return vec3(this->total_force_external(i));
}

PenningTrap::PenningTrap(double B_0, double V_0, double d, double t)
{
    this->B_0 = B_0;
    this->V_0 = V_0;
    this->d = d;
    this->t = t;
    this->perturbation = [](double t) { return 1.; };
}

PenningTrap::PenningTrap(uint i, double B_0, double V_0, double d, double t)
    : PenningTrap::PenningTrap(B_0, V_0, d)
{
    for (size_t j = 0; j < i; j++) {
        this->particles.push_back(
            Particle(vec3(vec3().randn() * .1 * this->d),
                     vec3(vec3().randn() * .1 * this->d)));
    }
}

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

void PenningTrap::set_pertubation(double f, double omega_V)
{
    this->perturbation = [f, omega_V](double t) {
        return 1 + f * std::cos(omega_V * t);
    };
}

void PenningTrap::reinitialize(double f, double omega_V, double t)
{
    this->t = t;
    this->set_pertubation(f, omega_V);
    Particle *p;

    for (size_t i = 0; i < this->particles.size(); i++) {
        p = &this->particles[i];
        p->v_vec = vec3().randn() * .1 * this->d;
        p->v_vec = vec3().randn() * .1 * this->d;
    }
}

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

void PenningTrap::evolve_RK4(double dt, bool particle_interaction)
{
    Particle *p;

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

    vec3 (PenningTrap::*force)(uint) =
        particle_interaction ? &PenningTrap::total_force
                             : &PenningTrap::total_force_no_interaction;

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

    // Allocating takes a long time, so reuse sim_arr if possible
    if (this->k_v.size() != 4 || this->k_r.size() != 4
        || this->k_v[0].size() != size || this->k_r[0].size() != size) {
        this->k_v = sim_arr(4, sim_cols(size));
        this->k_r = sim_arr(4, sim_cols(size));
    }

    // Each k_{i+1} is dependent on k_i, so outer loop is not parallelizable
    for (size_t i = 0; i < 4; i++) {
        // Inner loop is able to be parallelized
#pragma omp parallel for private(p)
        for (size_t j = 0; j < size; j++) {
            this->k_v[i][j] = (this->*force)(j) / this->particles[j].m;
            this->k_r[i][j] = this->particles[j].v_vec;

            p = &tmp_particles[j];

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

    this->t += dt;
}

void PenningTrap::evolve_forward_euler(double dt, bool particle_interaction)
{
    size_t size = this->particles.size();
    vec3 force_res[size];
    Particle *p;

    vec3 (PenningTrap::*force)(uint) =
        particle_interaction ? &PenningTrap::total_force
                             : &PenningTrap::total_force_no_interaction;

    // Calculating the force for each particle is independent and therefore
    // a good candidate for parallel execution
#pragma omp parallel for
    for (size_t i = 0; i < size; i++) {
        force_res[i] = (this->*force)(i);
    }

    // Updating the particles is also independent, so we can parallelize
    // this as well
#pragma omp parallel for
    for (size_t i = 0; i < size; i++) {
        p = &this->particles[i];
        p->r_vec += dt * p->v_vec;
        p->v_vec += dt * force_res[i] / p->m;
    }

    this->t += dt;
}

simulation_t PenningTrap::simulate(double time, uint steps, std::string method,
                                   bool particle_interaction)
{
    Particle *p;
    double dt = time / (double)steps;

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

    simulation_t res{sim_arr(size, sim_cols(steps)),
                     sim_arr(size, sim_cols(steps))};

    void (PenningTrap::*func)(double, bool);
    if (method == "rk4") {
        func = &PenningTrap::evolve_RK4;
    } else if (method == "euler") {
        func = &PenningTrap::evolve_forward_euler;
    } else {
        std::cout << "Not a valid method!" << std::endl;
        abort();
    }

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

    return res;
}

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

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

    std::ofstream ofile;

    // Writing each particle to its own file is independent and can be run in
    // parallel.
#pragma omp parallel for private(ofile)
    for (size_t i = 0; i < this->particles.size(); i++) {
        ofile.open(path + "particle_" + std::to_string(i) + "_r.txt");
        for (vec3 &vec : res.r_vecs[i]) {
            ofile << scientific_format(vec(0), 10, 8) << ','
                  << scientific_format(vec(1), 10, 8) << ','
                  << scientific_format(vec(2), 10, 8) << '\n';
        }
        ofile.close();
        ofile.open(path + "particle_" + std::to_string(i) + "_v.txt");
        for (vec3 &vec : res.v_vecs[i]) {
            ofile << scientific_format(vec(0), 10, 8) << ','
                  << scientific_format(vec(1), 10, 8) << ','
                  << scientific_format(vec(2), 10, 8) << '\n';
        }
        ofile.close();
    }
}

double PenningTrap::fraction_of_particles_left(double time, uint steps,
                                               std::string method,
                                               bool particle_interaction)
{
    double dt = time / (double)steps;

    void (PenningTrap::*func)(double, bool);
    if (method == "rk4") {
        func = &PenningTrap::evolve_RK4;
    } else if (method == "euler") {
        func = &PenningTrap::evolve_forward_euler;
    } else {
        std::cout << "Not a valid method!" << std::endl;
        abort();
    }

    for (size_t j = 0; j < steps; j++) {
        (this->*func)(dt, particle_interaction);
    }

    int particles_left = 0;

    // A reduction is perfect here
#pragma omp parallel for reduction(+ : particles_left)
    for (size_t i = 0; i < this->particles.size(); i++) {
        if (arma::norm(this->particles[i].r_vec) < this->d) {
            particles_left++;
        }
    }

    return (double)particles_left / (double)this->particles.size();
}