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

using namespace std;

//Data structures
vector<double> x;

vector<double> u;
vector<double> E;
vector<double> p;
vector<double> rho;

vector<vector <double> > Y_new;
vector<vector <double> > Y_old;

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

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

//Write a function to pick the upwind flux vector based off of u
vector<double> get_Flux(int i , int j){
  
  vector<double> F_uv(3);

  /////////////////////////////////////////////////////////////
  //Replace this with some code to implement our upwind flux//
  ///////////////////////////////////////////////////////////

  F_uv[0] = 0.0;
  F_uv[1] = 0.0;
  F_uv[2] = 0.0;

  /////////////////////////////////////////

  return F_uv;
}


//Write a function that will update Y_old and Y_new at each timestep, and then update the primitive variables
void update(){
  
  //Update the stencil

  
  //Interior points
  for(int i = 1; i < N-1; i++){
    vector<double> flux_R = get_Flux(i,i+1);
    vector<double> flux_L = get_Flux(i-1,i);
    
    //V*ddtu = -(F_+ - F_-)
    for(int a = 0; a < 3; a++){
      Y_new[i][a] = Y_old[i][a] - (dt/dx)*(flux_R[a]-flux_L[a]);
    }
  }
  //Left-most cell
  {
    vector<double> flux_R = get_Flux(0,1);
    vector<double> flux_L = get_Flux(N-1,0);
    
    //V*ddtu = -(F_+ - F_-)
    for(int a = 0; a < 3; a++){
      Y_new[0][a] = Y_old[0][a];
    }
  }
  //Right-most cell
  {
    vector<double> flux_R = get_Flux(N-1,0);
    vector<double> flux_L = get_Flux(N-2,N-1);
    
    //V*ddtu = -(F_+ - F_-)
    for(int a = 0; a < 3; a++){
      Y_new[N-1][a] = Y_old[N-1][a];
    }
  }
  
  //Now kick everything up one timestep
  for(int i = 0; i < N; i++){
    for(int a = 0; a < 3; a++){
      Y_old[i][a] = Y_new[i][a];
    }
  }

  //Update primitive variables from conservative variables
  for(int i = 0; i < N; i++){
    rho[i] = Y_new[i][0];
    u[i]   = Y_new[i][1]/rho[i];
    E[i]   = Y_new[i][2]/rho[i];
    p[i]   = rho[i]*(E[i] - 0.5*u[i]*u[i])*(gamma0 - 1.0);
  }

}

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 = 1.0/double(N);
  tfinal = 0.1;
  dt = 0.00001;
  gamma0 = 1.4;

  cout << "Initial condition CFL = " << sqrt(gamma0/0.125)*dt/dx << endl;

  //Initialize data structures
  x.resize(N);

  //Conserved variables and fluxes
  Y_new.resize(N);
  Y_old.resize(N);


  //Primitive variables
  u.resize(N);
  E.resize(N);
  p.resize(N);
  rho.resize(N);
  for(int i = 0; i < N; i++){
    Y_new[i].resize(3);
    Y_old[i].resize(3);
  }
  for(int i = 0; i < N; i++){
    x[i] = (0.5+double(i))*dx;
  }

  //Data structures initialized - now set up initial condition on primitive variables
  for(int i = 0; i < N; i++){
    if( x[i] < 0.5 ){
      rho[i] = 1.0;
      u[i]   = 0.0;
      p[i]   = 1.0;
    }
    else{
      rho[i] = 0.125;
      u[i]   = 0.0;
      p[i]   = 0.1;
    }

    //Evaluate energy from equation of state
    E[i] = p[i]/(rho[i]*(gamma0 - 1.0)) + 0.5*u[i]*u[i];
  }

  //Use initial condition on primitive variables to initialize conserved variables
  for(int i = 0; i < N; i++){
    //density = rho
    Y_old[i][0] = rho[i];
    //momentum = rho*U;
    Y_old[i][1] = rho[i]*u[i];
    //energy = rho*E
    Y_old[i][2] = rho[i]*E[i];
  }



  cout << "Beginning timestepping." << endl;
  t = 0;
  while(t < tfinal){
    update();
    t += dt;
  }
  
  cout << "Timestepping complete - outputing final state for post-processing." << endl;
  writeFields();


  cout << "Finished." << endl;

  return 0;
}