#include <sys/stat.h>
#include <fstream>
#include <iomanip>
#include <string>
#include <iostream>
#include <octave/oct.h>
#include <octave/dColVector.h>
#include <octave/dMatrix.h>
#include <octave/EIG.h>
#include "channelflow/flowfield.h"
#include "channelflow/vector.h"
#include "channelflow/chebyshev.h"
#include "channelflow/tausolver.h"
#include "channelflow/dns.h"
#include "channelflow/symmetry.h"

#include "channelflow/quicksort.h"
#include "octutils.h"

using namespace std;
using namespace channelflow;

// This program calculates eigenvalues of fixed point of plane Couette flow
// using Arnoldi iteration. The ideas and algorithm are based on Divakar
// Viswanath, "Recurrent motions within plane Couette turbulence",
// J. Fluid Mech.</em> 580 (2007), http://arxiv.org/abs/physics/0604062.
// Notation and details of the algorithm are based on Lecture 33 of
// "Numerical Linear Algebra" by Trefethen and Bau. Production runs began
// in early December 2006.

// Arnoldi iteration estimates the eigenvalues of a matrix A by iteratively
// constructing a QR decomposition of a matrix whose columns are
// Ab, A^2 b, etc., where b is an arbitrary starting vector.

// For the present problem, A is a discretized form of PCF dynamics,
// linearized about the equilibrium solutions, and b is the discretized
// form of an arbitrary starting velocity field.

// In continuous terms, the matrix A corresponds to a differential operator
// L and b to a perturbation du from the PCF equilibrium u*. Let F^T be the
// time-T Navier-Stokes map f^T : u(t) -> u(t+T), and let u* be an
// equilibrium solution of f^T; u* = f^T(u*). The define L as
// L du = f^T(ustar + du) - f^T(ustar) =  f^T(ustar + du) - ustar
//
// This program computes L du, L^2 du , L^3 du, etc. with a DNS algorithm
// and translates between velocity fields du^n = L^n du and vectors
// v^n = A^n b via an orthonormal basis set {Phi_i} for the space of velocity
// fields: v_i = (du, Phi_i), du = v_i Phi_i

// See my research blog circa 2006-12-02 for a discussion of why it's better
// to compute eigenvalues of the map u(t+T) = f^T(u(t)) than of the flow
// du/dt = f(u).

// f^T(u) = time-T DNS integration of u,

// Last modified: John F. Gibson, Wed Oct 10 16:43:13 EDT 2007

void f(const FlowField& u, Real T, FlowField& f_u, Real Reynolds, DNSFlags& flags,
       TimeStep& dt, int& fcount);

void Df(const FlowField& u, const FlowField& fu, const FlowField& du, FlowField& Df_du,
	Real T, const FieldSymmetry& tau, Real Reynolds, DNSFlags& flags,
	TimeStep& dt, Real eps, bool centerdiff, int& fcount);

