#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

const double pi = 3.141592653589793238462643383279502;

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;
};

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 set_bdry_conds(interpolator<RBF> &u, interpolator<RBF> &v);


//********************** 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<dirichlet_op> D(new dirichlet_op);
    zero_func Z;
    //FIXME: Turns this into a purely interpolating interpolator,
    //without a BVP.
    shared_ptr<linear_BVP2> 
      u_bvp_init(new linear_BVP2(Omega, I, D, Z, Z)),
      v_bvp_init(new linear_BVP2(Omega, I, D, Z, Z));
      
    
    //init the interps.
    using namespace rbf;
    conical::set_n(7); thin_plate_spline::set_n(6); 
    c_infty_rbf::set_epsilon(0.01);

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

    interpolator<RBF_TYPE> 
      u0(u_bvp_init), 
      v0(v_bvp_init), 
      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();

    //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;

    //Negative because of the way the F's below are written.
    double dt = -0.01;

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

    //main loop, timestepping with RK4.
    size_t maxiter = 5;
    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);

      k1[0].interpolate(f_u.at());
      k1[1].interpolate(f_v.at());
      k1[2].interpolate(f_h.at());

      set_bdry_conds(k1[0],k1[1]);
      
      cout << "k1 done!" << endl;

      //debug Euler
      u0 = k1[0]; v0 = k1[1]; h0 = k1[2]; continue;
           
      //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));

      k2[0].interpolate(f_u.at());
      k2[1].interpolate(f_v.at());
      k2[2].interpolate(f_h.at());
 
      set_bdry_conds(k2[0],k2[1]);

      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));

      k3[0].interpolate(f_u.at());
      k3[1].interpolate(f_v.at());
      k3[2].interpolate(f_h.at());
      
      set_bdry_conds(k3[0],k3[1]);
      
      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]);

      k4[0].interpolate(f_u.at());
      k4[1].interpolate(f_v.at());
      k4[2].interpolate(f_h.at());
      
      set_bdry_conds(k4[0],k4[1]);

      cout << "k4 done!" << endl;
     
      //Grand finale
      u0.interpolate((u0 + (k1[0] + 2*k2[0] + 2*k3[0] + k4[0])/6).at());
      v0.interpolate((v0 + (k1[1] + 2*k2[1] + 2*k3[1] + k4[1])/6).at());
      h0.interpolate((h0 + (k1[2] + 2*k2[2] + 2*k3[2] + k4[2])/6).at());

      set_bdry_conds(u0,v0);
    }

    return 0;
  }
  catch(error exc){
    utils::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>
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));
}
 
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;
  }


  if(false){ //debug
    double r, th, N, M;
    N = 30; M  = 70;
    linalg::matrix out(static_cast<size_t>(N),static_cast<size_t>(M));
    size_t i,j;
    for(r = 0, j=1;  j <= N; r += 1/(N-1),j++){
      for(th = 0,i=1; i <= M; th += 2*pi/(M-1),i++){
	linalg::vector p(2); 
	p(1) = r*cos(th); 
	p(2) = r*sin(th);
	out(j,i) = u(p);
      }
    }
  
    ofs << out;
  }//debug
}


template<typename RBF>
void set_bdry_conds(interpolator<RBF> & u, interpolator<RBF> & v)
{
  bvp::bdry_values u_bdry(u.at()), v_bdry(v.at());
  bvp::normals N=u.get_normals(), M=v.get_normals();
  if(N != M){
    failure exc;
    exc.reason="u and v have incompatible domains. This should never"
      " happen.";
    exc.line=__LINE__; exc.file=__FILE__; throw exc;
  }

  //Project the vector (u,v) onto the perp of (N(1),N(2)) = (-N(2),N(1)).
  bvp::bdry_values u_bdry_new = u_bdry*N(2)*N(2) - v_bdry*N(1)*N(2);
  bvp::bdry_values v_bdry_new = v_bdry*N(1)*N(1) - u_bdry*N(1)*N(2);
  
  u.set_bdry_values(u_bdry_new);
  v.set_bdry_values(v_bdry_new);
}