// License declaration at end of file

#include <iostream>
#include <fstream>
#include <iomanip>
#include <string>
#include <sys/stat.h>

#include "channelflow/flowfield.h"
#include "channelflow/vector.h"
#include "channelflow/chebyshev.h"
#include "channelflow/tausolver.h"
#include "channelflow/dns.h"
#include "channelflow/utilfuncs.h"

using namespace std;
using namespace channelflow;

void printdiagnostics(FlowField& u, const DNS& dns, Real t, const TimeStep& dt, Real nu, Real umin,
		      bool vardt, bool pl2norm, bool pchnorm, bool pdissip, bool pforcing,
		      bool pdiverge, bool pUbulk, bool pubulk, bool pdPdx, bool channel, bool pcfl);

int main(int argc, char* argv[]) {

  string purpose("integrate plane Couette or channel flow from a given "
		 "initial condition and save velocity fields to disk.");

  ArgList args(argc, argv, purpose);

  const Real T0       = args.getreal("-T0", "--T0", 0.0, "start time");
  const Real T1       = args.getreal("-T1", "--T1", 100, "end time");
  const bool vardt    = args.getflag("-vdt", "--variabledt", "adjust dt to keep CFLmin<=CFL<CFLmax");
  const Real dtarg    = args.getreal("-dt",     "--dt", 0.03125, "timestep");
  const Real dtmin    = args.getreal("-dtmin",  "--dtmin", 0.001, "minimum time step");
  const Real dtmax    = args.getreal("-dtmax",  "--dtmax", 0.05,  "maximum time step");
  const Real dT       = args.getreal("-dT",     "--dT", 1.0, "save interval");
  const Real CFLmin   = args.getreal("-CFLmin", "--CFLmin", 0.40, "minimum CFL number");
  const Real CFLmax   = args.getreal("-CFLmax", "--CFLmax", 0.60, "maximum CFL number");
  const string stepstr= args.getstr ("-ts", "--timestepping", "sbdf3", "timestepping algorithm, one of "
				     "[cnfe1 cnab2 smrk2 sbdf1 sbdf2 sbdf3 sbdf4]");
  const string nonlstr= args.getstr ("-nl", "--nonlinearity", "rot", "method of calculating nonlinearity, "
				     "one of [rot conv div skew alt]");
  const string outdir = args.getpath("-o", "--outdir", "data/", "output directory");
  const string label  = args.getstr ("-l", "--label", "u", "output field prefix");
  const Real Reynolds = args.getreal("-R", "--Reynolds", 400, "Reynolds number");
  const bool channel  = args.getflag("-c", "--channel", "channelflow instead of plane Couette");
  const bool bulkvel  = args.getflag("-b", "--bulkvelocity","hold bulk velocity fixed, not pressure gradient");
  const Real dPdx     = args.getreal("-P", "--dPdx",   0.0, "value for fixed pressure gradient");
  const Real Ubulk    = args.getreal("-U", "--Ubulk",  0.0, "value for fixed bulk velocity");
  const bool pcfl     = args.getflag("-cfl", "--cfl",         "print CFL number each dT");
  const bool pl2norm  = args.getflag("-l2",  "--l2norm",      "print L2Norm(u) each dT");
  const bool pchnorm  = args.getbool("-ch", "--chebyNorm", true, "print chebyNorm(u) each dT");
  const bool pdissip  = args.getflag("-D",  "--dissipation",  "print dissipation each dT");
  const bool pforcing = args.getflag("-I",  "--input",        "print power input each dT");
  const bool pdiverge = args.getflag("-dv", "--divergence",   "print divergence each dT");
  const bool pubulk   = args.getflag("-u",  "--ubulk",        "print ubulk each dT");
  const bool pUbulk   = args.getflag("-Up", "--Ubulk-print",  "print Ubulk each dT");
  const bool pdPdx    = args.getflag("-p",  "--pressure",     "print pressure gradient each dT");
  const Real umin     = args.getreal("-u",  "--umin",   0.0,  "stop if chebyNorm(u) < umin");
  const string uname  = args.getstr (1, "<flowfield>", "initial condition");

  args.check();
  args.save("./");
  args.save(outdir);
  mkdir(outdir);

  cout << "Constructing u,q, and optimizing FFTW..." << endl;
  FlowField u(uname);
  u.optimizeFFTW(FFTW_PATIENT);

  const int Nx = u.Nx();
  const int Ny = u.Ny();
  const int Nz = u.Nz();
  const Real Lx = u.Lx();
  const Real Lz = u.Lz();
  const Real a = u.a();
  const Real b = u.b();

  FlowField q(Nx,Ny,Nz,1,Lx,Lz,a,b);
  TimeStep dt(dtarg, dtmin, dtmax, dT, CFLmin, CFLmax, vardt);
  const bool inttime = (abs(dT - int(dT)) < 1e-12) ? true : false;
  const Real nu = 1.0/Reynolds;

  // Construct Navier-Stoke Integrator
  DNSFlags flags;
  flags.timestepping = s2stepmethod(stepstr);
  flags.dealiasing   = DealiasXZ;
  flags.nonlinearity = s2nonlmethod(nonlstr);
  flags.constraint   = bulkvel ? BulkVelocity : PressureGradient;
  flags.baseflow     = channel ? Parabolic : PlaneCouette;
  flags.dPdx         = dPdx;
  flags.Ubulk        = Ubulk;

  cout << "constructing DNS..." << endl;
  DNS dns(u, nu, dt, flags, Real(T0));

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

    printdiagnostics(u, dns, t, dt, nu, umin, vardt, pl2norm, pchnorm, pdissip,
		     pforcing, pdiverge, pUbulk, pubulk, pdPdx, channel, pcfl);

    u.binarySave(outdir + label + t2s(t, inttime));

    dns.advance(u, q, dt.n());

    if (dt.adjust(dns.CFL()))
      dns.reset_dt(dt);
  }
  cout << "done!" << endl;
}

