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

using namespace std;
using namespace channelflow;

int main() {

  cout << "================================================================\n";
  cout << "This program loads from disk a velocity field from a previous \n";
  cout << "integrates it with variable time-stepping, and computes some \n";
  cout << "statistics." << endl << endl;
  cout << setprecision(17);

  // Define flow parameters
  const Real Reynolds = 100.0;
  const Real nu = 1.0/Reynolds;

  // Define integration parameters
  const Real dtmin = 0.01;
  const Real dtmax = 0.06;
  const Real CFLmin = 0.40;
  const Real CFLmax = 0.60;
  const Real dT  = 1.0;    // plot interval
  const Real T0  = 100.0;   // start time
  const Real T1  = 200.0;   // stop time

  // Load velocity, modified pressure, and base flow from disk.
  FlowField u("data-couette/u100");
  FlowField q(u.Nx(), u.Ny(), u.Nz(), 1, u.Lx(), u.Lz(), u.a(), u.b());

  // Get y-domain information from velocity field.
  Real a = u.a();
  Real b = u.b();
  int Ny = u.Ny();

  // Set integration parameters.
  DNSFlags flags;
  flags.baseflow     = PlaneCouette;
  flags.timestepping = SBDF3;
  flags.initstepping = CNRK2;
  flags.nonlinearity = Rotational;
  flags.dealiasing = DealiasXZ;
  flags.constraint = PressureGradient; // enforce constant pressure gradient
  flags.dPdx = 0.0;
  flags.Ubulk = 0.0;

  // Construct a variable time step object. See below for usage.
  TimeStep dt((dtmin+dtmax)/2, dtmin, dtmax, dT, CFLmin, CFLmax);

  // Construct a DNS.
  DNS dns(u, nu, dt, flags, T0);

  // Construct Chebyshev expansion for computing mean kx,kz=0,0 Fourier profile
  ChebyCoeff u00mean(Ny,a,b,Spectral);
  Real dragmean = 0.0;
  int count = 0;

  // Time stepping loop

  for (Real t=T0; t<T1; t += dT) {

    // Get real part of kx=kz=0 Fourier component u00(y) and compute drag
    ChebyCoeff u00 = Re(u.profile(0,0,0));
    ChebyCoeff du00dy = diff(u00);
    Real drag = nu*(du00dy.eval_a() + du00dy.eval_b());

    u00mean += u00;
    dragmean += drag;
    ++count;

    // Save stuff
    string time = i2s(int(t));
    u00.save("uprofile00_"+time);
    Re(u.profile(1,2,0)).save("uprofile12_"+time);

    cout << "        t == " << t << endl;
    cout << "       dt == " << dt << endl;
    cout << "      CFL == " << dns.CFL() << endl;
    cout << "L2Norm(u) == " << L2Norm(u) << endl;
    cout << "     drag == " << drag << endl;

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

    // If CFL number is outside the range CFLmin,CFLmax, adjust dt within
    // limits dtmin,dtmax, and subject to dT/dt = n for some integer n.
    if (dt.adjust(dns.CFL())) {
      cout << "adjusting timestep" << endl;
      dns.reset_dt(dt);
    }
  }

  // Compute means
  dragmean /= count;
  u00mean /= count;

  // Fourier-transform u00mean
  u00mean.makePhysical();
  u00mean.save("u00mean");
  cout << "mean drag == " << dragmean << endl;
}