#include <fstream>
#include <iomanip>
#include <string>
#include <iostream>
#include "channelflow/flowfield.h"
#include "channelflow/diffops.h"
#include "channelflow/symmetry.h"
#include "channelflow/utilfuncs.h"

using namespace std;
using namespace channelflow;

Real energy(FlowField& u, int i, bool normalize=true);

// returns S,A, for symmetric,antisymmetric to eps accuracy
// returns s,a, for symmetric,antisymmetric to sqrt(eps) accuracy
char symmetric(const FieldSymmetry& s, const FlowField& u, Real& serr, Real& aerr, Real eps=1e-7);

Real project(const FlowField& u, const FieldSymmetry& s, int sign);

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

  string purpose("prints information about a FlowField");

  ArgList args(argc, argv, purpose);
  const bool geom     = args.getflag("-g",  "--geometry", "show geometrical properties");
  const bool mean     = args.getflag("-m",  "--mean",     "show mean properties");
  const bool norm     = args.getflag("-n",  "--norm",     "show norm properties");
  const bool spec     = args.getflag("-sp", "--spectral", "show spectral properties");
  const bool symm     = args.getflag("-sy", "--symmetry", "show symmetry properties");
  const bool dyn      = args.getflag("-d",  "--dynamic",  "show dynamical properties");
  const bool nrgy     = args.getflag("-e",  "--energy",   "show energy properties");

  const Real Reynolds = args.getreal("-R",  "--Reynolds", 400, "Reynolds number");
  const Real eps      = args.getreal("-eps", "--eps", 1e-7, "spectral noise floor");

  const Real dt       = args.getreal("-dt", "--dt", 0.03125, "integration time for du/dt estimate");

  const Real T        = args.getreal("-T", "--T", dt, "integration time for du/dt estimate");

  const string nlmeth = args.getstr("-nl", "--nonlinearity", "rot", "method for  u grad u [rot or skw]");

  const int  digits   = args.getint("-dg", "--digits", 6, "digits of output for reals");
  const string uname  = args.getstr(1, "<flowfield>", "input field");
  args.check();

  bool all = (!geom && !mean && !norm && !spec && !symm && !dyn && !nrgy) ? true : false;

  FlowField u(uname);
  u.makeSpectral();
  const Real unorm = L2Norm(u);
  cout << setprecision(digits);

  if (all || geom) {
    cout << "--------------Geometry------------------" << endl;
    cout << "Nx == " << u.Nx() << endl;
    cout << "Ny == " << u.Ny() << endl;
    cout << "Nz == " << u.Nz() << endl;
    cout << "Nd == " << u.Nd() << endl;
    cout << "Lx == " << u.Lx() << endl;
    cout << "Ly == " << u.Ly() << endl;
    cout << "Lz == " << u.Lz() << endl;
    cout << " a == " << u.a() << endl;
    cout << " b == " << u.b() << endl;
    cout << "lx == " << u.Lx()/(2*pi) << endl;
    cout << "lz == " << u.Lz()/(2*pi) << endl;
    cout << "alpha == " << 2*pi/u.Lx() << endl;
    cout << "gamma == " << 2*pi/u.Lz() << endl;
    cout << endl;
  }
  if (all || mean) {
    for (int i=0; i<u.Nd(); ++i) {
      cout << "mean u[" << i << "] == " <<  Re(u.profile(0,0,i)).mean() << endl;
    }
  }
  if (all || symm) {
    cout << "--------------Symmetry-------------------\n";
    FieldSymmetry s1( 1, 1,-1, 0.5, 0.0);
    FieldSymmetry s2(-1,-1, 1, 0.5, 0.5);
    FieldSymmetry s3(-1,-1,-1, 0.0, 0.5);
    FieldSymmetry r1( 1, 1,-1);
    FieldSymmetry r2(-1,-1, 1);
    FieldSymmetry tx( 1, 1, 1, 0.5, 0.0);
    FieldSymmetry tz( 1, 1, 1, 0.0, 0.5);
    FieldSymmetry txz(1, 1, 1, 0.5, 0.5);

    Real unorm2 = L2Norm2(u);
    cout.setf(ios::left);
    cout << "fractions of energy in symmetric/antisymmetric subspaces:" << endl;
    cout << "s1  symm: " << setw(12) << project(u, s1, 1)/unorm2;
    cout << "    anti: " << setw(12) << project(u, s1,-1)/unorm2 << endl;
    cout << "s2  symm: " << setw(12) << project(u, s2, 1)/unorm2;
    cout << "    anti: " << setw(12) << project(u, s2,-1)/unorm2 << endl;
    cout << "s3  symm: " << setw(12) << project(u, s3, 1)/unorm2;
    cout << "    anti: " << setw(12) << project(u, s3,-1)/unorm2 << endl;
    cout << "r1  symm: " << setw(12) << project(u, r1, 1)/unorm2;
    cout << "    anti: " << setw(12) << project(u, r1,-1)/unorm2 << endl;
    cout << "r2  symm: " << setw(12) << project(u, r2, 1)/unorm2;
    cout << "    anti: " << setw(12) << project(u, r2,-1)/unorm2 << endl;
    cout << "tx  symm: " << setw(12) << project(u, tx, 1)/unorm2;
    cout << "    anti: " << setw(12) << project(u, tx,-1)/unorm2 << endl;
    cout << "tz  symm: " << setw(12) << project(u, tz, 1)/unorm2;
    cout << "    anti: " << setw(12) << project(u, tz,-1)/unorm2 << endl;
    cout << "txz symm: " << setw(12) << project(u, txz, 1)/unorm2;
    cout << "    anti: " << setw(12) << project(u, txz,-1)/unorm2 << endl;
    cout << endl;
    cout.unsetf(ios::left);
  }

  if (all || norm) {
    cout << "--------------Norms---------------------" << endl;

    cout << " L2Norm(u)    == " <<  unorm << endl;
    for (int i=0; i<u.Nd(); ++i)
      cout << " L2Norm(u[" << i << "]) == " << sqrt(energy(u,i)) << endl;

    cout << " divNorm(u)   == " << divNorm(u) << endl;
    cout << "  bcNorm(u)   == " <<  bcNorm(u) << endl;
    cout << endl;
  }

  if (all || spec) {
    cout << "------------Spectrum--------------------" << endl;
    int kxmax=0;
    int kzmax=0;
    int kymax=0;
    cout << "u.kxmin() == " << u.kxmin() << endl;
    cout << "u.kxmax() == " << u.kxmax() << endl;
    cout << "u.kzmin() == " << u.kzmin() << endl;
    cout << "u.kzmax() == " << u.kzmax() << endl;
    cout << "u.kxminDealiased() == " << u.kxminDealiased() << endl;
    cout << "u.kxmaxDealiased() == " << u.kxmaxDealiased() << endl;
    cout << "u.kzminDealiased() == " << u.kzminDealiased() << endl;
    cout << "u.kzmaxDealiased() == " << u.kzmaxDealiased() << endl;

    for (int mx=0; mx<u.Mx(); ++mx) {
      const int kx = u.kx(mx);
      for (int mz=0; mz<u.Mz(); ++mz) {
	const int kz = u.kz(mz);
	BasisFunc prof = u.profile(mx,mz);
	if (L2Norm(prof) > eps) {
	  kxmax = Greater(kxmax, abs(kx));
	  kzmax = Greater(kzmax, abs(kz));
	}
	for (int ky=0; ky<u.Ny(); ++ky) {
	  Real sum = 0.0;
	  for (int i=0; i<u.Nd(); ++i)
	    sum += abs(prof[i][ky]);
	  if (sum > eps)
	    kymax = Greater(kymax, ky);
	}
      }
    }
    cout << "u.dealiased() == " << (u.dealiased() ? "true" : "false") << endl;
    cout << "Energy over " << eps << " is confined to : \n";
    cout << " |kx| <= " << kxmax << endl;
    cout << " |ky| <= " << kymax << endl;
    cout << " |kz| <= " << kzmax << endl;


    cout << "Minimum   aliased grid : "
	 << 2*(kxmax+1) << " x "
	 << kymax+1 << " x "
	 << 2*(kzmax+1) << endl;

    cout << "Minimum dealiased grid : "
	 << 3*(kxmax+1) << " x "
	 << kymax+1 << " x "
	 << 3*(kzmax+1) << endl;

    cout << endl;
  }

  if (all || nrgy) {
    ChebyCoeff U(u.Ny(), u.a(), u.b(), Spectral);
    U[1] = 1;
    FlowField uU(u);
    uU += U;

    cout << "dissip (u)    == " << dissipation(u) << endl;
    cout << "forcing(u)    == " << forcing(u) << endl;
    cout << "energy (u)    == " << 0.5*L2Norm2(u) << endl;
    cout << "dissip (u+U)  == " << dissipation(uU) << endl;
    cout << "forcing(u+U)  == " << forcing(uU) << endl;
    cout << "energy (u+U)  == " << 0.5*L2Norm2(uU) << endl;
  }

  if (all || dyn) {
    cout << "------------Dynamics--------------------" << endl;
    // Construct Navier-Stoke Integrator
    DNSFlags flags;
    flags.timestepping = SBDF3;
    flags.dealiasing   = DealiasXZ;
    flags.constraint   = PressureGradient;
    flags.verbosity    = PrintTicks;
    flags.baseflow     = PlaneCouette;
    flags.dPdx         = 0.0;

    if (nlmeth == "skw")
      flags.nonlinearity = SkewSymmetric;
    else if (nlmeth == "div")
      flags.nonlinearity = Divergence;
    else if (nlmeth == "cnv")
      flags.nonlinearity = Convection;
    else if (nlmeth == "alt")
      flags.nonlinearity = Alternating;
    else
      flags.nonlinearity = Rotational;


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

    const int Nx = u.Nx();
    const int Ny = u.Ny();
    const int Nz = u.Nz();
    const Real Lx = u.Lx();
    const Real  a = u.a();
    const Real  b = u.b();
    const Real Lz = u.Lz();
    //const Real dt = (T < 0.01) ? T : 0.01;

    const int N = iround(T/dt);

    // Construct base flow for plane Couette: U(y) = y
    Vector y = chebypoints(Ny,a,b);

    FlowField q(Nx,Ny,Nz,1,Lx,Lz,a,b);
    DNS dns(u, nu, dt, flags, 0.0);

    cout << "computing du/dt: " << flush;
    FlowField dudt(u);
    dns.advance(dudt, q, N);
    cout << endl;
    dudt -= u;
    dudt *= 1.0/(N*dt);
    cout << "L2Norm(u)     == " << L2Norm(u) << endl;
    cout << "L2Norm(du/dt) == " << L2Norm(dudt) << endl;

    FlowField dudx;
    FlowField dudz;
    xdiff(u,dudx);
    zdiff(u,dudz);

    Real dudxnorm = L2Norm(dudx);
    Real dudznorm = L2Norm(dudz);

    cout << "L2Norm(du/dx) == " << dudxnorm << endl;
    cout << "L2Norm(du/dz) == " << dudznorm << endl;

    Real ueps = 1e-12;
    dudx *= (dudxnorm < ueps) ? 0.0 : 1.0/dudxnorm;
    dudz *= (dudznorm < ueps) ? 0.0 : 1.0/dudznorm;

    Real ax = L2IP(dudt,dudx);
    Real az = L2IP(dudt,dudz);

    FlowField Pu(u);
    Pu.setToZero();

    FlowField utmp;
    utmp = dudx;
    utmp *= ax;
    Pu += utmp;

    utmp = dudz;
    utmp *= az;
    Pu += utmp;

    Real Punorm   = L2Norm(Pu);
    Real waviness = Punorm/L2Norm(dudt);
    Real theta = acos(L2IP(Pu,dudt)/(L2Norm(dudt)*Punorm));

    cout << endl;
    cout << " waviness == " << waviness << endl;
    cout << "    theta == " << theta << " == " << theta*180/pi << " degrees\n";
    cout << endl;
    cout << "(waviness == fraction of dudt within dudx,dudz tangent plane)\n";
    cout << "(   theta == angle between dudt and dudx,dudz tangent plane)\n";
    cout << endl;
  }
}

