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

using namespace std;
using namespace channelflow;

int main() {

  cout << "================================================================\n";
  cout << "This program integrates a plane Poiseuille from a random\n";
  cout << "initial condition and computes a number of turbulent statistics\n";
  cout << "The mean velocity profile can be plotted against the logarithmic\n";
  cout << "\"law of the wall\" with the matlab script walllaw.m\n";
  cout << endl << endl;

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

  const Real Lx = 4*pi/3; // 4.1888
  const Real Lz = 2*pi/3; // 2.0724

  //const Real Lx = 2*pi;   // 6.2832
  //const Real Lz =   pi;   // 6.2832
  const Real a= -1.0;
  const Real b=  1.0;
  const Real bsub=  -0.8;

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

  const Real CFLmin = 0.20;
  const Real CFLmax = 0.40;
  const Real dtmin  = 0.0001;
  const Real dtmax  = 0.01;
  const Real T0 = 0.0;   // integration start time
  const Real T1 = 100;   // grow turbulence from perturbations: T0 < t <T1
  const Real T2 = 200;   // take statistics: T1 <t <T2
  const Real dT = 1.0;   // plotting and statistics interval

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

  const Real perturbMag = 0.10;
  const Real decay = 0.7; // spectral decay of perturbation profiles
  const int kxmax=3;     // maximum Fourier mode for perturbations
  const int kzmax=3;

  const char sp= ' ';

  cout << setprecision(6);
  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);

  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.0,decay);
  u *= perturbMag/L2Norm(u);

  cout << "div(u)/L2Norm(u)  == " << divNorm(u)/L2Norm(u) << endl;

  FlowField tmp(Nx,Ny,Nz,6,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, nu, dt, flags, T0);
  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;
  ******************/

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

  TurbStats stats(Ubase, nu);

  mkdir("data-walllaw");
  for (Real t=T0; t<=T2; t += dT) {

    if (dt.adjust(dns.CFL())) {
      cerr << "Resetting dt to " << dt << ", new CFL == " << dt.CFL() << endl;
      dns.reset_dt(dt);
    }
    cout << "          t == " << t << endl;
    cout << "         dt == " << dt << endl;
    cout << "        CFL == " << dns.CFL() << endl;
    cout << " L2Norm2(u) == " << L2Norm2(u) << endl;
    cout << "divNorm2(u) == " << divNorm2(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) % 40 == 0 && t != T0) {
      cout << "saving flowfields..." << endl;
      u.binarySave("data-walllaw/u"+i2s(int(t)));
      //q.binarySave("q"+i2s(int(t)));
      cout << "done" << endl;
    }

    ChebyCoeff u00 = Re(u.profile(0,0,0));
    u00.makePhysical();
    u00.save("u00");

    /******************************
    u.makePhysical();
    u.saveSlice(1,0,Ny/2, "uslice");
    u.saveSlice(1,1,Ny/2, "vslice");
    u.saveSlice(1,2,Ny/2, "wslice");
    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();
    u.saveSpectrum("uspec");

    for (int kx=-kxmax; kx<=kxmax; ++kx) {
      const int mx = u.mx(kx);
      for (int kz=0; kz<=kzmax; ++kz) {
	const 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 (t >= T1) {
      stats.addData(u,tmp);
      cout << "centerline Re = " << stats.centerlineReynolds() << endl;
      cout << " parabolic Re = " << stats.parabolicReynolds() << endl;
      cout << "      bulk Re = " << stats.bulkReynolds() << endl;
      cout << "        ustar = " << stats.ustar() << endl;
      cout << "        bsub+ = " << stats.ustar()/nu * (bsub-a) << endl;

      stats.msave("uu");
      stats.msave("uustar", true);

      Real ustar = stats.ustar();
      Vector yp = stats.yplus();
      yp.save("yp");

      ChebyCoeff Umean = stats.U();
      Umean.makeSpectral(trans);
      ChebyCoeff Umeany = diff(Umean);
      Umean.makePhysical(trans);
      Umeany.makePhysical(trans);
      Umean.save("Umean");
      Umeany.save("Umeany");

      Umean /= ustar;
      Umean.save("Uplus");

      ChebyCoeff ubase = stats.ubase();
      ubase.save("ubase");

      ChebyCoeff uv = stats.uv();
      uv.save("uv");
      save(ustar, "ustar");
      save(nu, "nu");
    }
    dns.advance(u, q, dt.n());
    cout << endl;

  }
  cout << "done!" << endl;
}