void checkConjugacy(const ComplexColumnVector& u, const ComplexColumnVector& v);

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

  string purpose("computes eigenvalues of equilibrium or orbit with Arnoldi iteration");
  ArgList args(argc, argv, purpose);

  const Real EPS_arn  = args.getreal("-ek", "--epsKrylov", 1e-10, "min. condition # of Krylov vectors");
  const Real EPS_du   = args.getreal("-edu","--epsdu", 1e-7, "magnitude of Arnoldi perturbation");
  const bool centdiff = args.getflag("-cd", "--centerdiff", "centered differencing to estimate differentials");
  const int  seed     = args.getint ("-sd",  "--seed",       1, "seed for random number generator");
  const Real smooth   = args.getreal("-s",  "--smoothness", 0.4, "smoothness of initial perturb, 0 < s < 1");
  const int  Narnoldi = args.getint ("-Na", "--Narnoldi", 200, "number of Arnoldi iterations");
  const int  Nsave    = args.getint ("-Ne", "--Neigval", 30, "number of eigenvalues/functions to save");
  const bool l2basis  = args.getflag("-b",  "--l2basis", "use an explicit L2 basis");
  const bool s1       = args.getflag("-s1", "--s1", "restrict to s1-symmetric subspace");
  const bool s2       = args.getflag("-s2", "--s2", "restrict to s2-symmetric subspace");
  const bool s3       = args.getflag("-s3", "--s3", "restrict to s3-symmetric subspace");

  const bool flow     = args.getflag("-fl","--flow", "compute eigenvalues of flow rather than map");
  const Real T        = args.getreal("-T",  "--maptime", 10.0, "integration time for Navier-Stokes.");
  const bool sx       = args.getflag("-sx", "--x-sign", "change u,x sign after time integration");
  const bool sy       = args.getflag("-sy", "--y-sign", "change v,y sign");
  const bool sz       = args.getflag("-sz", "--z-sign", "change w,z sign");
  const Real ax       = args.getreal("-ax", "--ax-shift", 0.0, "translate by ax*Lx");
  const Real az       = args.getreal("-az", "--az-shift", 0.0, "translate by az*Lz");
  const bool anti     = args.getflag("-a",  "--anti",     "apply antisymmetry rather than symmetry");

  const string duname = args.getstr ("-du", "--perturb", "", "initial perturbation field, random if unset");
  const string outdir = args.getpath("-o", "--outdir", "./", "output directory");
  const Real Reynolds = args.getreal("-R", "--Reynolds", 400, "Reynolds number");
  const bool vardt    = args.getflag("-vdt","--variableDt", "adjust dt to keep CFLmin<=CFL<CFLmax");
  const string nonlstr= args.getstr ("-nl", "--nonl", "rotational", "method for u dot grad u calculation. See README for options.");
  const string stepstr= args.getstr ("-ts", "--timestepping", "sbdf3", "timestepping algorithm. See README for options.");
  const Real dtarg    = args.getreal("-dt","--timestep",  0.03125,   "(initial) integration timestep");
  const Real dtmin    = args.getreal("-dtmin", "--dtmin", 0.001,  "minimum timestep");
  const Real dtmax    = args.getreal("-dtmax", "--dtmax", 0.04,   "maximum timestep");
  const Real dTCFL    = args.getreal("-dTCFL", "--dTCFL", 1.0,    "interval between CFL checks");
  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 uname  = args.getstr(1, "<flowfield>", "the equilibrium field");

  args.check();
  args.save(outdir);

  srand48(seed);
  //const Real nu = 1.0/Reynolds;
  const Real decay = 1.0-smooth;
  const int prec = 16;

  // Define DNS algorithm
  DNSFlags flags;
  flags.baseflow     = PlaneCouette;
  flags.timestepping = s2stepmethod(stepstr);
  flags.dealiasing   = DealiasXZ;
  flags.nonlinearity = s2nonlmethod(nonlstr);
  flags.constraint   = PressureGradient;
  flags.dPdx         = 0.0;
  flags.verbosity    = Silent;
  cout << "flags == " << flags << endl;

  //BasisFlags bflags;
  //bflags.aBC = Diri;
  //bflags.bBC = Diri;
  //bflags.zerodivergence = true;
  //bflags.orthonormalize = true;

  int fcount = 0;

  FlowField u(uname);
  const int Nx = u.Nx();
  const int Ny = u.Ny();
  const int Nz = u.Nz();

  const Real Lx = u.Lx();
  const Real ya = u.a();
  const Real yb = u.b();
  const Real Lz = u.Lz();

  const int kxmin = -u.kxmaxDealiased();
  const int kxmax =  u.kxmaxDealiased();
  const int kzmin =  0;
  const int kzmax =  u.kzmaxDealiased();

  cout << setprecision(17);
  cout << "   Nx == " << Nx << endl;
  cout << "   Ny == " << Ny << endl;
  cout << "   Nz == " << Nz << endl;
  cout << "kxmin == " << kxmin << endl;
  cout << "kxmax == " << kxmax << endl;
  cout << "kzmin == " << kzmin << endl;
  cout << "kzmax == " << kzmax << endl;
  //cout << "building basis (this takes a long time)..." << endl;

  // Used to communicate timestep parameters to integration functions.
  cout << "dt     == " << dtarg << endl;
  cout << "dtmin  == " << dtmin << endl;
  cout << "dtmax  == " << dtmax << endl;
  cout << "CFLmin == " << CFLmin << endl;
  cout << "CFLmax == " << CFLmax << endl;
  TimeStep dt(dtarg, dtmin, dtmax, dTCFL, CFLmin, CFLmax, vardt);
  dt.adjust_for_T(T);

  FieldSymmetry tau(sx,sy,sz,ax,az,anti);

  // Construct real-valued ON div-free dirichlet basis set
  //vector<RealProfile> basis
  //= realBasis(Ny, kxmax, kzmax, Lx, Lz, ya, yb, bflags);

  //int M = basis.size(); // dimension of linear space
  //cout << M << "-dimensional state space." << endl;

  // Set up DNS operator ("A" in Arnoldi A*b terms)
  cout << "setting up DNS and initial fields..." << endl;
  ChebyTransform trans(Ny);

  const Real eps = EPS_du/L2Norm(u);
  if (flags.nonlinearity == LinearAboutField)
    u.setToZero();
  FlowField fu(u);
  cout << "computing tau f(u,T)..." << endl;
  f(u, T, fu, Reynolds, flags, dt, fcount);
  fu *= tau;


  /*************************
  FlowField dudt = fu;
  dudt -= u;
  dudt *= 1.0/T;
  cout << "L2Norm(u*)        == " << L2Norm(u) << endl;
  cout << "L2Norm(f(u*))     == " << L2Norm(fu) << endl;
  cout << "L2Dist(u*,f(u*))  == " << L2Dist(u,fu) << endl;
  cout << "L2Norm(du* /dt)    == " << L2Norm(dudt) << endl;
  *************************/

  // Set random initial function ("b" in Arnoldi A*b terms)
  FlowField u_du(Nx, Ny, Nz, 3, Lx, Lz, ya, yb);    // u == u(t)
  FlowField fu_du(Nx, Ny, Nz, 3, Lx, Lz, ya, yb);   // u(t+T) == f^T(u(t))
  FlowField Df_du(Nx, Ny, Nz, 3, Lx, Lz, ya, yb);   // Df du == f(u*+du) - f(u*)

  // Set du  = EPS_du (random unit perturbation))
  FlowField du(Nx, Ny, Nz, 3, Lx, Lz, ya, yb);

  du = u;
  du -= fu;
  cout << "L2Norm(u - tau f^T(u)) == " << L2Norm(du) << endl;
  du.setToZero();

  if (duname.length() == 0) {
    du.addPerturbations(kxmax, kzmax, 1.0, decay);
    du.addPerturbation(0, 0, 1.0, decay);
  }
  else
    du = FlowField(duname);

  du *= EPS_du/L2Norm(du);
  cout << "Post-rescaling..." << endl;

  /************************************
  if (indir != "") {
    for (int n=0; n<Nsave; ++n)
      du += FlowField(indir + "ef" + i2s(n));
  }
  **************************************/

  vector<RealProfileNG> basis;
  if (l2basis || s1 || s2 || s3) {
    cout << "constructing basis set (takes a long time)..." << endl;
    basis = realBasisNG( Ny, kxmax, kzmax, Lx, Lz, ya, yb);

    vector<FieldSymmetry> symms;
    if (s1) symms.push_back(FieldSymmetry( 1, 1,-1, 0.5, 0.0));
    if (s2) symms.push_back(FieldSymmetry(-1,-1, 1, 0.5, 0.5));
    if (s3) symms.push_back(FieldSymmetry(-1,-1,-1, 0.0, 0.5));
    if (s1 || s2 || s3)
      selectSymmetries(basis, symms);
    orthonormalize(basis);
  }

  // Project field du onto basis, producing vector b
  ColumnVector b; //(M);
  field2vector(du,b, basis);
  //field2vector(du,b);
  b *= 1.0/L2Norm(b);

  Arnoldi arnoldi(b, Narnoldi, EPS_arn);

  // Initialize output streams
  ofstream osLamconv((outdir + "Lamconv.asc").c_str());
  ofstream oslamconv;
  if (flow) {
    oslamconv.open((outdir + "lamconv.asc").c_str());
    oslamconv << setprecision(17);
  }

  for (int n=0; n<Narnoldi; ++n) {
    cout << "======================================================" << endl;
    cout << n << "th iteration:" << endl;

    // Compute v = Aq in Arnoldi iteration terms, where q test vector
    // In fluid terms, this is
    //      L du == F^T(u + du) - F^T(u)
    // where |du| << |u|, so that RHS expression is approximately linear
    // and where f^T is the forward map for time T of Navier-Stokes

    // Relation of v and du :
    //   du == eps * v_m phi_m
    //    v == 1/eps (du, phi_m)

    const ColumnVector& q = arnoldi.testVector();

    cout << "building field du..." << endl;
    vector2field(q, du, basis);

    // Compute Df du == 1/eps (f(u + eps du) - f(u)) or centerdiff equiv
    Df(u, fu, du, Df_du, T, tau, Reynolds, flags, dt, eps, centdiff, fcount);

    cout << "computing vector..." << endl;
    ColumnVector Aq; //(M);
    field2vector(Df_du, Aq, basis);

    // Send the product Aq back to Arnoldi
    arnoldi.iterate(Aq);

    // Now do some output: Lambda, the eigenvalues of map f^T,
    ComplexColumnVector Lambda = arnoldi.ew();
    ComplexMatrix Vn = arnoldi.ev();
    int nsave=lesser(Nsave, n+1);

    cout << "\nabs, arg Lambda == " << endl;
    for (int j=0; j<nsave; ++j) {
      cout << setw(prec+5) << abs(Lambda(j)) << setw(prec+5) << arg(Lambda(j)) << endl;
      osLamconv << Re(Lambda(j)) << ' ' << Im(Lambda(j)) << ' ';
    }
    for (int j=nsave; j<Nsave; ++j)
      osLamconv << "-99 0 ";
    osLamconv << endl;

    ofstream osLam((outdir + "Lambda.asc").c_str());
    osLam << setprecision(17);
    for (int j=0; j<n; ++j)
      osLam     << Re(Lambda(j)) << ' ' << Im(Lambda(j)) << endl;
    osLam.close();

    ofstream osLam2((outdir + "LambdaAbsArg.asc").c_str());
    osLam2 << setprecision(17);
    for (int j=0; j<n; ++j)
      osLam2     << abs(Lambda(j)) << ' ' << arg(Lambda(j)) << endl;
    osLam2.close();

    // Convert eigvals of map F^T (Lambda) to eigvals of flow f (lambda)
    // Lambda == e^{lambda T}.
    // lambda == 1/T log Lambda
    //        == 1/T (log(abs(Lambda)) + i arg(Lambda))
    //
    if (flow) {
      ComplexColumnVector lambda(n+1);
      for (int j=0; j<=n; ++j)
	lambda(j) = (1.0/T) * log(Lambda(j));

      cout << "\nlambda == " << endl;
      for (int j=0; j<nsave; ++j) {
	cout      << Re(lambda(j)) << ' ' << Im(lambda(j)) << endl;
	oslamconv << Re(lambda(j)) << ' ' << Im(lambda(j)) << ' ';
      }
      // For file lamconv.asc, supplement computed eigvals with bogus values
      // -99 + 0 i so that all rows have same number of columns
      for (int j=nsave; j<Nsave; ++j)
	oslamconv << "-99 0 ";
      oslamconv << endl;

      // Save all estimated eigenvalues
      ofstream oslam((outdir + "lambda.asc").c_str());
      oslam << setprecision(17);
      for (int j=0; j<n; ++j)
      oslam     << Re(lambda(j)) << ' ' << Im(lambda(j)) << endl;
      oslam.close();
    }

    // Print the period and char. mult of first oscillatory mdoe
    /*****************************************************
    for (int j=0; j<nsave; ++j) {
      if (abs(Im(lambda(j)) > 0.0)) {
	Real period  = abs(2*pi/Im(lambda(j)));
	Real magnify = exp(Re(lambda(j)*abs(period)));
	cout << "period  = " << period << endl;
	cout << "magnify = " << magnify << endl;
	break;
      }
    }
    *************************************************/

    // Reconstruct and save top ten eigenfunctions
    for (int j=0; j<nsave; ++j) {

      string sj1 = i2s(j+1);
      string sj2 = i2s(j+2);

      // Real-valued eigenvalue
      if (Im(Lambda(j)) == 0.0) {
	FlowField ef(u); // set ef to have same size as u
	ColumnVector ev = real(Vn.column(j));
	//vector2field(ev, ef);
	vector2field(ev, ef, basis);
	ef *= 1.0/L2Norm(ef);
	ef.binarySave(outdir + "ef" + sj1);
	save(Lambda(j), outdir + "eW" + sj1);
	if (flow) {
	  Complex lambda = log(Lambda(j)) / T;
	  save(lambda, outdir + "ew" + sj1);
	}
      }
      // Complex eigenvalue pair
      else if (j<nsave-1) {
	Complex LambdaA = Lambda(j);
	Complex LambdaB = Lambda(j+1);
	if (LambdaA != conj(LambdaB)) {
	  cerr << "Warning! Non-conjugate complex eigenvalues!" << endl;
	  cerr << "abs, arg Lambda(" << sj1 << ") == "
	       << setw(prec+5) << abs(LambdaA) << ' '
	       << setw(prec+5) << arg(LambdaA) << endl;
	  cerr << "abs, arg Lambda(" << sj2 << ") == "
	       << setw(prec+5) << abs(LambdaB)
	       << setw(prec+5) << arg(LambdaB) << endl;
	}
	checkConjugacy(Vn.column(j), Vn.column(j+1));

	ColumnVector evA = real(Vn.column(j));
	ColumnVector evB = imag(Vn.column(j));
	//evA += real(Vn.column(j+1));
	//evB -= imag(Vn.column(j+1));

	FlowField efA(u); // set ef to have same size as u
	FlowField efB(u); // set ef to have same size as u
	vector2field(evA, efA, basis);
	vector2field(evB, efB, basis);
	Real c = 1.0/sqrt(L2Norm2(efA) + L2Norm2(efB));
	efA *= c;
	efB *= c;
	efA.binarySave(outdir + "ef" + sj1);
	efB.binarySave(outdir + "ef" + sj2);
	save(LambdaA, outdir + "eW" + sj1);
	save(LambdaB, outdir + "eW" + sj2);

	if (flow) {
	  save(log(LambdaA) / T, outdir + "ew" + sj1);
	  save(log(LambdaB) / T, outdir + "ew" + sj2);
	}
	++j;

      }
    }
  }
}

