/** @file WaveSimulation.hpp
 *
 *  @author Cory Alexander Balaton (coryab)
 *  @author Janita Ovidie Sandtrøen Willumsen (janitaws)
 *
 *  @version 0.1
 *
 *  @brief The definition of the WaveSimulation class
 *
 *  @bug No known bugs
 * */
#ifndef __WAVE_SIMULATION__
#define __WAVE_SIMULATION__

#include "constants.hpp"
#include "literals.hpp"

#include <armadillo>
#include <cstdint>

/** @brief Simulate the evolution of a wave packet in 2 + 1 domensions.*/
class WaveSimulation {
private:
    int32_t M;     ///< The size
    int32_t N;     ///< The size of the inner part
    arma::cx_mat V; ///< The potential matrix
    arma::cx_mat U; ///< The particle wave matrix
    arma::sp_cx_mat B; ///< The B matrix
    arma::sp_cx_mat A; ///< The A matrix
    double h; ///< The step size in both x and y direction
    double dt; ///< The step size int the time direction
    double T; ///< The end time

    /**  @brief Initialize the U matrix using an unormalized Gaussian wave
     * packet.
     *
     *  @param x_c The center of the packet in the x direction.
     *  @param y_c The center of the packet in the y direction.
     *  @param sigma_x The The initial width in the x direction.
     *  @param sigma_y The The initial width in the y direction.
     *  @param p_x The wave packet momentum in the x direction.
     *  @param p_y The wave packet momentum in the y direction.
     * */
    void initialize_U(double x_c, double y_c, double sigma_x, double sigma_y,
                      double p_x, double p_y);

    /** @brief Initialize the V matrix.
     *
     *  @param thickness The thickness of the wall in the x direction.
     *  @param pos_x The center of the wall in the x direction.
     *  @param ap_sep The separation between each aperture.
     *  @param ap The aperture width.
     *  @param slits The number of slits.
     * */
    void initialize_V(double thickness, double pos_x,
                      double aperture_separation, double aperture,
                      uint32_t slits);

    /** @brief Initialize the V matrix with no wall.
     * */
    void initialize_V();

    /** @brief Initialize the A matrix according to the Crank-Nicolson method
     * */
    void initialize_A();

    /** @brief Initialize the B matrix according to the Crank-Nicolson method
     * */
    void initialize_B();

    /** @brief Write the U matrix in a single line to the file buffer given
     *
     *  @param ofile The file buffer to write to
     * */
    void write_U(std::ofstream &ofile);
public:
    /** @brief Constructor for the WaveSimulation class.
     *
     *  @param h The step size in the x and y direction.
     *  @param dt The step size in the temporal dimension.
     *  @param T The total time to simulate.
     *  @param x_c The center of the packet in the x direction.
     *  @param y_c The center of the packet in the y direction.
     *  @param sigma_x The The initial width in the x direction.
     *  @param sigma_y The The initial width in the y direction.
     *  @param p_x The wave packet momentum in the x direction.
     *  @param p_y The wave packet momentum in the y direction.
     *  @param thickness The thickness of the wall in the x direction.
     *  @param pos_x The center of the wall in the x direction.
     *  @param ap_sep The separation between each aperture.
     *  @param ap The aperture width.
     *  @param slits The number of slits.
     * */
    WaveSimulation(double h, double dt, double T, double x_c, double y_c,
                   double sigma_x, double sigma_y, double p_x, double p_y,
                   double thickness, double pos_x, double ap_sep, double ap,
                   uint32_t slits);

    /** @brief Constructor for the WaveSimulation class with no wall.
     *
     *  @param h The step size in the x and y direction.
     *  @param dt The step size in the temporal dimension.
     *  @param T The total time to simulate.
     *  @param x_c The center of the packet in the x direction.
     *  @param y_c The center of the packet in the y direction.
     *  @param sigma_x The The initial width in the x direction.
     *  @param sigma_y The The initial width in the y direction.
     *  @param p_x The wave packet momentum in the x direction.
     *  @param p_y The wave packet momentum in the y direction.
     * */
    WaveSimulation(double h, double dt, double T, double x_c, double y_c,
                   double sigma_x, double sigma_y, double p_x, double p_y);

    /** @brief Evolve a step forward in time
     * */
    void step();

    /** @brief Evolve the wave packet until the time T has been reached and 
     *  write U to file.
     *
     *  @param outfile The name of the file to write to
     *  @param write_each_step Boolean for deciding to write each step to file 
     *  or just the last step
     * */
    void simulate(std::string outfile, bool write_each_step = false);
    
    /** @brief Evolve the wave packet and write U to fileto file at each time step in
     *  the vector given.
     *
     *  @param outfile The name of the file to write to
     *  @param steps What time steps to write U to file.
     * */
    void simulate(std::string outfile, std::vector<double> &steps);

    /** @brief Write the deviation of the sum of the probability of U from 1
     *
     *  @param outfile The name of the file to write to
     *  to file*/
    void probability_deviation(std::string outfile,
                               bool write_each_step = false);
};

#endif