#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 kwantxi;

#define RBF_TYPE kwantxi::conical

//Will solve the system
//
//    u_t = -u*u_x - v*u_y - g*h_x         =: g1(u,v,h)
// 
//    v_t = -u*v_x - v*v_y - g*h_y         =: g2(u,v,h)
//
//    h_t = -u*h_x - h*u_x - v*h_y - h*v_y =: g3(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 directly computing the next timestep
//by
//
// u_{t+1} = u_t + /\t*g1(u_t, v_t, h_t) =: f_u(u_t, v_t, h_t, /\t)
//
// v_{t+1} = v_t + /\t*g2(u_t, v_t, h_t) =: f_v(u_t, v_t, h_t, /\t)
//
// h_{t+1} = h_t + /\t*g3(u_t, v_t, h_t) =: f_h(u_t, v_t, h_t, /\t)
//
//Let's hope some waves propagate here. :-)

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

template<typename RBF>
class Fgen : public realfunc{
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:
  double dt;
  interpolator<RBF> u;
  interpolator<RBF> v;
  interpolator<RBF> h;
};

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:
  double at(const point& p) 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:
  double at(const point& p) 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:
  double at(const point& p) const;
};

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(&kwantxi::errorHandler);

  try{
    using namespace std;
    using boost::shared_ptr;
    
    map<point, double> h_init = kwantxi::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 = 0.0005;

    //init the interps.
    conical::set_n(5); thin_plate_spline::set_n(6); 
    c_infty_kwantxi::set_epsilon(0.01);

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

    interpolator<RBF_TYPE> 
      u1, u0(u_bvp_init), 
      v1, v0(v_bvp_init), 
      h1, h0(h_init);
    u1 = u0; v1 = v0; h1 = 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 = 20;
    for(size_t i = 1; i <= maxiter; i++){ 
      cout << "Now on iteration #" << i << endl;
      save_interp(Omega,u0,i,"u");
      save_interp(Omega,v0,i,"v");
      save_interp(Omega,h0,i,"h");

      f_h.set_interps(u0,v0,h0);
      f_u.set_interps(u0,v0,h0);
      f_v.set_interps(u0,v0,h0);

      h1.set_f(f_h);      
      u1.set_f(f_u);
      v1.set_f(f_v);

      //Enforce boundary conditions iteratively
      bdry_iter(u1,v1,Omega);
      
      u0 = u1;
      v0 = v1;
      h0 = h1;
    }

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

}

//******************* 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>
double Fu<RBF>::at(const point& p) const{
  return u(p) - dt*( 
		    u(p)*u.d(p,1) +
		    v(p)*u.d(p,2) +
		    g*h.d(p,1)
		    );
}

template<typename RBF>
double Fv<RBF>::at(const point& p) const{
  return v(p) - dt*( 
		    u(p)*v.d(p,1) +
		    v(p)*v.d(p,2) +
		    g*h.d(p,2)
		    );
}

template<typename RBF>
double Fh<RBF>::at(const point& p) const{
  return h(p) - dt*(
		    u(p)*h.d(p,1) +
		    h(p)*u.d(p,1) +
		 
		    v(p)*h.d(p,2) +
		    h(p)*v.d(p,2)
		    );
}

iter_neumann::iter_neumann(size_t u_or_v_in)
{
  if(u_or_v_in < 1 or u_or_v_in > 2){
    kwantxi::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>
kwantxi::vector at_bdry(const interpolator<RBF>& u,
		       const boost::shared_ptr<domain> Omega)
{
  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{    
    kwantxi::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);
}