#include "classes.hpp"
#include "constants.hpp"
#include <iostream>

// Initialize particle with position and velocity
// Mass in atomic mass units (u), charge in elementary charges (e)
// Positions in meters, velocities in meters/second
Particle::Particle(arma::vec pos, arma::vec vel, double m, double q) 
    : position(pos * constants::meter), velocity(vel * constants::meter / constants::second), mass(m * constants::amu), charge(q * constants::elementary_charge) {}

// Initialize particle at origin with zero velocity
Particle::Particle(double m, double q) 
    : position(arma::vec({0,0,0})), velocity(arma::vec({0,0,0})), mass(m * constants::amu), charge(q * constants::elementary_charge) {}

std::string Particle::info() const {
    std::string info_str = "Particle (m = " + std::to_string(mass) + " u, q = " + std::to_string(charge) + " e)\n";
    info_str += "Position (um): [" + std::to_string(position(0)) + ", " + std::to_string(position(1)) + ", " + std::to_string(position(2)) + "]\n";
    info_str += "Velocity (um/us): [" + std::to_string(velocity(0)) + ", " + std::to_string(velocity(1)) + ", " + std::to_string(velocity(2)) + "]\n";
    return info_str;
}

bool Particle::inside_trap(double d) const {
    return arma::norm(position) <= d;
}

// Empty Penning Trap with specified field parameters
// B0 in tesla, V0 in volt, d in meters
PenningTrap::PenningTrap(double B0_in, double V0_in, double d_in) 
    : B0(B0_in * constants::tesla), V0(V0_in * constants::volt), d(d_in * constants::meter) {}

// Add particle to trap
void PenningTrap::add_particle(Particle& p) {
    particles.push_back(p);
}

// Add n identical particles with mass m and charge q
void PenningTrap::add_n_identical_particles(int n, double m, double q) {
    for (int i = 0; i < n; ++i) {
        arma::vec pos = arma::vec(3).randn() * 0.1 * d; // Random position within 10% of d (unit: um)
        arma::vec vel = arma::vec(3).randn() * 0.1 * d; // Random velocity within 10% of d/s (unit: um/us)
        Particle p(pos / constants::meter, vel, m, q);
        add_particle(p);
    }
}

void PenningTrap::set_external_E_field(arma::vec (*E_field_func)(const arma::vec&, double, const PenningTrap&)) {
    _external_E_field = E_field_func;
    modified_fields = true;
}

void PenningTrap::set_external_B_field(arma::vec (*B_field_func)(const arma::vec&, double, const PenningTrap&)) {
    _external_B_field = B_field_func;
    modified_fields = true;
}


// Force on particle i from particle j due to Coulomb interaction
// Neglecting magnetic fields from moving charges
arma::vec PenningTrap::force_particle(int i, int j) const {
    if (!coulomb_interaction) {
        return arma::vec({0, 0, 0});
    }
    arma::vec r_i = particles[i].position;
    arma::vec r_j = particles[j].position;
    arma::vec r_ij = r_i - r_j;
    double distance = arma::norm(r_ij);
    if (distance == 0) {
        return arma::vec({0, 0, 0}); // Avoid division by zero
    }
    double force_magnitude_per_dist = (constants::k_e * particles[i].charge * particles[j].charge) / (distance * distance * distance);
    return force_magnitude_per_dist * r_ij;
}

// Total force on particle i from external fields
// Sum of electric and magnetic forces
arma::vec PenningTrap::total_force_external(int i, double t) const {
    arma::vec E = external_E_field(get_particle(i).position, t);
    arma::vec B = external_B_field(get_particle(i).position, t);
    arma::vec v = get_particle(i).velocity;
    arma::vec F_electric = get_particle(i).charge * E;
    arma::vec F_magnetic = get_particle(i).charge * arma::cross(v, B);
    return F_electric + F_magnetic;
}

// Total force on particle i from all other particles
arma::vec PenningTrap::total_force_particles(int i) const {
    arma::vec total_force = arma::vec({0, 0, 0});
    if (!coulomb_interaction) { // Add the check here as well to avoid unnecessary calls to force_particle
        return total_force;
    }
    for (size_t j = 0; j < particles.size(); ++j) {
        if (j != i) {
            total_force += force_particle(i, j);
        }
    }
    return total_force;
}

