#include <armadillo>
#include <ostream>

#ifndef CLASSES_HPP
#define CLASSES_HPP

class PenningTrap;

arma::vec standard_external_E_field(const arma::vec& r, double t, const PenningTrap& trap);
arma::vec oscillating_external_E_field(const arma::vec& r, double t, const PenningTrap& trap);
arma::vec standard_external_B_field(const arma::vec& r, double t, const PenningTrap& trap);
arma::vec limited_range_external_B_field(const arma::vec& r, double t, const PenningTrap& trap);

class Particle {
    public:
        arma::vec position;
        arma::vec velocity;
        double mass;
        double charge;
        Particle(arma::vec pos, arma::vec vel, double m, double q);
        Particle(double m, double q);
        std::string info() const;
        bool inside_trap(double d) const;

};

class PenningTrap {
    private:
        std::vector<Particle> particles;
        bool coulomb_interaction = true;
        bool modified_fields = false;
        arma::vec (*_external_E_field)(const arma::vec& r, double t, const PenningTrap& trap) = standard_external_E_field;
        arma::vec (*_external_B_field)(const arma::vec& r, double t, const PenningTrap& trap) = standard_external_B_field;
    public:
        double B0;
        double V0;
        double d;
        double f = 0.0; // Amplitude of oscillating potential
        double omega_V = 0.0; // Angular frequency of oscillating potential
        PenningTrap(double B0_in, double V0_in, double d_in);
        void add_particle(Particle& p);
        void add_n_identical_particles(int n, double m, double q);
        arma::vec external_E_field(const arma::vec& r, double t) const {
            return _external_E_field(r, t, *this);
        }
        arma::vec external_B_field(const arma::vec& r, double t) const {
            return _external_B_field(r, t, *this);
        }
        void set_external_E_field(arma::vec (*E_field_func)(const arma::vec&, double, const PenningTrap&));
        void set_external_B_field(arma::vec (*B_field_func)(const arma::vec&, double, const PenningTrap&));
        arma::vec force_particle(int i, int j) const;
        arma::vec total_force_external(int i, double t) const;
        arma::vec total_force_particles(int i) const;
        arma::vec total_force(int i, double t) const;
        double potential_energy(int i) const;
        double total_potential_energy() const;
        double kinetic_energy(int i) const;
        double total_kinetic_energy() const;
        double total_energy() const;
        Particle& get_particle(int i);
        const Particle& get_particle(int i) const;
        int num_particles() const;
        void enable_interactions() {
            coulomb_interaction = true;
        }
        void disable_interactions() {
            coulomb_interaction = false;
        }
        bool get_coulomb_interaction() const {
            return coulomb_interaction;
        }
        bool get_modified_fields() const {
            return modified_fields;
        }
        std::string info() const;
        int num_inside_trap() const;
};


std::ostream& operator<<(std::ostream& os, const Particle& p);

std::ostream& operator<<(std::ostream& os, const PenningTrap& trap);


#endif