Real energy(FlowField& u, int i, bool normalize) {
  assert(u.ystate() == Spectral);
  assert(u.xzstate()==Spectral);
  Real sum = 0.0;
  ComplexChebyCoeff prof(u.Ny(), u.a(), u.b(), Spectral);
  for (int mx=0; mx<u.Mx(); ++mx) {
    Real cz = 1.0; // cz = 2 for kz>0 to take account of kz<0 ghost modes
    for (int mz=0; mz<u.Mz(); ++mz) {
      for (int ny=0; ny<u.Ny(); ++ny)
	prof.set(ny, u.cmplx(mx,ny,mz,i));
      sum += cz*L2Norm2(prof, normalize);
      cz = 2.0;
    }
  }
  if (!normalize)
    sum *= u.Lx()*u.Lz();
  return sum;
}

char symmetric(const FieldSymmetry& s, const FlowField& u, Real& serr,Real& aerr, Real eps) {
  Real sqrteps= sqrt(eps);

  FlowField su = s(u);
  FlowField us, ua;
  ((us = u) += su) *= 0.5;
  ((ua = u) -= su) *= 0.5;
  serr = L2Dist(us, u);
  aerr = L2Dist(ua, u);

  // Check for symmetry
  if (serr < eps)
    return 'S';
  else if (serr < sqrteps)
    return 's';
  else if (aerr < eps)
    return 'A';
  else if (aerr < sqrteps)
    return 'a';
  else
    return ' ';
};

Real project(const FlowField& u, const FieldSymmetry& s, int sign) {
  FlowField Pu(u);
  if (sign<0)
    Pu -= s(u);
  else
    Pu += s(u);
  //Pu *= 0.5;
  return 0.25*L2Norm2(Pu);
}