#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>

using namespace std;

//Data structures
vector<double> x;
vector<double> uold;
vector<double> unew;

//Variables that we'll need to use
int N;
double t, tfinal, dt, dx; 
double a,nu; //advection and diffusion parameters;

//For writing data to disk
void writeFields(){
  ofstream write;
  write.open("num.dat");
  for(int i = 0; i < N; i++){
    write << x[i] << " " << unew[i] << endl;
  }
  write.close();

  write.open("exact.dat");
  for(int i = 0; i < N; i++){
    write << x[i] << " " << exp(-t)*sin(x[i]-t) << endl;
  }
  write.close();
}

//Write a function that will update unew and uold at each timestep
void update(){
  
  double dxi = 1.0/dx; //Inverse of dx

  //First update the stencil

  //Interior points
  for(int i = 1; i < N-1; i++){
    unew[i] = uold[i] + dt*(
			    -(uold[i]-uold[i-1])*dxi + (uold[i+1] - 2.0*uold[i] + uold[i-1])*dxi*dxi
			    );
  }
  //End points
    unew[0]   = uold[0]   + dt*(
			    -(uold[0]-uold[N-1])*dxi + (uold[1] - 2.0*uold[0] + uold[N-1])*dxi*dxi
			    );
    unew[N-1] = uold[N-1] + dt*(
			    -(uold[N-1]-uold[N-2])*dxi + (uold[0] - 2.0*uold[N-1] + uold[N-2])*dxi*dxi
			    );
  
  
    //Now kick everything up one timestep
    for(int i = 0; i < N; i++){
      uold[i] = unew[i];
    }

}

int main(){

  cout << "Beginning finite difference code for 1D advection-diffusion equation with periodic BCs." << endl;
  
  cout << "Pre-processing: Initialize parameters and data structures." << endl;
  cout << "Input N:";
  cin >> N;
  dx = 2.0*M_PI/double(N);
  dt = min(0.5*dx,0.25*dx*dx);
  tfinal = 0.5;

  //Initialize data structures
  x.resize(N);
  uold.resize(N);
  unew.resize(N);
  for(int i = 0; i < N; i++){
    x[i] = double(i)*dx;
    uold[i] = sin(x[i]);
    unew[i] = 0.0;
  }
    
  cout << "Beginning timestepping." << endl;
  t = 0;
  while(t < tfinal){
    update();
    t += dt;
  }

  cout << "Timestepping complete - outputing final state for post-processing." << endl;
  writeFields();

  cout << "Compute error compared to exact solution." << endl;
  double error = 0.0;
  for(int i = 0; i < N; i++){
    error += abs(unew[i] - exp(-t)*sin(x[i]-t));
  }
  cout << "(dt,dx,error) = (" << dt << "," << dx << "," << error/double(N) << ")"<< endl; 

  cout << "Finished." << endl;

  return 0;
}