// f^T(u) = time-T DNS integration of u
void f(const FlowField& u, Real T, FlowField& f_u, Real Reynolds, DNSFlags& flags,
       TimeStep& dt, int& fcount) {

  if (T<0) {
    cerr << "f: negative integration time T == " << T << endl;
    cerr << "returning f(u,T) == (1+abs(T))*u" << endl;
    f_u = u;
    f_u *= 1+abs(T);
    return;
  }

  // Special case #1: no time steps
  if (T == 0) {
    cout << "f: T==0, no integration, returning" << endl;
    f_u = u;
    return;
  }

  //cout << "f(u,T) : L2Norm(u) == " << L2Norm(u) << ",  T == " << T << flush;

  const Real nu = 1.0/Reynolds;
  FlowField p(u.Nx(), u.Ny(), u.Nz(), 1, u.Lx(), u.Lz(), u.a(), u.b());
  f_u = u;

  // Adjust dt for CFL if necessary
  DNS dns(f_u, nu, dt, flags);
  if (dt.variable()) {
    dns.advance(f_u, p, 1);
    dt.adjust(dns.CFL());
    dns.reset_dt(dt);
    f_u = u;
  }

  //cout << ", CFL == " << dns.CFL() << endl;
  //cout << "t == " << flush;

  //  t == current time in integration
  //  T == total integration time
  // dT == CFL-check interval
  // dt == DNS time-step
  //  N == T/dT, n == dT/dt;
  //  T == N dT, dT == n dt
  //  t == s dT (s is loop index)
  cout << "f(u,t) : t == " << flush;
  for (int s=1; s<=dt.N(); ++s) {
    dns.advance(f_u, p, dt.n());
    Real t = s*dt.dT();
    if (s % 10 == 0)
      cout << iround(t) << flush;
    else if (s % 2 == 0)
      cout << '.' << flush;
    if (dt.adjust(dns.CFL()))
      dns.reset_dt(dt);
  }
  cout << "  CFL == " << dns.CFL() << endl;
  ++fcount;
}

