/** @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 add_particle
 *  @todo Implement external_E_field
 *  @todo Implement external_B_field
 *  @todo Implement force_on_particle
 *  @todo Implement total_force_external
 *  @todo Implement total_force_particles
 *  @todo Implement total_force
 *  @todo Implement evolve_RK4
 *  @todo Implement evolve_forward_euler
 * */

#include "PenningTrap.hpp"
#include "constants.hpp"
#include <algorithm>
#include <stdexcept>

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

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

arma::vec PenningTrap::external_E_field(arma::vec r)
{
    arma::vec::fixed<3> res;
    double x = r(0), y = r(1), z = r(2);
    double f = this->V_0/2*this->d*this->d;
    res(0) = f*2*x;
    res(1) = f*2*y;
    res(2) = -f*4*z;

    return res;
}

arma::vec PenningTrap::external_B_field(arma::vec r)
{
    arma::vec::fixed<3> res;
    res(0) = 0.;
    res(1) = 0.;
    res(2) = this->B_0;

    return res;
}

arma::vec PenningTrap::force_on_particle(int i, int j)
{
    // Calculate the difference between the particles' position
    arma::vec::fixed<3> res = this->particles.at(i).r_vec 
        - this->particles.at(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.at(j).q/(norm*norm*norm);

    return res;
}

arma::vec PenningTrap::total_force_external(int i)
{
    Particle p = this->particles.at(i);
    arma::vec force = p.q
        *(this->external_E_field(p.r_vec) + this->external_B_field(p.r_vec));

    return force;
}

arma::vec PenningTrap::total_force_particles(int i)
{
    Particle p = this->particles.at(i);

    arma::vec::fixed<3> res;

    for (int j=0; j < this->particles.size(); j++) {
        if (i == j) {
            continue;
        }

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

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

    return res;
}

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

void PenningTrap::evolve_RK4(double dt)
{

}

void PenningTrap::evolve_forward_euler(double dt)
{

}