#include <iostream>
#include <fstream>
#include <map>
#include <set>
#include <sstream>

#include <boost/shared_ptr.hpp>

#include "include/linalg.hpp"
#include "include/rbf.hpp"
#include "include/bvp.hpp"
#include "include/error.hpp"
#include "include/utils.hpp"
#include "include/diff_op.hpp"
#include "include/interpolator.hpp"
#include "include/func.hpp"

using namespace linalg;
using namespace bvp;

#define RBF_TYPE rbf::conical

//Will solve the system
//
//    u_t = -u*u_x - v*u_y - g*h_x         =: f1(u,v,h)
// 
//    v_t = -u*v_x - v*v_y - g*h_y         =: f2(u,v,h)
//
//    h_t = -u*h_x - h*u_x - v*h_y - h*v_y =: f3(u,v,h)
//
//where g is the acceleration due to gravity, u(x,y,0) = v(x,y,0) = 0,
//and h starts out as an undisturbed surface high with a small
//disturbance. We'll do this by RK4.


const double g = 9.8; // m/s^2

template<typename RBF>
class Fgen {
public:
  void set_interps(const interpolator<RBF>& u0, 
		   const interpolator<RBF>& v0,
		   const interpolator<RBF>& h0);
  void set_u(const interpolator<RBF>& u0);
  void set_v(const interpolator<RBF>& v0);
  void set_h(const interpolator<RBF>& h0);
  void set_dt(double dt_in);
protected:
  interpolator<RBF> u;
  interpolator<RBF> v;
  interpolator<RBF> h;
  double dt;
};

template<typename RBF>
class Fu : public Fgen<RBF>{
  using Fgen<RBF>::u;
  using Fgen<RBF>::v;
  using Fgen<RBF>::h;
  using Fgen<RBF>::dt;
public:
  bvp::interp_values at() const;
};

template<typename RBF>
class Fv : public Fgen<RBF>{
  using Fgen<RBF>::u;
  using Fgen<RBF>::v;
  using Fgen<RBF>::h;
  using Fgen<RBF>::dt;
public:
  bvp::interp_values at() const;
};

template<typename RBF>
class Fh : public Fgen<RBF>{
  using Fgen<RBF>::u;
  using Fgen<RBF>::v;
  using Fgen<RBF>::h;
  using Fgen<RBF>::dt;
public:
  bvp::interp_values at() const;
};

//FIXME: This iteration must work on interpolators without boundary
class iter_neumann : public bdry_diff_op{
public:
  //If this is u's, set 1, for v, set 2.
  iter_neumann(size_t u_or_v_in);
  double at(const realfunc &f, const point &p, const vector &n) const;
private:
  size_t u_or_v;
  double at(const realfunc &f, const point &p) const {return f(p);};
};

template<typename RBF>
class Gu_or_v : public realfunc{
public:
  //specify the OTHER, whether u or v.
  Gu_or_v(size_t u_or_v_in, 
	  const shared_ptr<domain> Omega_in) : u_or_v(u_or_v_in),
					       Omega(Omega_in) {};
  double at(const point& p) const;
  void set_other(const interpolator<RBF>& other_in){other = other_in;};
private:
  size_t u_or_v;
  const shared_ptr<domain> Omega;
  interpolator<RBF> other;
};

class zero_func : public realfunc{
public:
  double at(const point& p) const{p.size(); return 0;};
};

//************ Function declarations ****************************

template<typename RBF>
void save_interp(const shared_ptr<domain> Omega, 
		 const interpolator<RBF>& u, 
		 size_t n, string which_interp);

template<typename RBF>
void bdry_iter(interpolator<RBF> &u, interpolator<RBF> &v, 
	       const boost::shared_ptr<domain> Omega);


