/** @file main.cpp
 *
 *  @author Cory Alexander Balaton (coryab)
 *  @author Janita Ovidie Sandtrøen Willumsen (janitaws)
 *
 *  @version 1.0
 *
 *  @brief The main program for this project
 *
 *  @bug No known bugs
 * */

#include <cmath>
#include <complex>
#include <fstream>
#include <omp.h>
#include <string>
#include <vector>

#include "PenningTrap.hpp"
#include "constants.hpp"
#include "utils.hpp"

#define PARTICLES 100
#define N 40000

// Particles used for testing
Particle p1(vec3{20., 0., 20.}, vec3{0., 25., 0.}); ///< Particle 1
Particle p2(vec3{25., 25., 0.}, vec3{0., 40., 5.}); ///< Particle 2

/** @brief The analytical solution for particle p1
 *
 *  @param t Time
 *
 *  @return vec3
 * */
vec3 analytical_solution_particle_1(double t)
{
    double w_0 = T / CA_MASS;
    double w_z2 = (50. * V / 1000.) / (CA_MASS * 500. * 500.);
    double w_p = (w_0 + std::sqrt(w_0 * w_0 - 2. * w_z2)) / 2.;
    double w_n = (w_0 - std::sqrt(w_0 * w_0 - 2. * w_z2)) / 2.;
    double A_p = (25. + w_n * 20.) / (w_n - w_p);
    double A_n = -(25. + w_p * 20.) / (w_n - w_p);
    std::cout << A_p << "," << A_n << std::endl;
    std::complex<double> f =
        A_p * std::exp(std::complex<double>(0., -w_p * t))
        + A_n * std::exp(std::complex<double>(0., -w_n * t));
    vec3 res{std::real(f), std::imag(f), 20. * std::cos(std::sqrt(w_z2) * t)};
    return res;
}

/** @brief Simulate a single particle over the period of 50 \f$ \mu s \f$.
 * */
void simulate_single_particle()
{
    // Initialize trap with particle 1
    PenningTrap trap(std::vector<Particle>{p1});

    double time = 50.; // microseconds

    // Simulate and write results to file
    trap.write_simulation_to_dir("output/simulate_single_particle", time, N,
                                 "rk4", false);
}

/** @brief Simulate 2 particles over the period of 50 \f$ \mu s \f$ with and
 *  without particle interactions.
 * */
void simulate_two_particles()
{
    // Initialize traps with particles
    PenningTrap trap_no_interaction(std::vector<Particle>{p1, p2});
    PenningTrap trap_with_interaction(std::vector<Particle>{p1, p2});

    double time = 50.; // microseconds

    // Simulate and write results to files
    trap_no_interaction.write_simulation_to_dir(
        "output/simulate_2_particles/no_interaction", time, N, "rk4", false);
    trap_with_interaction.write_simulation_to_dir(
        "output/simulate_2_particles/with_interaction", time, N);
}

/** @brief Simulate a single particle over 50 \f$ \mu s \f$ using different
 *  amount of steps and different methods.
 * */
void simulate_single_particle_with_different_steps()
{
    double time = 50.; // microseconds

    std::ofstream ofile;

    // Calculate relative error for RK4
    std::string path = "output/relative_error/RK4/";
    mkpath(path);
#pragma omp parallel for private(ofile)
    for (int i = 0; i < 4; i++) {
        int steps = 4000 * std::pow(2, i);
        std::cout << steps << std::endl;
        double dt = time / (double)steps;
        ofile.open(path + std::to_string(steps) + "_steps.txt");
        PenningTrap trap(std::vector<Particle>{p1});
        simulation_t res = trap.simulate(time, steps, "rk4", false);
        for (int i = 0; i < steps; i++) {
            ofile << arma::norm(res.r_vecs[0][i]
                                - analytical_solution_particle_1(dt * i))
                  << '\n';
        }
        ofile.close();
    }

    // Calculate relative error for forward Euler
    path = "output/relative_error/euler/";
    mkpath(path);
#pragma omp parallel for private(ofile)
    for (int i = 0; i < 4; i++) {
        int steps = 4000 * std::pow(2, i);
        double dt = time / (double)steps;
        ofile.open(path + std::to_string(steps) + "_steps.txt");
        PenningTrap trap(std::vector<Particle>{p1});
        simulation_t res = trap.simulate(time, steps, "euler", false);
        for (int i = 0; i < steps; i++) {
            ofile << arma::norm(res.r_vecs[0][i]
                                - analytical_solution_particle_1(dt * i))
                  << '\n';
        }
        ofile.close();
    }
}

