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


// Estimate Poincare section of unstable manifold of an equilibrium u*
// Poincare condition is (u-u*, e) = 0

// Integrate flow f^t(u) with u(0) = u* + eps v,
// Save datapoints when (f^t(u)-u*, e) = 0
//

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

  string purpose("compute crossings of a Poincare section in a pre-computed trajectory");

  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 Real dT       = args.getreal("-dT",     "--dT", 1.0, "save interval");
  const string udir   = args.getpath("-d", "--datadir", "data/", "flowfield series directory");
  const string odir   = args.getpath("-o", "--poindir", "poincare/", "output dir for Poincare crossings");
  const string tag    = args.getstr("-tag", "--tag", "", "tag for saved fields: u0${tag}, u1${tag}, etc");

  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 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");
  const string nonlstr= args.getstr ("-nl", "--nonlinearity", "rot", "method of calculating nonlinearity");
  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 Real theta    = args.getreal("-th", "--theta",  0.0, "poincare condition (u(t)-u*, e0 cos theta + e1 sin theta) = 0");
  const string etxname = args.getstr(5, "<flowfield>", "a taux-symmetric basis function");
  const string etzname = args.getstr(4, "<flowfield>", "a tauz-symmetric basis function");
  const string e0name  = args.getstr(3, "<flowfield>", "e0, in plane of complex oscillation");
  const string e1name  = args.getstr(2, "<flowfield>", "e1, in plane of complex oscillation ");
  const string eqbname = args.getstr(1, "<flowfield>", "fixed pt of poincare map");

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

  cout << "Constructing u,q, and optimizing FFTW..." << endl;
  FlowField ustar(eqbname);
  FlowField etx(etxname);
  FlowField etz(etzname);
  FlowField e1(e1name);
  FlowField e0(e0name);

  e0 *= cos(theta)/L2Norm(e0);
  e1 *= sin(theta)/L2Norm(e1);
  FlowField epoincare = e0;
  epoincare += e1;

  FieldSymmetry taux(1,1,1,0.5,0);
  FieldSymmetry tauz(1,1,1,0,0.5);
  FieldSymmetry tauxz(1,1,1,0.5,0.5);

  cout << "Checking that etx and etz are antisymmetric in taux and tauz" << endl;
  cout << "L2IP(etx, taux*etx)/L2Norm2(etx) == " << L2IP(etx, taux*etx)/L2Norm2(etx) << endl;
  cout << "L2IP(etz, taux*etz)/L2Norm2(etz) == " << L2IP(etz, tauz*etz)/L2Norm2(etz) << endl;

  TimeStep dt(dtarg, dtmin, dtmax, dT, CFLmin, CFLmax, vardt);
  const bool inttime = (abs(dT - int(dT)) < 1e-12) && (abs(T0 - int(T0)) < 1e-12) ? true : false;
  const Real nu = 1.0/Reynolds;

  // Construct time-stepping algorithm
  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;
  flags.verbosity    = Silent;

  // Poincare condition is c == (u-u*, e) == 0 or (u,e) - (u*,e) == 0.

  string uname = udir + "u" + t2s(T0, inttime);
  FlowField uprev(uname);
  FieldSymmetry tau; // identity

  // Move uprev into 1st quadrant: (uprev, etx)=>0 && (uprev, etz)=>0
  if (L2IP(uprev, etx) <0) {
    uprev *= taux;
    tau *= taux;
  }
  if (L2IP(uprev, etz) <0) {
    uprev *= taux;
    tau *= tauz;
  }

  Real c1star = L2IP(ustar, epoincare);
  Real c1prev = L2IP(uprev, epoincare) - c1star;

  int crossing=0;  // count number of crossings
  char tab = '\t';

  ofstream os((odir + "crossingtimes-" + tag + ".asc").c_str());
  os << setprecision(17);

  cout << t2s(T0,inttime) << tab << c1prev << endl;

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

    FlowField u(udir + "u" + t2s(t, inttime));

    // Apply same translation as applied to uprev
    u *= tau;

    // Check if u has moved into new quadrant, if so, translate both
    if (L2IP(u, etx) <0) {
      cout << "Crossed x quadrant; translating everything by taux..." << endl;
      cout << "Before     (u, etx) == " << L2IP(u,etx) << endl;
      u     *= taux;
      uprev *= taux;
      tau   *= taux;
      c1prev = L2IP(uprev, epoincare) - c1star;
      cout << "After (taux u, etx) == " << L2IP(u,etx) << endl;
   }
    if (L2IP(u, etz) <0) {
      cout << "Crossed z quadrant; translating everything by tauz..." << endl;
      cout << "Before     (u, etz) == " << L2IP(u,etz) << endl;
      u     *= tauz;
      uprev *= tauz;
      tau   *= tauz;
      c1prev = L2IP(uprev, epoincare) - c1star;
      cout << "After (tauz u, etz) == " << L2IP(u,etz) << endl;
    }

    Real c1 = L2IP(u, epoincare) - c1star;

    cout << t2s(t,inttime) << tab << c1 << endl;

    // If we cross the Poincare section, back up and reintegrate
    // checking condition every integration time step dt.
    if ((c1prev<=0 && 0<c1) || (c1prev>=0 && 0>c1)) {

      bool increasing = (c1prev<=0 && 0<c1) ? true : false;

      cout << "====================================" << endl;
      cout << "found a crossing..." << endl;

      // Allocate arrays to store time, velocity, and Poincare conditions
      // at three consecutive time steps for quadratic interpolation. E.g.
      array<Real> s(3);      // s[n] = t - dT + n dt
      array<Real> d(3);      // d[n] = c1(s[n])
      array<FlowField> v(3); // v[n] = v(s[n])

      s[0] = t-dT;
      s[1] = 0.0;
      s[2] = 0.0;
      v[0] = uprev;
      d[0] = c1prev;
      d[1] = 0.0;
      d[2] = 0.0;


      cout << "d:";
      for (int i=0; i<3; ++i)
	cout << d[i] << tab;
      cout << endl;
      cout << "s:";
      for (int i=0; i<3; ++i)
	cout << s[i] << tab;
      cout << endl;


      FlowField q(uprev.Nx(), uprev.Ny(), uprev.Nz(), 1,
		  uprev.Lx(), uprev.Lz(), uprev.a(), uprev.b());

      cout << "constructing DNS..." << endl;

      DNS dns(v[0], nu, dt, flags, t);
      cout << "finished DNS..." << endl;

      int count=1; // need three data points for quadratic interpolation

      for (; s[0] <= t+dt; ) {

	//cout << "time  shifts..." << endl;
	// Time-shift veclocity and Poincare condition arrays
	// v[2] <- v[1] <- v[0], same w d
	for (int n=2; n>0; --n) {
	  s[n] = s[n-1];
	  v[n] = v[n-1];
	  d[n] = d[n-1];
	}
	//cout << "time step..." << endl;
	dns.advance(v[0], q, 1); // take one step of length dt
	d[0] = L2IP(v[0], epoincare) - c1star;
	s[0] += dt;

	//cout << "d:";
	//for (int i=0; i<3; ++i)
	//cout << d[i] << tab;
	//cout << endl;
	//cout << "s:";
	//for (int i=0; i<3; ++i)
	//cout << s[i] << tab;
	//cout << endl;

	cout << s[0] << tab << d[0] << endl;

	//cout << "crossing check..." << endl;

	// Check for Poincare section crossing
	if (++count >= 3 && ((d[2]<=0 && 0<d[0]) || (d[2]>=0 && 0>d[0]))) {

	  cout << "Found a crossing. Interpolating ..." << endl;
	  // Interpolate v as a function of d at value d==0

	  FlowField upoincare = quadraticInterpolate(v, d, 0);
	  upoincare.makeSpectral();
	  if (increasing)
	    upoincare.binarySave(odir + "uP" + i2s(crossing++) + tag);
	  else
	    upoincare.binarySave(odir + "uM" + i2s(crossing) + tag);


	  // output time of crossing
	  Real crosstime = quadraticInterpolate(s, d, 0);
	  os << crosstime << endl;

	  cout << "Estimated poincare crossing: " << endl;
	  cout << "        time == " << crosstime << endl;
	  Real cross = L2IP(upoincare, epoincare) - c1star;
	  cout << "(u-ustar,e) == " << cross << endl;

	  //cfpause();
	  break;
	}
      }
    }
    uprev = u;
    c1prev = c1;
  }
  cout << "done!" << endl;
}

/* poincare.cpp: compute crossings of Poincare section
 * channelflow-1.3
 *
 * Copyright (C) 2001-2009  John F. Gibson
 *
 * gibson@cns.physics.gatech.edu
 *
 * 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
 */