// Total force on particle i from all other particles and external fields
arma::vec PenningTrap::total_force(int i, double t) const {
    return total_force_external(i, t) + total_force_particles(i);
}

// Potential energy of particle i in the trap
double PenningTrap::potential_energy(int i) const {
    if (modified_fields) {
        throw std::runtime_error("Potential energy calculation not valid with modified external fields.");
    }
    double q = particles[i].charge;
    arma::vec r = particles[i].position;
    double V = (V0 / (d * d)) * (r(2) * r(2) - 0.5 * (r(0) * r(0) + r(1) * r(1)));
    if (coulomb_interaction) {
        // Add contributions from other particles
        for (size_t j = 0; j < particles.size(); ++j) {
            if (j != i) {
                arma::vec r_j = particles[j].position;
                double distance = arma::norm(r - r_j);
                if (distance != 0) {
                    V += (constants::k_e * particles[j].charge) / distance;
                }
            }
        }
    }
    return q * V;
}

// Total potential energy of all particles in the trap
double PenningTrap::total_potential_energy() const {
    if (modified_fields) {
        throw std::runtime_error("Potential energy calculation not valid with modified external fields.");
    }
    double total_U = 0.0;
    for (size_t i = 0; i < particles.size(); ++i) {
        total_U += potential_energy(i);
    }
    return total_U;
}

// Kinetic energy of particle i
double PenningTrap::kinetic_energy(int i) const {
    double m = particles[i].mass;
    arma::vec v = particles[i].velocity;
    return 0.5 * m * arma::dot(v, v);
}
// Total kinetic energy of all particles in the trap
double PenningTrap::total_kinetic_energy() const {
    double total_K = 0.0;
    for (size_t i = 0; i < particles.size(); ++i) {
        total_K += kinetic_energy(i);
    }
    return total_K;
}

// Total energy of all particles in the trap
double PenningTrap::total_energy() const {
    return total_kinetic_energy() + total_potential_energy();
}


// Getter function for particles
Particle& PenningTrap::get_particle(int i) {
    return particles[i];
}

const Particle& PenningTrap::get_particle(int i) const {
    return particles[i];
}

// Number of particles in trap
int PenningTrap::num_particles() const{
    return particles.size();
}


std::string PenningTrap::info() const {
    std::string info_str = "Penning Trap (B0 = " + std::to_string(B0) + " u/(us*e), V0 = " + std::to_string(V0) + " u*um^2/(us^2*e), d = " + std::to_string(d) + " um)\n";
    info_str += "Number of particles: " + std::to_string(num_particles()) + "\n";
    info_str += "Coulomb interactions: " + std::string(coulomb_interaction ? "enabled" : "disabled") + "\n";
    return info_str;
}

int PenningTrap::num_inside_trap() const {
    int count = 0;
    for (int i = 0; i < num_particles(); ++i) {
        if (particles[i].inside_trap(d)) {
            count++;
        }
    }
    return count;
}

// Electric field at position r and time t
// Electric field from V = V0 * (z^2 - (x^2 + y^2)/2) / d^2
arma::vec standard_external_E_field(const arma::vec& r, double t, const PenningTrap& trap) {
    return (trap.V0 / (trap.d * trap.d)) * arma::vec({r(0), r(1), -2 * r(2)});
}

// Magnetic field at position r and time t
// Constant magnetic field in z-direction
arma::vec standard_external_B_field(const arma::vec& r, double t, const PenningTrap& trap) {
    return arma::vec({0, 0, trap.B0});
}

arma::vec oscillating_external_E_field(const arma::vec& r, double t, const PenningTrap& trap) {
    if (arma::norm(r) > trap.d) {
        return arma::vec({0, 0, 0}); // No electric field outside the trap
    }
    return standard_external_E_field(r, t, trap) * (1 + trap.f * cos(trap.omega_V * t));
}

arma::vec limited_range_external_B_field(const arma::vec& r, double t, const PenningTrap& trap) {
    if (arma::norm(r) > trap.d) {
        return arma::vec({0, 0, 0}); // No magnetic field outside the trap
    }
    return standard_external_B_field(r, t, trap);
}

std::ostream& operator<<(std::ostream& os, const Particle& p) {
    return os << p.info();
}
std::ostream& operator<<(std::ostream& os, const PenningTrap& trap) {
    return os << trap.info();
}