/** @brief Simulate 100 particles over 50 \f$ \mu s \f$.
 * */
void simulate_100_particles()
{
    PenningTrap trap((unsigned)100);

    double time = 50.; // microseconds

    // trap.write_simulation_to_dir("output/simulate_100_particles", time, N,
    //"rk4", false);
    trap.simulate(time, N, "rk4", true);
}

/** @brief Simulate 100 particles over 500 \f$ \mu s \f$ using a time
 *  dependent potential.
 *
 *  @details The simulation sweeps over different frequencies in [0.2, 2.5]
 * MHz.
 *
 * */
void potential_resonance_wide_sweep()
{
    double time = 500.;

    double amplitudes[]{.1, .4, .7};

    double freq_start = .2;
    double freq_end = 2.5;
    double freq_increment = .02;
    size_t freq_iterations =
        (size_t)((freq_end - freq_start) / freq_increment) + 1;

    double res[4][freq_iterations];

    std::string path = "output/time_dependent_potential/";
    mkpath(path);

    std::ofstream ofile;

#pragma omp parallel for
    // Insert frequencies
    for (size_t i = 0; i < freq_iterations; i++) {
        res[0][i] = freq_start + freq_increment * i;
    }

#pragma omp parallel
    {
        // Each thread creates a PenningTrap instance and reuses it throughout
        // the sweep.
        PenningTrap trap((uint)100);
#pragma omp for collapse(2)
        for (size_t i = 0; i < 3; i++) {
            for (size_t j = 0; j < freq_iterations; j++) {
                // Reset particles and give new time dependent potential.
                trap.reinitialize(amplitudes[i], res[0][j]);
                res[i + 1][j] =
                    trap.fraction_of_particles_left(time, N, "rk4", false);
            }
        }
    }

    // Write results to file
    ofile.open(path + "wide_sweep.txt");
    for (size_t i = 0; i < freq_iterations; i++) {
        ofile << res[0][i] << ',' << res[1][i] << ',' << res[2][i] << ','
              << res[3][i] << '\n';
    }
    ofile.close();
}

/** @brief Simulate 100 particles over 500 \f$ \mu s \f$ using a time
 *  dependent potential.
 *
 *  @details The simulation sweeps over different frequencies in [1., 1.7]
 * MHz.
 *
 * */
void potential_resonance_narrow_sweep()
{
    double time = 500.;

    double amplitudes[]{.1, .4, .7};

    double freq_start = 1.1;
    double freq_end = 1.7;
    double freq_increment = .002;
    size_t freq_iterations =
        (size_t)((freq_end - freq_start) / freq_increment) + 1;

    double res[4][freq_iterations];

    std::string path = "output/time_dependent_potential/";
    mkpath(path);

    std::ofstream ofile;

#pragma omp parallel for
    // Insert frequencies
    for (size_t i = 0; i < freq_iterations; i++) {
        res[0][i] = freq_start + freq_increment * i;
    }

#pragma omp parallel
    {
        // Each thread creates a PenningTrap instance and reuses it throughout
        // the sweep.
        PenningTrap trap((uint)100);
#pragma omp for collapse(2)
        for (size_t i = 0; i < 3; i++) {
            for (size_t j = 0; j < freq_iterations; j++) {
                // Reset particles and give new time dependent potential.
                trap.reinitialize(amplitudes[i], res[0][j]);
                res[i + 1][j] =
                    trap.fraction_of_particles_left(time, N, "rk4", false);
            }
        }
    }

    // Write results to file
    ofile.open(path + "narrow_sweep.txt");
    for (size_t i = 0; i < freq_iterations; i++) {
        ofile << res[0][i] << ',' << res[1][i] << ',' << res[2][i] << ','
              << res[3][i] << '\n';
    }
    ofile.close();
}

/** @brief Simulate 100 particles over 500 \f$ \mu s \f$ using a time
 *  dependent potential.
 *
 *  @details The simulation sweeps over different frequencies in [1., 1.7]
 * MHz.
 *
 * */