void printdiagnostics(FlowField& u, const DNS& dns, Real t, const TimeStep& dt, Real nu,
		      Real umin, bool vardt, bool pl2norm, bool pchnorm, bool pdissip, bool pforcing,
		      bool pdiverge, bool pUbulk, bool pubulk, bool pdPdx, bool channel, bool pcfl) {

    // Printing diagnostics
    cout << "           t == " << t << endl;
    if (vardt)    cout << "          dt == " << Real(dt) << endl;
    if (pl2norm)  cout << "   L2Norm(u) == " << L2Norm(u) << endl;

    if (pchnorm || umin !=0.0) {
      Real chnorm = chebyNorm(u);
      cout << "chebyNorm(u) == " << chnorm << endl;
      if (chnorm < umin) {
	cout << "Exiting: chebyNorm(u) < umin." << endl;
	exit(0);
      }
    }
    if (pdissip)  {
      Complex U(1,0);
      u.cmplx(0,1,0,0) += U;
      cout << " dissip(u+U) == " << dissipation(u) << endl;
      u.cmplx(0,1,0,0) -= U;
    }
    if (pforcing) cout << "  input(u+U) == " << 1 + forcing(u) << endl;
    if (pdiverge) cout << "  divNorm(u) == " << divNorm(u) << endl;
    if (pUbulk)   cout << "       Ubulk == " << dns.Ubulk() << endl;
    if (pubulk)   cout << "       ubulk == " << Re(u.profile(0,0,0)).mean() << endl;
    if (pdPdx)    cout << "        dPdx == " << dns.dPdx() << endl;
    if (channel) {
      cout << "  Ubulk*h/nu == " << dns.Ubulk()/nu << endl;
      cout << " Uparab*h/nu == " << 1.5*dns.Ubulk()/nu << endl;
    }
    cout << "         CFL == " << dns.CFL() << endl;
}

/* couette.cpp: time-integration class for spectral Navier-Stokes simulator
 * channelflow-1.1 PCF-utils
 *
 * Copyright (C) 2001-2007  John F. Gibson
 *
 * gibson@cns.physics.gatech.edu  jfg@member.fsf.org
 *
 * Center for Nonlinear Science
 * School of Physics
 * Georgia Institute of Technology
 * Atlanta, GA 30332-0430
 * 404 385 2509
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation version 2
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, U
 */