void Df(const FlowField& u, const FlowField& fu, const FlowField& du, FlowField& Df_du,
	Real T, const FieldSymmetry& tau, Real Reynolds, DNSFlags& flags,
	TimeStep& dt, Real eps, bool centerdiff, int& fcount) {

  FlowField u_du(du);
  FlowField fu_du(du);

  if (!centerdiff) {

    cout << "Computing f^T(u + du) for approximation L du == f^T(u + du) - f^T(u)" << endl;
    u_du  = du;
    u_du *= eps;
    u_du += u;

    fu_du = u_du;
    f(u_du, T, fu_du, Reynolds, flags, dt, fcount);
    fu_du *= tau;

    Df_du = fu_du;
    Df_du -= fu;

    cout << "          L2Norm(u_du) == " << L2Norm(u_du) << endl;
    cout << "            L2Norm(du) == " << L2Norm(du) << endl;
    cout << "L2Norm(f(u))           == " << L2Norm(fu) << endl;
    cout << "L2Norm(f(u + du))      == " << L2Norm(fu_du) << endl;
    cout << "L2Norm(f(u + du)-f(u)) == " << L2Norm(Df_du) << endl;

    Df_du *= 1.0/eps;
  }
  else {

    cout << "Computing f^T(u +/- du/2) for approximation L du == f^T(u + du/2) - f^T(u - du/2)" << endl;

    // Compute f^T(u+du)
    u_du  = du;
    u_du *= eps/2;
    u_du += u;

    fu_du = u_du;
    f(u_du, T, fu_du, Reynolds, flags, dt, fcount);
    fu_du *= tau;
    Df_du = fu_du;

    // Compute f^T(u-du)
    u_du  = du;
    u_du *= -eps/2;
    u_du += u;

    fu_du = u_du;
    f(u_du, T, fu_du, Reynolds, flags, dt, fcount);
    fu_du *= tau;
    Df_du -= fu_du;

    // Complete approximation L du == f^T(u + du/2) - f^T(u - du/2)
    Df_du *= 1.0/eps;
  }
}


void checkConjugacy(const ComplexColumnVector& u, const ComplexColumnVector& v) {
  ColumnVector f = real(u);
  ColumnVector g = real(v);
  Real rerr = L2Dist(f,g);
  f = imag(u);
  g = imag(v);
  g *= -1;
  Real ierr = L2Dist(f,g);
  if (rerr + ierr > 1e-14*(L2Norm(u)+L2Norm(v))) {
    cout << "error : nonconjugate u,v, saving these vectors" << endl;
    save(u, "uerr");
    save(v, "verr");
  }
}