/** @file WaveSimulation.cpp
 *
 *  @author Cory Alexander Balaton (coryab)
 *  @author Janita Ovidie Sandtrøen Willumsen (janitaws)
 *
 *  @version 0.1
 *
 *  @brief Implementation of the WaveSimulation class.
 *
 *  @bug No known bugs
 * */
#include "WaveSimulation.hpp"

#include "utils.hpp"

#include <cmath>
#include <complex>
#include <cstdint>
#include <cstdlib>

WaveSimulation::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)
{
    this->dt = dt;
    this->h = h;
    this->T = T;
    this->M = 1. / h;
    this->N = M - 2;
    this->build_V(thickness, pos_x, ap_sep, ap, slits);
    this->initialize_U(x_c, y_c, sigma_x, sigma_y, p_x, p_y);
    this->build_A();
    this->build_B();
}

WaveSimulation::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)
{
    this->dt = dt;
    this->h = h;
    this->T = T;
    this->M = 1. / h;
    this->N = M - 2;
    this->build_V(0.,0.,0.,0., 0);
    this->initialize_U(x_c, y_c, sigma_x, sigma_y, p_x, p_y);
    this->build_A();
    this->build_B();

}

void WaveSimulation::solve(std::ofstream &ofile)
{
    ofile << this->N << '\n';
    uint32_t iterations = this->T / this->dt;
    for (size_t i = 0; i < iterations; i++) {
        this->write_U(ofile);
        this->step();
    }
}

void WaveSimulation::step()
{
    arma::cx_vec tmp = this->B * this->U.as_col();
    arma::spsolve(this->U, this->A, tmp);
}

void WaveSimulation::build_A()
{
    // Create the diagonal
    arma::cx_vec diagonal(this->N * this->N);

    // Set diagonal values
    std::complex<double> r = (1._i * this->dt) / (2 * h * h);
    for (size_t i = 0; i < diagonal.size(); i++) {
        diagonal(i) = 1. + 4. * r + (1._i * this->dt / 2.) * this->V(i);
    }

    // Create the submatrix
    arma::cx_mat sub_matrix(this->N, this->N, arma::fill::zeros);
    sub_matrix.diag(-1).fill(-r);
    sub_matrix.diag(1).fill(-r);

    // Set the size of A
    this->A.set_size(this->N * this->N, this->N * this->N);

    // Fill in the values in the submatrix diagonal
    for (size_t i = 0; i < this->A.n_cols; i += this->N) {
        this->A.submat(i, i, i + this->N - 1, i + this->N - 1) = sub_matrix;
    }

    // Fill the last sub/sup-diagonals
    this->A.diag() = diagonal;
    this->A.diag(-this->N).fill(-r);
    this->A.diag(this->N).fill(-r);
}

void WaveSimulation::build_B()
{
    std::complex<double> r = (1._i * this->dt) / (2 * h * h);

    // Create the diagonal
    arma::cx_vec diagonal(this->N * this->N);

    for (size_t i = 0; i < diagonal.size(); i++) {
        diagonal(i) = 1. - 4. * r - (1._i * this->dt / 2.) * this->V(i);
    }

    // Create the submatrix
    arma::cx_mat sub_matrix(this->N, this->N, arma::fill::zeros);
    sub_matrix.diag(-1).fill(r);
    sub_matrix.diag(1).fill(r);

    // Set the size of B
    this->B.set_size(this->N * this->N, this->N * this->N);

    // Fill in the values in the submatrix diagonal
    for (size_t i = 0; i < this->B.n_cols; i += this->N) {
        this->B.submat(i, i, i + this->N - 1, i + this->N - 1) = sub_matrix;
    }

    // Fill the last sub/sup-diagonals
    this->B.diag() = diagonal;
    this->B.diag(-this->N).fill(r);
    this->B.diag(this->N).fill(r);
}

void WaveSimulation::initialize_U(double x_c, double y_c, double sigma_x,
                                  double sigma_y, double p_x, double p_y)
{
    this->U.set_size(this->N, this->N);
    double x, y, diff_x, diff_y;
    std::complex<double> sum = 0.;
    for (size_t j = 0; j < this->U.n_cols; j++) {
        x = j * h;
        diff_x = x - x_c;
        for (size_t i = 0; i < this->U.n_rows; i++) {
            y = i * h;
            diff_y = y - y_c;
            this->U(i, j) =
                std::exp(-(diff_x * diff_x) / (2. * sigma_x * sigma_x)
                         - (diff_y * diff_y) / (2. * sigma_y * sigma_y)
                         + p_x * x * 1._i + p_y * y * 1._i);
            sum += this->U(i, j) * std::conj(this->U(i, j));
        }
    }
    if (std::abs(sum.imag()) > 1e-7) {
        abort();
    }
    double norm = 1. / std::sqrt(sum.real());

    this->U.for_each([norm](std::complex<double> &el) { el *= norm; });
}

void WaveSimulation::write_U(std::ofstream &ofile)
{
    this->U.for_each(
        [&ofile](std::complex<double> el) { ofile << el << '\t'; });
    ofile << '\n';
}

void WaveSimulation::build_V(double thickness, double pos_x,
                             double aperture_separation, double aperture,
                             uint32_t slits)
{
    this->V.set_size(this->N, this->N);
    this->V.fill(0.);
    if (slits == 0) {
        return;
    }
    arma::cx_vec res;

    if (slits % 2 == 0) {
        res = arma::cx_vec(aperture_separation / this->h,
                           arma::fill::value(1e10));
        for (size_t i = 0; i < slits; i += 2) {
            res = arma::join_cols(
                res, arma::cx_vec(aperture / this->h, arma::fill::zeros));
            res = arma::join_cols(
                arma::cx_vec(aperture / this->h, arma::fill::zeros), res);
            res =
                arma::join_cols(res, arma::cx_vec(aperture_separation / this->h,
                                                  arma::fill::value(1e10)));
            res = arma::join_cols(arma::cx_vec(aperture_separation / this->h,
                                               arma::fill::value(1e10)),
                                  res);
        }
    }
    else {
        res = arma::cx_vec(aperture / this->h, arma::fill::value(0));
        for (size_t i = 0; i < slits - 1; i += 2) {
            res =
                arma::join_cols(res, arma::cx_vec(aperture_separation / this->h,
                                                  arma::fill::value(1e10)));
            res = arma::join_cols(arma::cx_vec(aperture_separation / this->h,
                                               arma::fill::value(1e10)),
                                  res);
            res = arma::join_cols(
                res, arma::cx_vec(aperture / this->h, arma::fill::zeros));
            res = arma::join_cols(
                arma::cx_vec(aperture / this->h, arma::fill::zeros), res);
        }
    }
    if (res.size() > this->N) {
        abort();
    }
    uint32_t fill = (this->N - res.size()) / 2;
    res = arma::join_cols(arma::cx_vec(fill, arma::fill::value(1e10)), res);
    res = arma::join_cols(res, arma::cx_vec(fill + ((this->N - res.size()) % 2),
                                            arma::fill::value(1e10)));

    uint32_t start = pos_x / this->h - thickness / this->h / 2;
    for (size_t i = 0; i < thickness / this->h; i++) {
        this->V.col(start + i) = res;
    }
}