#include "classes.hpp"
#include <vector>
#include <armadillo>

#ifndef SOLVERS_HPP
#define SOLVERS_HPP


arma::cube to_cube(const std::vector<arma::mat>& vec);

void save_to_file(const std::string& filename, const arma::cube& data);
void save_to_file(const std::string& filename, const arma::vec& data);


class Solver {
    protected:
        PenningTrap& trap;
        double dt;
        int steps;
        std::vector<arma::mat> positions;
        std::vector<arma::mat> velocities;
        std::vector<int> inside_trap_count;
        std::vector<double> potential_energies;
        std::vector<double> kinetic_energies;
        std::vector<double> total_energies;
        virtual void step() = 0;
        void record_position();
        void record_velocity();
        void record_energies();
        void record_trap_count();
    public:
        Solver(PenningTrap& trap_in, double dt_in);
        void simulate(int num_steps);
        void simulate_time(double total_time) {
            int num_steps = total_steps(total_time);
            simulate(num_steps);
        }
        std::vector<arma::mat> get_positions();
        std::vector<arma::mat> get_velocities();
        std::vector<double> get_potential_energies();
        std::vector<double> get_kinetic_energies();
        std::vector<double> get_total_energies();
        std::vector<int> get_inside_trap_count();
        int get_current_step();
        int total_steps(double total_time);
        void save(std::string base_filename);
        void save_counts(std::string base_filename);
        void record_current_step() {
            record_position();
            record_velocity();
            if (!trap.get_modified_fields()) {
                record_energies();
            }
            record_trap_count();
        }
        bool record_positions_and_velocities = true;
        bool record_inside_trap_count = true;
};

class EulerSolver : public Solver {
    public:
        EulerSolver(PenningTrap& trap_in, double dt_in);
        void step() override;
};
class RK4Solver : public Solver {
    public:
        RK4Solver(PenningTrap& trap_in, double dt_in);
        void step() override;
};
class VelocityVerletSolver : public Solver {
    private:
        std::vector<arma::vec> accelerations;
    public:
        VelocityVerletSolver(PenningTrap& trap_in, double dt_in);
        void step() override;
};


class BorisSolver : public Solver {
    public:
        BorisSolver(PenningTrap& trap_in, double dt_in);
        void step() override;
};

class AnalyticalSolver : public Solver {
    private:
        std::vector<double> w_0;
        std::vector<double> w_z;
        std::vector<double> w_plus;
        std::vector<double> w_minus;
        std::vector<double> A_plus;
        std::vector<double> A_minus;
        std::vector<double> z_0;
    public:
        AnalyticalSolver(PenningTrap& trap_in, double dt_in);
        void step() override;
};

#endif