#include <iostream>
#include <fstream>
#include <sstream>
#include <iomanip>
#include "channelflow/flowfield.h"
#include "channelflow/dns.h"

//#include "helperfuncs.h"

using namespace std;

int main() {

  const int Nx=64;
  const int Ny=65;
  const int Nz=32;

  const Real Lx=4*pi/3;
  const Real a= -1.0;
  const Real b=  1.0;
  const Real Lz=2*pi/3;

  const Real Reynolds = 4000.0;
  const Real nu = 1.0/Reynolds;

  const Real CFLmin = 0.50;
  const Real CFLmax = 0.70;
  const Real dtmin  = 0.002;
  const Real dtmax  = 0.02;
  const Real T0 = 0;    // start time
  const Real T1 = 100;  // end time
  const Real dT = 1.0;  // plotting and statistics interval

  DNSFlags flags;
  flags.timestepping = SBDF3;
  flags.dealiasing   = DealiasXZ;

  const Real perturbMag = 0.10;
  const Real decay = 0.7; // cheb spectral decay of perturb profiles
  const int kxmax=3;      // maximum Fourier mode for perturbations
  const int kzmax=3;
 
  const char sp= ' ';
  const char nl= '\n';

  cout << setprecision(4);
  Vector x = periodicpoints(Nx, Lx);
  Vector y = chebypoints(Ny,a,b);
  Vector z = periodicpoints(Nz, Lz);
  x.save("x");
  y.save("y");
  z.save("z");
  
  ChebyTransform trans(Ny);

  ChebyCoeff Ubase(Ny,a,b,Physical);
  for (int ny=0; ny<Ny; ++ny) 
    Ubase[ny] = 1.0 - square(y[ny]);
  Ubase.save("Ubase");
  Ubase.makeSpectral(trans);

  FlowField u(Nx,Ny,Nz,3,Lx,Lz,a,b);
  FlowField q(Nx,Ny,Nz,1,Lx,Lz,a,b);
  u.addPerturbations(kxmax,kzmax,1,decay);
  u *= perturbMag/L2Norm(u);
  //FlowField u("u80");
  //FlowField q("q80");
  
  FlowField omega(Nx,Ny,Nz,3,Lx,Lz,a,b);

  cout << "optimizing FFTW..." << flush;
  fftw_loadwisdom();
  u.optimizeFFTW();
  fftw_savewisdom();
  cout << "done" << endl;

  cout << "constructing DNS..." << flush;
  TimeStep dt((dtmin+dtmax)/2, dtmin, dtmax, dT, CFLmin, CFLmax); 
  DNS dns(u, Ubase, nu, dt, flags, T0);
  dns.reset_Ubulk(2.0/3.0);
  cout << "done" << endl;
  
  // print header info in fmodes.asc file

  ofstream modestream("fmodes.asc");
  ofstream dragstream("drags.asc");

  modestream << "% ";
  for (int kx=-kxmax; kx<=kxmax; ++kx) 
    for (int kz=0; kz<=kzmax; ++kz) 
      modestream << kx << ',' << kz << sp;
  modestream << nl;
  
  for (Real t=T0; t<=T1; t += dT) {
    cout << "          t == " << t << endl;
    cout << "         dt == " << dt << endl;
    cout << "        CFL == " << dns.CFL() << endl;
    cout << " L2Norm2(u) == " << L2Norm2(u) << endl;
    cout << "divNorm2(u) == " << divNorm(u)/L2Norm(u) << endl;  
    cout << "      Ubulk == " << dns.Ubulk() << endl;
    cout << "      ubulk == " << Re(u.profile(0,0,0)).mean()/2 << endl;
    cout << "       dPdx == " << dns.dPdx() << endl;
    

    if (int(t) % 10 == 0 && t != T0) {
      cout << "saving flowfields..." << endl;
      u.binarySave("u"+i2s(int(t)));
      q.binarySave("q"+i2s(int(t)));
      cout << "done" << endl;
    }

    u.makePhysical();
    u.saveSlice(2,0,0,"uside");
    u.saveSlice(2,1,0,"vside");
    u.saveSlice(2,2,0,"wside");
    u.saveSlice(0,0,0,"usec");
    u.saveSlice(0,1,0,"vsec");
    u.saveSlice(0,2,0,"wsec");
    u.makeSpectral();

    curl(u, omega);
    omega.saveSlice(0,0,0,"omsec0");
    omega.saveSlice(0,1,0,"omsec1");
    omega.saveSlice(0,2,0,"omsec2");
    for (int kx=-kxmax; kx<=kxmax; ++kx) {
      int mx = u.mx(kx);
      for (int kz=0; kz<=kzmax; ++kz) {
	int mz = u.mz(kz);
	BasisFunc u_kxkz = u.profile(mx,mz);
	modestream << L2Norm(u_kxkz) << sp;
      }
    }
    modestream << endl;
    ChebyCoeff dudy = diff(Re(u.profile(0,0,0)));
    dragstream << nu*dudy.eval_a() << sp << nu*dudy.eval_b() << endl;

    if (dt.adjust(dns.CFL())) {
      cerr << "Resetting dt to " << dt << ", new CFL == " << dt.CFL() << endl;
      dns.reset_dt(dt);
    }
    dns.advance(u, q, dt.n());
    cout << endl;
  }
  cout << "done!" << endl;
}