/** @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>
#include <vector>

// Initializers
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+1) * h;
        diff_x = x - x_c;
        for (size_t i = 0; i < this->U.n_rows; i++) {
            y = (i+1) * 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));
        }
    }

    // The imaginary part of sum should be 0.
    if (std::abs(sum.imag()) > 1e-7) {
        abort();
    }

    double norm = 1. / std::sqrt(sum.real());

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

void WaveSimulation::initialize_V()
{
    this->V.set_size(this->N, this->N);
    this->V.fill(0.);
}

void WaveSimulation::initialize_V(double thickness, double pos_x, double ap_sep,
                                  double ap, uint32_t slits)
{
    this->V.set_size(this->N, this->N);
    this->V.fill(0.);
    if (slits == 0) {
        return;
    }

    arma::cx_vec res;

    // Differentiate between odd and even number of slits.
    uint32_t ap_points, ap_sep_points;

    if (slits % 2 == 0) {
        ap_points = ap / this->h;
        ap_sep_points = ap_sep / this->h
                        + ((uint32_t)(ap_sep / this->h) % 2 != this->N % 2);

        res = arma::cx_vec(ap_sep_points, arma::fill::value(1e10));

        for (size_t i = 0; i < slits; i += 2) {
            res = arma::join_cols(res,
                                  arma::cx_vec(ap_points, arma::fill::zeros));
            res = arma::join_cols(arma::cx_vec(ap_points, arma::fill::zeros),
                                  res);
            res = arma::join_cols(
                res, arma::cx_vec(ap_sep_points, arma::fill::value(1e10)));
            res = arma::join_cols(
                arma::cx_vec(ap_sep_points, arma::fill::value(1e10)), res);
        }
    }
    else {
        ap_points =
            ap / this->h + ((uint32_t)(ap / this->h) % 2 != this->N % 2);
        ap_sep_points = ap_sep / this->h;

        res = arma::cx_vec(ap_points, arma::fill::value(0));

        for (size_t i = 0; i < slits - 1; i += 2) {
            res = arma::join_cols(
                res, arma::cx_vec(ap_sep_points, arma::fill::value(1e10)));
            res = arma::join_cols(
                arma::cx_vec(ap_sep_points, arma::fill::value(1e10)), res);
            res = arma::join_cols(res,
                                  arma::cx_vec(ap_points, arma::fill::zeros));
            res = arma::join_cols(arma::cx_vec(ap_points, 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)));

    // Subtract one to compensate for the border being gone
    uint32_t start = (pos_x / this->h - 1) - thickness / this->h / 2;
    for (size_t i = 0; i < thickness / this->h; i++) {
        this->V.col(start + i) = res;
    }
}

void WaveSimulation::initialize_A()
{
    DEBUG("Inside initialize_A");
    // Create the diagonal
    arma::cx_vec diagonal(this->N * this->N);

    DEBUG("Before setting diagonal values");
    // 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);
    }

    DEBUG("Create submatrix");
    // 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);

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

    DEBUG("Fill in values");
    // 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;
    }

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

void WaveSimulation::initialize_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);
}

// Constructors
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->initialize_V(thickness, pos_x, ap_sep, ap, slits);
    this->initialize_U(x_c, y_c, sigma_x, sigma_y, p_x, p_y);
    this->initialize_A();
    this->initialize_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)
{
    DEBUG("Inside constructor");
    this->dt = dt;
    this->h = h;
    this->T = T;
    this->M = 1. / h;
    this->N = M - 2;
    DEBUG("Before V");
    this->initialize_V();
    DEBUG("Before U");
    this->initialize_U(x_c, y_c, sigma_x, sigma_y, p_x, p_y);
    DEBUG("Before A");
    this->initialize_A();
    DEBUG("Before B");
    this->initialize_B();
    DEBUG("After B");
}

// Public methods
void WaveSimulation::step()
{
    // Create the right hand side of Ax = b
    arma::cx_vec tmp = this->B * this->U.as_col();

    // Solve Ax = b
    arma::spsolve(this->U, this->A, tmp);
}

void WaveSimulation::simulate(std::string outfile, bool write_each_step)
{
    // Create path and proceed if successful.
    if (!utils::mkpath(utils::dirname(outfile))) {
        abort();
    }

    // Open file
    std::ofstream ofile;
    ofile.open(outfile);

    // Write the size of the matrix on the first line
    ofile << this->N << '\n';

    uint32_t iterations = this->T / this->dt;

    if (write_each_step) {
        for (size_t i = 0; i < iterations; i++) {
            this->write_U(ofile);
            this->step();
        }
    }
    else {
        for (size_t i = 0; i < iterations; i++) {
            this->step();
        }
    }
    this->write_U(ofile);
    ofile.close();
}
void WaveSimulation::simulate(std::string outfile, std::vector<double> &steps)
{
    // Create path and proceed if successful.
    if (!utils::mkpath(utils::dirname(outfile))) {
        abort();
    }

    // Open file
    std::ofstream ofile;
    ofile.open(outfile);

    // Write the size of the matrix on the first line
    ofile << this->N << '\n';
    
    uint32_t iterations;

    for (size_t i=0; i < steps.size(); i++) {
        if (i == 0) {
            iterations = steps[i] / this->dt;
        }
        else {
            iterations = (steps[i] - steps[i-1]) / this->dt;
        }

        for (size_t j=0; j < iterations; j++) {
            this->step();
        }
        this->write_U(ofile);
    }

    ofile.close();
}

void WaveSimulation::probability_deviation(std::string outfile,
                                           bool write_each_step)
{
    // Create path and proceed if successful.
    if (!utils::mkpath(utils::dirname(outfile))) {
        abort();
    }

    // Open file
    std::ofstream ofile;
    ofile.open(outfile);

    uint32_t iterations = this->T / this->dt;

    double sum;

    if (write_each_step) {
        for (size_t i = 0; i < iterations; i++) {
            sum = arma::accu(this->U % arma::conj(this->U)).real();
            ofile << i*this->dt << '\t' << std::abs(1. - sum) << '\n';
            this->step();
        }
    }
    else {
        for (size_t i = 0; i < iterations; i++) {
            this->step();
        }
    }
    sum = arma::accu(this->U % arma::conj(this->U)).real();
    ofile << this->T << '\t' << std::abs(1. - sum) << '\n';
    ofile.close();
}

void WaveSimulation::write_U(std::ofstream &ofile)
{
    // Write each element to file in column-major order.
    this->U.for_each(
        [&ofile](std::complex<double> el) { ofile << el << '\t'; });
    ofile << '\n';
}