// Script to simulate tunnelling in 1-D
// Isaac Lenton
// August 2017

#include <fstream>
#include <vector>
#include <complex>
#include <cmath>
#include <fftw3.h>      // For Fourier transforms, link with -lfftw3
using std::complex;

const double barrierWidth = 0.12;           // Width of barrier
const double barrierHeight = 2.1;           // Height of barrier
const unsigned animLength = 1000;           // Number of output frames (files)

const double xmax = 10.0;                   // Half width of the simulation
const double dt = 1.0e-2;                   // Time step size
const unsigned N = 512;                     // Simulation spatial resolution

// Implementation of Matlab's fftshift function
void fftshift(std::vector<complex<double>>& vec) {
  std::vector<complex<double>>::iterator mid = vec.begin() + vec.size()/2;
  std::swap_ranges(vec.begin(), mid, mid);
}

int main(int argc, char** argv) {

  std::vector<complex<double>> kprop(N);
  std::vector<complex<double>> xprop(N);
  std::vector<complex<double>> wavefunction(N);
  std::vector<complex<double>> pwavefunction(N);
  std::complex<double> cJ(0.0, 1.0);

  // Create a plan for the Fourier transforms
  fftw_plan forward, backward;
  forward = fftw_plan_dft_1d(N, (double(*)[2]) wavefunction.data(),
      (double(*)[2]) pwavefunction.data(), FFTW_FORWARD, FFTW_ESTIMATE);
  backward = fftw_plan_dft_1d(N, (double(*)[2]) pwavefunction.data(),
      (double(*)[2]) wavefunction.data(), FFTW_BACKWARD, FFTW_ESTIMATE);

  std::ofstream fpp("potential.dat");
  fpp << "# x\tPotential\n";

  for (unsigned i = 0; i < N; ++i) {
    double x = 2*i*xmax/N - xmax;
    double p = 2*M_PI*i/xmax - N*M_PI/xmax;

    double potential = 0.25*x*x + barrierHeight*exp(-pow(x/barrierWidth, 2.0));

    // Calculate the propagators
    kprop[i] = exp(-cJ*dt*p*p)/pow(N, 0.5);
    xprop[i] = exp(-cJ*dt*potential)/pow(N, 0.5);

    // Calculate initial wave-function
    wavefunction[i] = sqrt(sqrt(1.0/(2.0*M_PI)))*exp(-0.25*pow(x+4.0, 2.0));

    // Write the potential to file
    fpp << x << '\t' << potential << '\n';
  }
  fpp.close();
  fftshift(kprop);

  for (unsigned i = 0; i < animLength; ++i) {
    // Evolve the wave-function
    fftw_execute(forward);
    for (unsigned k = 0; k < N; ++k) pwavefunction[k] *= kprop[k];
    fftw_execute(backward);
    for (unsigned k = 0; k < N; ++k) wavefunction[k] *= xprop[k];

    // Write the wave-function to file
    char buf[128];
    snprintf(buf, 128, "output_%06d.dat", i);
    std::ofstream fp(buf);
    fp << "# x\t|Wavefunction|^2\n";
    for (unsigned k = 0; k < N; ++k) {
      double x = 2*k*xmax/N - xmax;
      fp << x << '\t' << pow(abs(wavefunction[k]), 2.0) << '\n';
    }
  }

  // Cleanup the plans
  fftw_destroy_plan(forward);
  fftw_destroy_plan(backward);

  return 0;
}