//********************** Main ******************************************
int main()
{  
  gsl_set_error_handler(&error_handling::errorHandler);

  try{
    using namespace std;
    using boost::shared_ptr;
    
    map<point, double> h_init = utils::read_pd_map("data/h_init.map");
      
    shared_ptr<domain> Omega(new domain("data/circ_intr.matrix",
					"data/circ_bdry.matrix",
					"data/circ_nrml.matrix"));

    shared_ptr<Id_op> I(new Id_op);
    shared_ptr<iter_neumann> 
      u_bdry_op(new iter_neumann(1)), 
      v_bdry_op(new iter_neumann(2));
    zero_func Z;
    shared_ptr<linear_BVP2> 
      u_bvp_init(new linear_BVP2(Omega, I, u_bdry_op, Z, Z)),
      v_bvp_init(new linear_BVP2(Omega, I, v_bdry_op, Z, Z));
      
   
    double dt = 1e-2;

    //init the interps.
    using namespace rbf;
    conical::set_n(5); thin_plate_spline::set_n(6); 
    c_infty_rbf::set_epsilon(0.01);

    cout << "Initialising... please wait." << endl;

    interpolator<RBF_TYPE> 
      u1, u0(u_bvp_init), 
      v1, v0(v_bvp_init), 
      h1, h0(Omega, h_init);

    //Precompute some RBFs
    u0.precompute_ev(); v0.precompute_ev(), h0.precompute_ev();
    u0.precompute_d1(); v0.precompute_d1(), h0.precompute_d1();


    u1 = u0; v1 = v0; h1 = h0;

    //Intermediate interpolators for RK4
    std::vector<interpolator<RBF_TYPE> >
      // FIXME: make the copy ctor for interps make different copies
      // of the bvp.
      k1(3,h0), k2(3,h0), k3(3,h0), k4(3,h0); 

    Fu<RBF_TYPE> f_u;
    Fv<RBF_TYPE> f_v;
    Fh<RBF_TYPE> f_h;

    f_u.set_dt(dt);
    f_v.set_dt(dt);
    f_h.set_dt(dt);

    //main loop
    size_t maxiter = 1000;
    for(size_t i = 1; i <= maxiter; i++){ 
      cout << "Now on iteration #" << i << endl;
      if(i % 1 == 0){
	save_interp(Omega,u0,i,"u");
	save_interp(Omega,v0,i,"v");
	save_interp(Omega,h0,i,"h");
      }

      //k1 
      f_u.set_interps(u0,v0,h0);
      f_v.set_interps(u0,v0,h0);
      f_h.set_interps(u0,v0,h0);
      
      cout << "k1 done!" << endl;

      bdry_iter(k1[0],k1[1],Omega);
           
      //k2
      f_u.set_interps(u0+(k1[0]/2), v0+(k1[1]/2), h0+(k1[2]/2));
      f_v.set_interps(u0+(k1[0]/2), v0+(k1[1]/2), h0+(k1[2]/2));
      f_h.set_interps(u0+(k1[0]/2), v0+(k1[1]/2), h0+(k1[2]/2));
 
      bdry_iter(k2[0],k2[1],Omega);

      cout << "k2 done!" << endl;

      //k3      
      f_u.set_interps(u0+(k2[0]/2), v0+(k2[1]/2), h0+(k2[2]/2));
      f_v.set_interps(u0+(k2[0]/2), v0+(k2[1]/2), h0+(k2[2]/2));
      f_h.set_interps(u0+(k2[0]/2), v0+(k2[1]/2), h0+(k2[2]/2));
      
      bdry_iter(k3[0],k3[1],Omega);
      
      cout << "k3 done!" << endl;

      //k4
      f_u.set_interps(u0+k3[0], v0+k3[1], h0+k3[2]);
      f_v.set_interps(u0+k3[0], v0+k3[1], h0+k3[2]);
      f_h.set_interps(u0+k3[0], v0+k3[1], h0+k3[2]);
      
      bdry_iter(k4[0],k4[1],Omega);

      cout << "k4 done!" << endl;
     
      //Grand finale
      u1.set_f(u0 + (k1[0] + 2*k2[0] + 2*k3[0] + k4[0])/6);
      v1.set_f(v0 + (k1[1] + 2*k2[1] + 2*k3[1] + k4[1])/6);
      h1.set_f(h0 + (k1[2] + 2*k2[2] + 2*k3[2] + k4[2])/6);
     
      //Enforce boundary conditions iteratively
      bdry_iter(u1,v1,Omega);
      
      u0 = u1;
      v0 = v1;
      h0 = h1;
    }

    return 0;
  }
  catch(error exc){
    utils::show_exception(exc);
    return ;
  }

}

//******************* Function definitions ***************************

