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

#include <iostream>
#include <ctime>

using namespace std;

void print_first_particles(PenningTrap& trap, int num_particles) {
    for (int i = 0; i < num_particles && i < trap.num_particles(); ++i) {
        cout << "Particle " << i + 1 << ": " << endl << trap.get_particle(i) << endl;
    }
}

void save_timing(clock_t start, clock_t end, const string& filename) {
    double cpu_time_used = ((double)(end - start)) / CLOCKS_PER_SEC;
    ofstream file(filename);
    if (file.is_open()) {
        file << "CPU time used: " << cpu_time_used << " seconds" << endl;
        file.close();
    } else {
        cerr << "Unable to open file for writing timing information." << endl;
    }
}

int main(int argc, char* argv[]) {
    argparse::ArgumentParser program("many_particles");
    program.add_argument("-v", "--verbose")
        .help("Enable verbose output of first particles")
        .flag();
    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")
        .default_value(100)
        .scan<'i', int>();
    program.add_argument("-i", "--disable-interactions")
        .help("Disable Coulomb interactions")
        .default_value(true)
        .implicit_value(false);
    program.add_argument("-O", "--oscillating-potential")
        .help("Enable oscillating potential")
        .flag();
    program.add_argument("-f", "--amplitude")
        .help("Amplitude of oscillating potential")
        .default_value(0.0)
        .scan<'f', double>();
    program.add_argument("-w", "--frequency")
        .help("Angular frequency of oscillating potential")
        .default_value(0.0)
        .scan<'f', double>();
    program.add_argument("-r", "--reduced-output")
        .help("Reduce output for oscillating potential to only final count")
        .flag();
    program.add_argument("-E", "--euler")
        .help("Use Euler solver instead of RK4")
        .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();
    program.add_argument("-t", "--timing")
        .help("Enable timing of the simulation")
        .flag();

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

    bool verbose = program.get<bool>("--verbose");

    int N = program.get<int>("--steps");
    int n_particles = program.get<int>("--num_particles");
    bool interactions = program.get<bool>("--disable-interactions");
    bool oscillating_potential = program.get<bool>("--oscillating-potential");
    bool use_euler = program.get<bool>("--euler");
    bool use_velocity_verlet = program.get<bool>("--velocity-verlet");
    bool use_boris = program.get<bool>("--boris");
    bool timing = program.get<bool>("--timing");

    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();
    }


    bool reduced_output = program.get<bool>("--reduced-output"); // Only relevant if oscillating potential is enabled
    if (oscillating_potential) {
        trap.set_external_E_field(oscillating_external_E_field);
        trap.set_external_B_field(limited_range_external_B_field);
        trap.f = program.get<double>("--amplitude");
        trap.omega_V = program.get<double>("--frequency");
    }

    trap.add_n_identical_particles(n_particles, 40.078 * constants::amu, 1 * constants::elementary_charge);
    if (verbose) {
        cout << trap.info() << endl;
        print_first_particles(trap, 10); 
    }
       

    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 = "VelocityVerlet";
    } else if (use_boris) {
        solver = make_unique<BorisSolver>(trap, dt);
        solver_name = "Boris";
    } else {
        solver = make_unique<RK4Solver>(trap, dt);
        solver_name = "RK4";
    }
    clock_t start, end;
    start = clock();
    if (oscillating_potential) {
        solver->record_positions_and_velocities = false; // Disable recording for performance
        if (reduced_output) {
            solver->record_inside_trap_count = false; // Disable recording for performance
        }
        solver->simulate(N);
        end = clock();
        string folder = "results/" + solver_name + "/many_particles_" + to_string(trap.num_particles()) + "_particles_" + to_string(int(N)) + "_steps_" + (interactions ? "enabled" : "disabled") + "_interactions_oscillating_potential";
        string filename = "f" + to_string(trap.f) + "_w" + to_string(trap.omega_V);
        filesystem::create_directories(folder);

        if (reduced_output) {
            solver->record_current_step(); // Record only final step
        }
        solver->save_counts(folder + "/" + filename);
        cout << "Final number of particles inside trap: " << trap.num_inside_trap() << " out of " << trap.num_particles() << endl;
        if (timing) {
            save_timing(start, end, folder + "/" + filename + "_timing.txt");
        }
        return 0;

    }
    solver->simulate(N);
    end = clock();
    string folder = "results/" + solver_name;
    filesystem::create_directories(folder);
    string filename = "/many_particles_" + to_string(trap.num_particles()) + "_particles_" + to_string(int(N)) + "_steps_" + (interactions ? "enabled" : "disabled") + "_interactions";
    solver->save(folder + filename);
    if (verbose) {
        vector<arma::mat> positions = solver->get_positions();
        cout << "Final positions of particles: " << endl;
        print_first_particles(trap, 10);
    }
    if (timing) {
        save_timing(start, end, folder + filename + "_timing.txt");
    }

    return 0;
}

