#include <fstream>
#include <iomanip>
#include <string>
#include <iostream>
#include <octave/oct.h>
#include <octave/dColVector.h>
#include <octave/dMatrix.h>

#include "channelflow/flowfield.h"
#include "channelflow/diffops.h"
#include "channelflow/symmetry.h"

#include "utilities.h"

using namespace std;
using namespace channelflow;

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

// Return a translation that maximizes the s-symmetry of u. 
// (i.e. return tau // that minimizes L2Norm(s(tau(u)) - tau(u))).
FieldSymmetry optimizePhase(const FlowField& u, const FieldSymmetry& s,
			    int Nsteps, Real residual, Real damp, 
			    bool verbose, Real x0, Real z0);

int main(int argc, char* argv[]) {
  string purpose("translate a field to maximize its shift-reflect "
		 "and/or shift-rotate symmetry");

  ArgList args(argc, argv, purpose);

  const bool s1opt   = args.getflag("-s1", "--shift-reflect", "optimize x-phase for s1 symmetry");
  const bool s2opt   = args.getflag("-s2", "--shift-rotate",  "optimize x-phase for s2 symmetry");
  const bool a1opt   = args.getflag("-a1", "--anti-shift-reflect", "optimize x-phase for s1 anti-symmetry");
  const bool a2opt   = args.getflag("-a2", "--anti-shift-rotate",  "optimize x-phase for s2 anti-symmetry");
  const bool force   = args.getflag("-f",  "--force", "force symmetries by projection after optimization");
  const Real ax      = args.getreal("-ax", "--axshift", 0.0, "initial guess for x translation: x -> x+ax*Lx");
  const Real az      = args.getreal("-az", "--axshift", 0.0, "initial guess for z translation: z -> z+az*Lz");
  const bool verbose = args.getflag("-v",  "--verbose", "print results during Newton search");
  //const bool verbose = args.getflag("-v",  "--verbose", "print results during Newton search");
  const int  Nsteps  = args.getint ("-N",  "--Nsteps", 10, "max # Newton steps");
  const Real eps     = args.getreal("-e",  "--eps", 1e-14, "stop Newton search when err<eps");
  const Real damp    = args.getreal("-d",  "--damping", 1.0, "damping factor for Newton steps");
  //const Real x0      = args.getreal("-x",  "--xtrans", 0.0, "initial guess for x-translation");
  //const Real z0      = args.getreal("-z",  "--ztrans", 0.0, "initial guess for z-translation");
  //const Real bound   = args.getreal("-b",  "--bound",  1.0, "stop if |x-x0|/Lx > bound");
  const string uname = args.getstr(2, "<infield>", "input field");
  const string oname = args.getstr(1, "<outfield>", "output field");

  if ((s1opt && a1opt) || (s2opt && a2opt)) {
    cerr << "Sorry, can't optimize for a symmetry and its antisymmetry.\n";
    cerr << "Please choose only one of -s1,-a1 and one of -s2,-a2" << endl;
    exit(1);
  }
  args.check();
  args.save("./");

  FlowField u(uname);
  FlowField u0(u);
  FlowField su;

  FieldSymmetry s1( 1, 1,-1, 0.5, 0);
  FieldSymmetry s2(-1,-1, 1, 0.5, 0.5);
  FieldSymmetry s3(-1,-1,-1, 0.0, 0.5);

  FieldSymmetry a(-1);
  FieldSymmetry a1 = a*s1;
  FieldSymmetry a2 = a*s2;
  FieldSymmetry a3 = a*s3;
  
  Real unorm2 = L2Norm2(u);
  
  cout << endl;
  cout.setf(ios::left);
  cout << "Initial energy decomposition:" << 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 << endl;
  cout.unsetf(ios::left);
  cout << "\nL2Dist(u,u0) == " << L2Dist(u,u0) << endl;

  if (ax != 0.0 || az != 0.0) {
    FieldSymmetry tau_guess(1,1,1,ax,az);
    u *= tau_guess;

    cout << "Applying initial guess at translations: u = tau(u)" << endl;
    cout << endl;
    Real unorm2 = L2Norm2(u);
    cout.setf(ios::left);
    cout << "Post-guess energy decomposition:" << 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 << endl;
    cout << "\nL2Dist(u,u0) == " << L2Dist(u,u0) << endl;
  }

  FieldSymmetry tau;
  FieldSymmetry t;
  if (s1opt) {
    t = optimizePhase(u, s1, Nsteps, eps, damp, verbose);
    u = t(u);
    tau *= t;
  }
  else if (a1opt) {
    t = optimizePhase(u, a1, Nsteps, eps, damp, verbose);
    u = t(u);
    tau *= t;
  }

  if (s2opt) {
    t = optimizePhase(u, s2, Nsteps, eps, damp, verbose);
    u = t(u);
    tau *= t;
  }
  else if (a2opt) {
    t = optimizePhase(u, a2, Nsteps, eps, damp, verbose);
    u = t(u);
    tau *= t;
  }

  cout << endl;
  cout.setf(ios::left);
  cout << "Optimized energy decomposition:" << 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 << endl;
  cout.unsetf(ios::left);
  cout << "Post-translation s1     symm err == " << L2Dist(u,s1(u)) << endl;
  
  if (force) {
    if (s1opt) {
      cout << "\nForcing s1 symmetry by projection." << endl;
      (u += s1(u)) *= 0.5;
    }
    else if (a1opt) {
      cout << "\nForcing s1 antisymmetry by projection." << endl;
      (u -= s1(u)) *= 0.5;
    }
    if (s2opt) {
      cout << "Forcing s2  symmetry by projection." << endl;
      (u += s2(u)) *= 0.5;
    }
    else if (a2opt) {
      cout << "Forcing s2 antisymmetry by projection." << endl;
      (u -= s2(u)) *= 0.5;
    }

    unorm2 = L2Norm2(u);
    cout << endl;
    cout << "Forced energy decomposition:" << endl;
    cout.setf(ios::left);
    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 << endl;
  }
  cout << "\nL2Dist(u,u0)      == " << L2Dist(u,u0) << endl;
  cout << setprecision(16);
  cout << "optimal translation == " << tau << endl;
  u.binarySave(oname);
}

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 L2Norm2(Pu);
}
    
  
// Return a translation that maximizes the s-symmetry of u. 
// (i.e. return tau // that minimizes L2Norm(s(tau(u)) - tau(u))).
FieldSymmetry optimizePhase(const FlowField& u, const FieldSymmetry& s,
			    int Nsteps, Real residual, Real damp, 
			    bool verbose, Real x0, Real z0) {

  FieldSymmetry tau;
  if (s.sx()==1 && s.sz()==1) 
    ;
  else if (s.sx() == -1 ^ s.sz() == -1) {

    // Newton search for x-translation that maximizes shift-rotate symmetry
    if (verbose) 
      cout << "\noptimizePhase: Newton search on translation" <<endl;
    
    int i   = (s.sx() == -1) ? 0 : 2; // are we optimizing on x or z?
    Real L  = (s.sx() == -1) ? u.Lx() : u.Lz(); 
    array<Real> x(3);
    x[0] = x0;
    x[2] = z0;

    for (int n=0; n<Nsteps; ++n) {
      tau = FieldSymmetry(x[0]/u.Lx(), x[2]/u.Lz());
      
      FlowField tu = tau(u);
      FlowField stu = s(tu);
      FlowField tux = diff(tu, i, 1);
      FlowField stux = diff(stu, i, 1);
      FlowField tuxx = diff(tux, i, 1);
      FlowField stuxx = diff(stux, i, 1);

      // g(x) is df/dx where f == (L2Norm(stu - tu)) and t = tau(x/Lx,0)
      Real g    =  L2IP(stux, tu) - L2IP(stu, tux);
      Real dgdx = 2*L2IP(stux, tux) - L2IP(stuxx, tu) - L2IP(stu, tuxx);
      
      if (verbose) {
	Real f = L2Dist(stu, tu);
	cout << endl;
	cout << "Newton step n  == " << n << endl;
	cout << "tau.ax         == " << tau.ax() << endl;
	cout << "tau.az         == " << tau.az() << endl;
	cout << "zeroable g     == " << g << endl;
	cout << "L2Norm(stu-tu) == " << f << endl;
      }
      if (abs(g) < residual) {
	cout << "\noptimizePhase: Newton search converged." << endl;
	break;
      }
    
      Real dx = -g/dgdx;
      x[i] += damp*dx;

      if (abs(x[i])/L > 1.0) {
	cout << "\nerror in optimizePhase(FlowField, FieldSymmetry) :" << endl;
	cout << "Newton search went beyond bounds. Returning identity." << endl;
	tau = FieldSymmetry();
	break;
      }
    }
  }
  else {
    cout << "\nerror in optimizePhase(FlowField, FieldSymmetry) :" << endl;
    cout << "Not yet implemented. Returning identity." << endl;
  }
    
  if (verbose) {
    FlowField tu = tau(u);
    FlowField stu = s(tu);
    cout << "\noptimal translation == " << tau << endl;
    cout << "Post-translation symmetry error == " << L2Dist(stu, tu) << endl;
  }
  return tau;
}