void potential_resonance_narrow_sweep_interaction()
{
    double time = 500.;

    double amplitudes[]{.1, .4, .7};

    double freq_start = 1.1;
    double freq_end = 1.7;
    double freq_increment = .002;
    size_t freq_iterations =
        (size_t)((freq_end - freq_start) / freq_increment) + 1;

    double res[4][freq_iterations];

    std::string path = "output/time_dependent_potential/";
    mkpath(path);

    std::ofstream ofile;

#pragma omp parallel for
    for (size_t i = 0; i < freq_iterations; i++) {
        res[0][i] = freq_start + freq_increment * i;
    }

#pragma omp parallel
    {
        // Each thread creates a PenningTrap instance and reuses it throughout
        // the sweep.
        PenningTrap trap((uint)100);
#pragma omp for collapse(2)
        for (size_t i = 0; i < 3; i++) {
            for (size_t j = 0; j < freq_iterations; j++) {
                // Reset particles and give new time dependent potential.
                trap.reinitialize(amplitudes[i], res[0][j]);
                res[i + 1][j] = trap.fraction_of_particles_left(time, N);
            }
        }
    }

    // Write results to file
    ofile.open(path + "narrow_sweep_interactions.txt");
    for (size_t i = 0; i < freq_iterations; i++) {
        ofile << res[0][i] << ',' << res[1][i] << ',' << res[2][i] << ','
              << res[3][i] << '\n';
    }
    ofile.close();
}

int main()
{
    int option = 1;
    bool chosen = false;

    system("clear");
    std::cout << "(1) All (default)\n"
              << "(2) Simulate single particle\n"
              << "(3) simulate 2 particles\n"
              << "(4) Simulate single particle with different time steps\n"
              << "(5) Simulate 100 particles\n"
              << "(6) Potential resonance wide sweep\n"
              << "(7) Potential resonance narrow sweep without particle "
                 "interactions\n"
              << "(8) Potential resonance narrow sweep with particle "
                 "interaction\n"
              << "Select what to run: ";
    std::cin >> std::noskipws;
    do {
        if (!(std::cin >> option) || option < 1 || option > 8) {
            std::cin.clear();
            std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
            system("clear");
            std::cout
                << "(1) All (default)\n"
                << "(2) Simulate single particle\n"
                << "(3) simulate 2 particles\n"
                << "(4) Simulate single particle with different time steps\n"
                << "(5) Simulate 100 particles\n"
                << "(6) Potential resonance wide sweep\n"
                << "(7) Potential resonance narrow sweep without particle "
                   "interactions\n"
                << "(8) Potential resonance narrow sweep with particle "
                   "interaction\n"
                << "Not a valid option, please enter a valid number: ";
        } else {
            chosen = true;
        }
    } while (!chosen);

    double start, end;

    system("clear");

    start = omp_get_wtime();
    switch (option) {
    case 1:
        std::cout << "Running simulate_single_particle\n";
        simulate_single_particle();

        std::cout << "Running simulate_two_particles\n";
        simulate_two_particles();

        std::cout << "Running simulate_single_particle_with_different_steps\n";
        simulate_single_particle_with_different_steps();

        std::cout << "Running simulate_100_particles\n";
        simulate_100_particles();

        std::cout << "Running potential_resonance_wide_sweep\n";
        potential_resonance_wide_sweep();

        std::cout << "Running potential_resonance_narrow_sweep\n";
        potential_resonance_narrow_sweep();

        std::cout << "Running potential_resonance_narrow_sweep_interaction\n";
        potential_resonance_narrow_sweep_interaction();
        break;
    case 2:
        std::cout << "Running simulate_single_particle\n";
        simulate_single_particle();
        break;
    case 3:
        std::cout << "Running simulate_two_particles\n";
        simulate_two_particles();
        break;
    case 4:
        std::cout << "Running simulate_single_particle_with_different_steps\n";
        simulate_single_particle_with_different_steps();
        break;
    case 5:
        std::cout << "Running simulate_100_particles\n";
        simulate_100_particles();
        break;
    case 6:
        std::cout << "Running potential_resonance_wide_sweep\n";
        potential_resonance_wide_sweep();
        break;
    case 7:
        std::cout << "Running potential_resonance_narrow_sweep\n";
        potential_resonance_narrow_sweep();
        break;
    case 8:
        std::cout << "Running potential_resonance_narrow_sweep_interaction\n";
        potential_resonance_narrow_sweep_interaction();
        break;
    }
    end = omp_get_wtime();

    std::cout << "Time: " << end - start << " seconds" << std::endl;

    return 0;
}