template<typename RBF>
void Fgen<RBF>::set_interps(const interpolator<RBF>& u0, 
			    const interpolator<RBF>& v0,
			    const interpolator<RBF>& h0){
  set_u(u0);
  set_v(v0);
  set_h(h0);
}

template<typename RBF>
void Fgen<RBF>::set_u(const interpolator<RBF>& u0)
{
  u = u0;
}

template<typename RBF>
void Fgen<RBF>::set_v(const interpolator<RBF>& v0)
{
  v = v0;
}

template<typename RBF>
void Fgen<RBF>::set_h(const interpolator<RBF>& h0)
{
  h = h0;
}

template<typename RBF>
void Fgen<RBF>::set_dt(double dt_in)
{
  dt = dt_in;
}


template<typename RBF>
interp_values Fu<RBF>::at() const
{
  return   dt*(u()*u.d(1) +  v()*u.d(2) +  g*h.d(1));
}

template<typename RBF>
interp_values Fv<RBF>::at() const
{
  return   dt*(u()*v.d(1) +  v()*v.d(2) +  g*h.d(2));
}

template<typename RBF>
interp_values Fh<RBF>::at() const
{
  return   dt*(u()*h.d(1) + h()*u.d(1) +   v()*h.d(2) + h()*v.d(2));
}

iter_neumann::iter_neumann(size_t u_or_v_in)
{
  if(u_or_v_in < 1 or u_or_v_in > 2){
    error_handling::badArgument exc;
    exc.reason = "Argument to iter_neumann constructor must be 1 or 2";
    exc.line = __LINE__;
    exc.file = __FILE__;
    throw exc;
  }
  u_or_v = u_or_v_in;
}

double iter_neumann::at(const realfunc &f, const point &p, 
			const vector &n) const
{
  return n(u_or_v)*f(p);
}

template<typename RBF>
double Gu_or_v<RBF>::at(const point& p) const
{
  point n = Omega -> get_normals().at(p);
  return -n(u_or_v)*other(p);
}


template<typename RBF>
void save_interp(const shared_ptr<domain> Omega, 
		 const interpolator<RBF>& u, 
		 size_t n, string which_interp)
{
  using namespace std;
  string filename = "results/" + which_interp;
  if(n < 10000)
    filename += "0";
  if(n < 1000)
    filename += "0";
  if(n < 100)
    filename += "0";
  if(n < 10)
    filename += "0";
  stringstream ss;
  string n_string;
  ss << n;
  ss >> n_string;
  filename += n_string + ".map";
  ofstream ofs(filename.c_str());

  slice s(1,2);
  matrix M(Omega -> get_interior().size() +
	   Omega -> get_boundary().size(), 3);
  size_t k = 1;
  for(set<point>::const_iterator i = Omega -> get_interior().begin();
      i != Omega -> get_interior().end(); i++){
    M(k,s) = *i;
    M(k,3) = u(*i);
    k++;
  }
  for(set<point>::const_iterator i = Omega -> get_boundary().begin();
      i != Omega -> get_boundary().end(); i++){
    M(k,s) = *i;
    M(k,3) = u(*i);
    k++;
  }
  ofs << M;
}

template<typename RBF>
linalg::vector at_bdry(const interpolator<RBF>& u,
		       const boost::shared_ptr<domain> Omega)
{
  using namespace linalg;

  vector out(Omega -> get_boundary().size());
  std::set<point>::const_iterator I = Omega -> get_boundary().begin();

  for(size_t i = 1; i <= out.size(); i++){
    out(i) = u(*I);
    I++;
  }

  return out;
}

template<typename RBF>
void bdry_iter(interpolator<RBF> &u, interpolator<RBF> &v, 
	       const boost::shared_ptr<domain> Omega)
{
  Gu_or_v<RBF_TYPE> gu(2, Omega), gv(1, Omega);
  double err = 1;
  do{    
    linalg::vector u_old, u_new, v_old, v_new;
    u_old = at_bdry(u,Omega);
    gu.set_other(v);
    u.set_g(gu);
    u_new = at_bdry(u,Omega);

    v_old = at_bdry(v,Omega);
    gv.set_other(u);
    v.set_g(gv);
    v_new = at_bdry(v,Omega);
    
    double err_u = norm(u_new - u_old);
    double err_v = norm(v_new - v_old);
    err = (err_u > err_v? err_u : err_v);
  }while(err > 1e-2);
}