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

using namespace std;
using namespace channelflow;

int main() {

  cout << "================================================================\n";
  cout << "This program integrates a plane Couette flow from a random\n";
  cout << "initial condition. Velocity fields are saved at intervals dT=1.0\n";
  cout << "in a data-couette/ directory, in channelflow's binary data file\n";
  cout << "format." << endl << endl;

  // Define gridsize
  const int Nx=24;
  const int Ny=33;
  const int Nz=48;

  // Define box size
  const Real Lx=1.75*pi;
  const Real a= -1.0;
  const Real b=  1.0;
  const Real Lz=1.2*pi;

  // Define flow parameters
  const Real Reynolds = 400.0;
  const Real nu = 1.0/Reynolds;
  const Real dPdx = 0.0;
  const Real Ubulk = 0.0;

  // Define integration parameters
  const int n = 25;      // take fifty steps between printouts
  const Real dt = 1.0/n;  // integration timestep
  const Real T  = 100.0; // integrate from t=0 to t=T

  // Define DNS parameters
  DNSFlags flags;
  flags.baseflow     = PlaneCouette;
  flags.timestepping = SBDF3;
  flags.initstepping = SMRK2;
  flags.nonlinearity = Rotational;
  flags.dealiasing   = DealiasXZ;
  flags.taucorrection = true;
  flags.constraint  = PressureGradient; // enforce constant pressure gradient
  flags.dPdx = dPdx;
  flags.Ubulk = Ubulk;

  // Define size and smoothness of initial disturbance
  Real spectralDecay = 0.5;
  Real magnitude  = 0.3;
  int kxmax = 3;
  int kzmax = 3;

  // Construct data fields: 3d velocity and 1d pressure
  cout << "building velocity and pressure fields..." << flush;
  FlowField u(Nx,Ny,Nz,3,Lx,Lz,a,b);
  FlowField q(Nx,Ny,Nz,1,Lx,Lz,a,b);
  cout << "done" << endl;

  // Perturb velocity field
  u.addPerturbations(kxmax,kzmax,1.0,spectralDecay);
  u *= magnitude/L2Norm(u);

  // Construct Navier-Stoke integrator, set integration method
  cout << "building DNS..." << flush;
  DNS dns(u, nu, dt, flags);
  cout << "done" << endl;

  mkdir("data-couette");
  for (Real t=0; t<=T; t += n*dt) {
    cout << "         t == " << t << endl;
    cout << "       CFL == " << dns.CFL() << endl;
    cout << " L2Norm(u) == " << L2Norm(u) << endl;
    cout << "divNorm(u) == " << divNorm(u) << endl;
    cout << "      dPdx == " << dns.dPdx() << endl;
    cout << "     Ubulk == " << dns.Ubulk() << endl;

    // Write velocity and modified pressure fields to disk
    u.binarySave("data-couette/u"+i2s(int(t)));
    q.binarySave("data-couette/q"+i2s(int(t)));

    // Take n steps of length dt
    dns.advance(u, q, n);
    cout << endl;
  }
}