// dsmc3d - Program to simulate a dilute gas using DSMC algorithm
// This version simulates planar Couette flow

#include "NumMeth.h"
#include "SortList3d.h"
#include "SampList3d.h"

double rand( long& seed );
double randn( long& seed );
int colider3d( Matrix3d& v, Matrix3d crmax, double tau, long& seed,
              Matrix3d selxtra, double coeff, SortList3d& sD );
void sorter3d( Matrix3d& position, double L, SortList3d &sD );
void mover3d( Matrix3d& position, Matrix3d& v, int npart, double L,
            double mpv, double vwall, double tau,
            Matrix3d& strikes, Matrix3d& delv, long& seed );
void sampler3d( Matrix3d& position, Matrix3d& v, int npart, double L,
            SampList3d& sampD );

int main() {

  //* Initialize constants  (particle mass, diameter, etc.)
  const double pi = 3.141592654;
  const double boltz = 1.3806e-23;    // Boltzmann's constant (J/K)
  double mass = 6.63e-26;       // Mass of argon atom (kg)
  double diam = 3.66e-10;       // Effective diameter of argon atom (m)
  double T = 273;               // Temperature (K)
  double density = 2.685e25;    // Number density of argon at STP (m^-3)
  double L = 5.0e-6;              // System size is one micron
  double Volume = L*L*L;        // Volume of the system
  cout << "Enter number of simulation particles: ";
  int npart; cin >> npart;
  double eff_num = density*L*L*L/npart;
  cout << "Each particle represents " << eff_num << " atoms" << endl;
  double mfp = Volume/(sqrt(2.0)*pi*diam*diam*npart*eff_num);
  cout << "System width is " << L/mfp << " mean free paths" << endl;
  double mpv = sqrt(2*boltz*T/mass);  // Most probable initial velocity
//Not needed anymore:
//  cout << "Enter wall velocity as Mach number: ";
//  double vwall_m; cin >> vwall_m;
//  double vwall = vwall_m * sqrt(5./3. * boltz*T/mass);
//  cout << "Wall velocities are " << -vwall << " and "
//      << vwall << " m/s" << endl;
double vwall = 0;

  //* Assign random positions and velocities to particles
  long seed = 1;       // Initial seed for rand (DO NOT USE ZERO)
  //Matrix3d x(npart); //Replaced with a position matrix
  Matrix3d position(npart,3), v(npart,3);
  int i;
  for( i=1; i<=npart; i++ ) {
    //x(i) = L*rand(seed);        // Assign random positions
    position(i,1) = L*rand(seed);//x coord
    position(i,2) = L*rand(seed);//y coord
    position(i,3) = L*rand(seed);//z coord
    // Initial velocities are Maxwell-Boltzmann distributed
    v(i,1) = sqrt(boltz*T/mass) * randn(seed);
    v(i,2) = sqrt(boltz*T/mass) * randn(seed);
    v(i,3) = sqrt(boltz*T/mass) * randn(seed);
    //Not what we need anymore
    // Add velocity gradient to the y-component
    //v(i,2) += vwall * 2*(x(i)/L - 0.5);
  }

  //* Initialize variables used for evaluating collisions
  int ncell_x = 20;                       // Number of cells
  int ncell_y = 20;                       // Number of cells
  int ncell_z = 20;                       // Number of cells
  int ncell_total = ncell_x*ncell_y*ncell_z;                       // Number of cells
  double tau = 0.2*(L/ncell_total)/mpv;       // Set timestep tau
  //Matrix3d vrmax(ncell_total), selxtra(ncell_total);  //We can't use Matrix3d here, we need an actual 3d matrix for these.
  //vrmax.set(3*mpv);    // Estimated max rel. speed
  //selxtra.set(0.0);       // Used by routine "colider"
  Matrix3d vrmax(ncell_x,ncell_y,ncell_z);
  Matrix3d selxtra(ncell_x,ncell_y,ncell_z);

  vrmax.set(3*mpv);
  selxtra.set(0);

  double coeff = 0.5*eff_num*pi*diam*diam*tau/(L*L*L/ncell_total);

  //* Declare object for lists used in sorting
  SortList3d sortData(ncell_x, ncell_y, ncell_z, npart);

    //Ignoring this for now.
  // Initialize object and variables used in statistical sampling
  SampList3d sampData(ncell_x,ncell_y,ncell_z);
  double tsamp = 0;             // Total sampling time
  Matrix3d dvtot(2), dverr(2);
  dvtot.set(0.0);               // Total momentum change at a wall
  dverr.set(0.0);               // Used to find error in dvtot

  //* Loop for the desired number of time steps
  int colSum = 0;         // Count total collisions
  Matrix3d strikes(2), strikeSum(2);
  strikeSum.set(0.0);     // Count strikes on each wall
  cout << "Enter total number of time steps: ";
  int istep, nstep; cin >> nstep;
  int col;
  for( istep = 1; istep<=nstep; istep++ ) {

    //* Move all the particles
    Matrix3d delv(2);  delv.set(0.0);
    mover3d( position, v, npart, L, mpv, vwall,
           tau, strikes, delv, seed );
    strikeSum(1) += strikes(1);
    strikeSum(2) += strikes(2);

    //* Sort the particles into cells
    sorter3d(position,L,sortData);

    //* Evaluate collisions among the particles
    col = colider3d(v,vrmax,tau,seed,selxtra,coeff,sortData);
    colSum += col;  // Increment collision count

    //* After initial transient, accumulate statistical samples
    if(istep > nstep/10) {
      sampler3d(position,v,npart,L,sampData);
      // Cummulative velocity change for particles striking walls
      dvtot(1) += delv(1);          dvtot(2) += delv(2);
      dverr(1) += delv(1)*delv(1);  dverr(2) += delv(2)*delv(2);
      tsamp += tau;
    }

    //* Periodically display the current progress
    if( (istep%1000) < 1 )  {
      cout << "Done " << istep << " of " << nstep << " steps; " <<
             colSum << " collisions" << endl;
      cout << "Total wall strikes: " << strikeSum(1) << " (left) "
           << strikeSum(2) << " (right)" << endl;
    }
  }

  //* Normalize the accumulated statistics
  int nsamp = sampData.nsamp;

/*  //Ignore for now.
for( i=1; i<=ncell; i++ ) {
    sampData.ave_n[i] *= (eff_num/(Volume/ncell))/nsamp;
    sampData.ave_ux[i] /= nsamp;
    sampData.ave_uy[i] /= nsamp;
    sampData.ave_uz[i] /= nsamp;
    sampData.ave_T[i] *= mass/(3*boltz*nsamp);
  }
*/


  dverr(1) = dverr(1)/(nsamp-1) - (dvtot(1)/nsamp)*(dvtot(1)/nsamp);
  dverr(1) = sqrt(dverr(1)*nsamp);
  dverr(2) = dverr(2)/(nsamp-1) - (dvtot(2)/nsamp)*(dvtot(2)/nsamp);
  dverr(2) = sqrt(dverr(2)*nsamp);

  //* Compute viscosity from drag force on the walls
  Matrix3d force(2), ferr(2);
  force(1) = (eff_num*mass*dvtot(1))/(tsamp*L*L);
  force(2) = (eff_num*mass*dvtot(2))/(tsamp*L*L);
  ferr(1) = (eff_num*mass*dverr(1))/(tsamp*L*L);
  ferr(2) = (eff_num*mass*dverr(2))/(tsamp*L*L);
  cout << "Force per unit area is" << endl;
  cout << "Left wall:  " << force(1) << " +/- " << ferr(1) << endl;
  cout << "Right wall: " << force(2) << " +/- " << ferr(2) << endl;
  double vgrad = 2*vwall/L;    // Velocity gradient
  double visc = 0.5*(-force(1)+force(2))/vgrad;  // Average viscosity
  double viscerr = 0.5*(ferr(1)+ferr(2))/vgrad;  // Error
  cout << "Viscosity = " << visc << " +/- " << viscerr
       << "N s/m^2" << endl;
  double eta = 5.*pi/32.*mass*density*(2./sqrt(pi)*mpv)*mfp;
  cout << "Theoretical value of viscoisty is " << eta
       << "N s/m^2" << endl;

/*  //Ignore for now until we know what to do about this.
  // Print out the plotting variables:
  //    xcell, ave_n, ave_ux, ave_uy, ave_uz, ave_T
  ofstream xcellOut("xcell.txt"), ave_nOut("ave_n.txt"),
           ave_uxOut("ave_ux.txt"), ave_uyOut("ave_uy.txt"),
           ave_uzOut("ave_uz.txt"), ave_TOut("ave_T.txt");
  for( i=1; i<=ncell; i++ ) {
    xcellOut << (i-0.5)*L/ncell << endl;
    ave_nOut << sampData.ave_n[i] << endl;
    ave_uxOut << sampData.ave_ux[i] << endl;
    ave_uyOut << sampData.ave_uy[i] << endl;
    ave_uzOut << sampData.ave_uz[i] << endl;
    ave_TOut << sampData.ave_T[i] << endl;
  }
*/



}
/***** To plot in MATLAB; use the script below ********************
load xcell.txt;  load ave_n.txt;     load ave_ux.txt;
load ave_uy.txt; load ave_uz.txt;    load ave_T.txt;
figure(1); clf;
plot(xcell,ave_n); xlabel('position');  ylabel('Number density');
figure(2); clf;
plot(xcell,ave_ux,xcell,ave_uy,xcell,ave_uz);
xlabel('position');  ylabel('Velocities');
legend('x-component','y-component','z-component');
figure(3); clf;
plot(xcell,ave_T); xlabel('position');  ylabel('Temperature');
******************************************************************/