/** @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

/** @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 = .0005;
    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)PARTICLES);
#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_fine.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 = .0005;
    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)PARTICLES);
#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_fine.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()
{
    double start, end;

    start = omp_get_wtime();

    potential_resonance_narrow_sweep();

    potential_resonance_narrow_sweep_interaction();

    end = omp_get_wtime();

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

    return 0;
}