#include <armadillo>
#include "classes.hpp"
#include "solvers.hpp"
#include "constants.hpp"
#include "argparse/argparse.hpp"

#include <iostream>

using namespace std;

int main(int argc, char* argv[]) {
    argparse::ArgumentParser program("two_particles");
    program.add_argument("-N", "--steps")
        .help("Number of time steps")
        .default_value(4000)
        .scan<'i', int>();
    program.add_argument("-n", "--num_particles")
        .help("Number of particles (either 1 or 2)")
        .default_value(2)
        .scan<'i', int>();
    program.add_argument("-i", "--disable-interactions")
        .help("Disable Coulomb interactions")
        .default_value(true)
        .implicit_value(false);
    program.add_argument("-E", "--euler")
        .help("Use Euler solver instead of RK4")
        .flag();
    program.add_argument("-A", "--analytical")
        .help("Use Analytical solver instead of RK4 (only valid for particles without interactions and specific initial conditions)")
        .flag();
    program.add_argument("-V", "--velocity-verlet")
        .help("Use Velocity Verlet solver instead of RK4")
        .flag();
    program.add_argument("-B", "--boris")
        .help("Use Boris solver instead of RK4")
        .flag();

    try {
        program.parse_args(argc, argv);
    } catch (const runtime_error& err) {
        cerr << err.what() << endl;
        cerr << program.help().str() << endl;
        exit(1);
    }

    double N = program.get<int>("--steps");
    int n_particles = program.get<int>("--num_particles");
    bool interactions = program.get<bool>("--disable-interactions");
    bool use_euler = program.get<bool>("--euler");
    bool use_analytical = program.get<bool>("--analytical");
    bool use_velocity_verlet = program.get<bool>("--velocity-verlet");
    bool use_boris = program.get<bool>("--boris");

    double B_0 = 1.0; // Tesla
    double V_0 = 0.025; // Volt
    double d = 500e-6; // meter

    PenningTrap trap(B_0, V_0, d);
    if (interactions) {
        trap.enable_interactions();
    }
    else {
        trap.disable_interactions();
    }
    cout << trap << endl;


    Particle p1(arma::vec({20e-6, 0, 20e-6}), arma::vec({0, 25, 0}), 40.078 * constants::amu, 1 * constants::elementary_charge);
    Particle p2(arma::vec({25e-6, 25e-6, 0}), arma::vec({0, 40, 5}), 40.078 * constants::amu, 1 * constants::elementary_charge);

    trap.add_particle(p1);
    if (n_particles > 1) {
        trap.add_particle(p2);
    }

    double dt = 50e-6 / N; // seconds

    unique_ptr<Solver> solver;
    string solver_name;

    if (use_euler) {
        solver = make_unique<EulerSolver>(trap, dt);
        solver_name = "euler";
    } else if (use_velocity_verlet) {
        solver = make_unique<VelocityVerletSolver>(trap, dt);
        solver_name = "velocity_verlet";
    } else if (use_boris) {
        solver = make_unique<BorisSolver>(trap, dt);
        solver_name = "boris";
    } else if (use_analytical) {
        solver = make_unique<AnalyticalSolver>(trap, dt);
        solver_name = "analytical";
    } else {
        solver = make_unique<RK4Solver>(trap, dt);
        solver_name = "rk4";
    }

    solver->simulate(N);
    solver->save("results/two_particles_" + std::to_string(n_particles) + (interactions ? "_with_interactions_" : "_no_interactions_") + std::to_string(N) + "_steps_" + solver_name);
    vector<arma::mat> positions = solver->get_positions();
    cout << "Final position of particle 1: " << endl;
    cout << positions.back().col(0) << endl;
    if (n_particles > 1) {
        cout << "Final position of particle 2: " << endl;
        cout << positions.back().col(1) << endl;
    }
}