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

//#include "helperfuncs.h"

using namespace std;
using namespace channelflow;

int main() {

  cout << "================================================================\n";
  cout << "This program integrates a plane Poiseuille (channel) flow from \n";
  cout << "a random initial initial condition and saves velocity fields\n";
  cout << "in a data-channel/ directory, in channelflow's binary data file\n";
  cour << "format." << endl << endl;
  const int Nx=128;
  const int Ny=65;
  const int Nz=128;

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

  const Real nu = 1.0/4000;

  const Real CFLmin = 0.40;
  const Real CFLmax = 0.60;
  const Real dtmin  = 0.002;
  const Real dtmax  = 0.04;
  const Real dt0    = 0.02;
  const bool variable_dt = true;
  const Real T0 = 0;    // start time
  const Real T1 = 300;  // end time
  const Real dT = 0.25;  // plotting and statistics interval

  DNSFlags flags;
  flags.baseflow     = Parabolic;
  flags.timestepping = SBDF3;
  flags.dealiasing   = DealiasXZ;
  flags.constraint   = BulkVelocity;
  flags.Ubulk        = 2.0/3.0;

  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");

  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(dt0, dtmin, dtmax, dT, CFLmin, CFLmax, variable_dt);
  DNS dns(u, 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;

  // prior to timestepping loop set Ubase(y) =  1-y^2 = [T_0(y) - T_2(y)]/2
  ChebyCoeff Ubase(Ny, a,b,Spectral);
  Ubase[0] = 0.5;
  Ubase[2] = -0.5;
  const Real h = (b-a)/2;
  const Real ycenter = (b+a)/2;

  mkdir("data-channel");
  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;

      // inside timestepping-loop set
    ChebyCoeff U = Re(u.profile(0,0,0));
    U += Ubase;

    Real Reynolds_center = U.eval(ycenter)*h/nu;
    Real Reynolds_mean   = U.mean()*h/nu;

    cout << " centerline Re  == " << Reynolds_center << endl;
    cout << " mean flow  Re  == " << Reynolds_mean << endl;

    const int it = iround(t);
    if (abs(Real(it)-t) < dt && it % 10 == 0) {
      cout << "saving flowfields..." << endl;
      u.binarySave("data-channel/u"+i2s(it));
      q.binarySave("data-channel/q"+i2s(it));
      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;
    ***************************/

    //ChebyCoeff u00 = Re(u.profile(0,0,0));
    //u00.save("u00_" + i2s(int(t)));

    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;
}