#include "hole_filling.h"
#include "IO.h"
#include "util.h"
#include <MeshReconstruction.h>
#include <OpenMesh/Core/Utils/PropertyManager.hh>


static std::vector<std::vector<HalfedgeHandle>> findHoles(MyMesh &mesh) {
	std::vector<std::vector<HalfedgeHandle>> holes;
	std::set<HalfedgeHandle> ignore;
	for (HalfedgeHandle it : mesh.halfedges()) {
		if (mesh.is_boundary(it)
			&& ignore.find(it) == ignore.end()) {
			holes.emplace_back();
			holes.back().push_back(it);
			ignore.insert(it);
			for (HalfedgeHandle it2 : HalfedgeLoopRange<MyMesh>(mesh, it)) {
				holes.back().push_back(it2);
				ignore.insert(it2);
			}
		}
	}
	return holes;
}


void fillHoleDumb(MyMesh &mesh, std::vector<HalfedgeHandle> &hole) {
	mesh.request_vertex_status();
	mesh.request_edge_status();
	mesh.request_face_status();

	Point center(0, 0, 0);
	size_t length = 0;
	for (HalfedgeHandle it : hole) {
		VertexHandle vert = mesh.to_vertex_handle(it);
		center += mesh.point(vert);
		length++;
	}
	center /= length;
	VertexHandle center_handle = mesh.new_vertex_dirty(center);
	mesh.set_color(center_handle, mesh.default_color);

	for (HalfedgeHandle it : hole) {
		mesh.add_face(mesh.from_vertex_handle(it),
					  mesh.to_vertex_handle(it),
					  center_handle);
	}
}


void fillHolesDumb(MyMesh &mesh) {
	mesh.holes = findHoles(mesh);
	for (auto hole : mesh.holes) {
		fillHoleDumb(mesh, hole);
	}
}


using namespace std ;
using namespace MeshReconstruction;

/* ************** Implicit RBF ************** */

Implicit_RBF::Implicit_RBF(vector<float> alpha, vector<float> beta, vector<MyMesh::Point> center) : _alpha(alpha), _beta(beta), _center(center)
{
    _n = _alpha.size() ;
    _d = _beta.size() ;
    if (_center.size() != _n)
        throw runtime_error("Inconsistent size of alpha and centers in Implicit_RBF constructor.");
}

// Computation of the value of the implicit surface at point X
float Implicit_RBF::val(MyMesh::Point X) const
{
    float res = 0 ;
    // Computation of the sum of RBF at centers
    for(int i=0; i<_n; i++)
    {
        res += _alpha.at(i) * myphi((X-_center.at(i)).norm()) ;
    }
    // Computation of the polynomial part
    for(int j=0; j<_d; j++)
    {
        res += _beta.at(j) * myp(j, X) ;
    }
    return res ;
}

/* ************** Hole Filling ************** */

Hole_Filling::Hole_Filling(MyMesh &mesh, float scale, float discr)
	: _mesh(mesh),
	  _scale(scale),
	  _discr(discr)
{
    _mesh.add_property(_vprop,  "vprop_flag");
    cout << "Starting hole filling ..." << endl ;
}

// ***** Computation of boundary and its neighborhood

MyMesh::HalfedgeHandle Hole_Filling::find_boundary_edge()
{
    MyMesh::HalfedgeIter he_it = _mesh.halfedges_begin();
    while ( (he_it != _mesh.halfedges_end()) && (!_mesh.is_boundary(*he_it)))
    {
        ++he_it ;
    }
    if (he_it != _mesh.halfedges_end())
        return *he_it ;
    else
        throw std::runtime_error("Boundary HE does not exist") ;
}

vector<MyMesh::VertexHandle> Hole_Filling::find_boundary(MyMesh::HalfedgeHandle heh)
{
    MyMesh::HalfedgeHandle heh_ini = heh ;
    vector<MyMesh::VertexHandle> boundary ;

    // Follow (and memorize) boundary edges
    do
    {
        boundary.push_back(_mesh.to_vertex_handle(heh));
        heh = _mesh.next_halfedge_handle(heh);
    } while (heh != heh_ini) ;
    return boundary ;
}

void Hole_Filling::init_mark_boundary(const vector<MyMesh::VertexHandle> & bnd)
{
    for (MyMesh::VertexIter v_it = _mesh.vertices_begin() ; v_it != _mesh.vertices_end(); ++v_it)
        _mesh.property(_vprop, *v_it) = false ;
    for (int i=0; i<bnd.size(); ++i)
        _mesh.property(_vprop, bnd.at(i)) = true ;
}

vector<MyMesh::VertexHandle> Hole_Filling::next_neighbors(const vector<MyMesh::VertexHandle> & bnd)
{
    // Visit bnd vertices to find and mark next circle
    vector<MyMesh::VertexHandle> next_bnd ;
    for (int i=0; i<bnd.size(); i++)
    {
        for (MyMesh::VertexVertexIter vv_it = _mesh.vv_iter(bnd.at(i));vv_it.is_valid();++vv_it)
        {
            if (_mesh.property(_vprop, *vv_it) == false) // new vertex
            {
                _mesh.property(_vprop, *vv_it) = true ;
                next_bnd.push_back(*vv_it);
            }
        }
    }
    return next_bnd ;
}

// ***** Computation of RBF

pair<pair<Eigen::MatrixXd &, Eigen::VectorXd &>, vector<MyMesh::Point> &> Hole_Filling::compute_approx_mat(vector<MyMesh::VertexHandle> vlist, double normal_scale)
{
    const int n(vlist.size()), d(10) ;
    Eigen::MatrixXd & A = *(new Eigen::MatrixXd(3*n+d,3*n+d)) ;
    Eigen::MatrixXd Phi(3*n,3*n);
    Eigen::MatrixXd P(3*n,d);
    Eigen::VectorXd & B = *(new Eigen::VectorXd(3*n+d));

    vector<MyMesh::Point> & pts_list = *(new vector<MyMesh::Point>);
    //Append vertices to pts_list
    for (int i=0; i<n; i++)
    {
        pts_list.push_back(_mesh.point(vlist.at(i))) ;
    }
    //Append vertices+normals to pts_list
    for (int i=0; i<n; i++)
    {
        pts_list.push_back(_mesh.point(vlist.at(i)) + _mesh.normal(vlist.at(i)) * normal_scale) ;
    }
    //Append vertices-normals to pts_list
    for (int i=0; i<n; i++)
    {
        pts_list.push_back(_mesh.point(vlist.at(i)) - _mesh.normal(vlist.at(i)) * normal_scale) ;
    }

    int nn = pts_list.size() ;
    // Compute corresponding B vector (0 / 1 / -1 )
    B << Eigen::VectorXd::Zero(n), Eigen::VectorXd::Ones(n), -Eigen::VectorXd::Ones(n), Eigen::VectorXd::Zero(d) ;

    // Fill Phi matrix
	//TODO
	for (int i = 0; i < 3*n; i++) {
		for (int j = 0; j <= i; j++) {
			Phi(i, j) = myphi((pts_list[i] - pts_list[j]).norm());
			Phi(j, i) = Phi(i, j);
		}
	}

    // Fill P matrix
    for (int i = 0; i < 3*n; i++) {
		for (int j = 0; j < d; j++) {
			P(i, j) = myp(j, pts_list[i]);
		}
	}

    // Set final A matrix
    /* A = Phi | P
     *      P' | 0    */
	A << Phi, P, P.transpose(), Eigen::MatrixXd::Zero(d, d);
    // TODO

    cout << "size of pts_list : " << nn << endl ;
    cout << "End computation of matrices" << endl ;

    return {{A,B},pts_list} ;
}

pair<vector<float>&, vector<float>&> Hole_Filling::solve_approx(const pair<Eigen::MatrixXd &, Eigen::VectorXd &> &p, int n, int d)
{
    Eigen::MatrixXd & A = p.first ;
    Eigen::VectorXd & B = p.second ;

    Eigen::VectorXd res = A.householderQr().solve(B) ;

    cout << "End of solver" << endl ;
    cout << "res : " << res.head(10) << endl ;
    vector<float> & alpha = *(new vector<float>);
    vector<float> & beta = *(new vector<float>);

    if (res.size() != (n+d))
    {
        cout << "taille du res : " << res.size() << endl ;
        throw std::runtime_error("Error in solve_approx") ;
    }
    for (int i=0; i<n; i++)
    {
        alpha.push_back(res(i)) ;
    }

    for (int j=0; j<d; j++)
    {
        beta.push_back(res(n+j)) ;
    }

    return {alpha, beta} ;
}

// ***** IO

void Hole_Filling::colorize_prop()
{
    for (MyMesh::VertexIter v_it = _mesh.vertices_begin(); v_it != _mesh.vertices_end() ; ++v_it)
    {
        if(_mesh.property(_vprop, *v_it) == true)
            _mesh.set_color(*v_it, MyMesh::Color(255, 0, 0)) ;
        else
			_mesh.set_color(*v_it, MyMesh::Color(200, 200, 200)) ;
    }
}

void Hole_Filling::colorize_verts(const vector<MyMesh::VertexHandle> &vlist)
{
    for (MyMesh::VertexIter v_it = _mesh.vertices_begin(); v_it != _mesh.vertices_end() ; ++v_it)
    {
           _mesh.set_color(*v_it, MyMesh::Color(200, 200, 200)) ;
    }

    for(int i=0; i<vlist.size(); i++)
    {
        _mesh.set_color(vlist.at(i), MyMesh::Color(255, 0, 0)) ;
    }
}

Rect3 Hole_Filling::estimate_BB(const vector<MyMesh::VertexHandle> &vlist)
{
    MyMesh::Point minp = _mesh.point(vlist.at(0)), maxp=minp ;
    for (int i=0; i<vlist.size(); i++)
    {
        minp = min(minp, _mesh.point(vlist.at(i))) ;
        maxp = max(maxp, _mesh.point(vlist.at(i))) ;
    }
    MyMesh::Point sizep(maxp-minp), centerp = (minp+maxp)/2 ;
    minp = centerp - _scale*sizep/2 ;
    sizep *= _scale ;

    Rect3 domain ;
    domain.min = {minp[0], minp[1], minp[2]} ;
    domain.size = {sizep[0], sizep[1], sizep[2]} ;
    return domain ;
}

Mesh Hole_Filling::poly_n_out(const Implicit_RBF &implicit, Rect3 domain)
{
    auto implicitEq = [&implicit](Vec3 const& pos)
    {
        MyMesh::Point X(pos.x, pos.y, pos.z) ;
        return implicit.val(X) ;
    };


    Vec3 cubeSize(domain.size*(_discr)) ;
    cout << "mind " << domain.min << endl ;
    cout << "sized" << domain.size << endl ;
    cout << "cubesize "<< cubeSize << endl ;

    return MarchCube(implicitEq, domain, cubeSize);
}


/* Computes a mesh's bounding box and stores it in a mesh property
 * named "bounding_box". */
static void computeMeshBoundingBox(MyMesh &mesh, Hole_Filling &hf) {
	try {
		auto mesh_bb = OpenMesh::getProperty<void, Rect3>
			(mesh, "bounding_box");
	} catch (const std::runtime_error &e) {
		auto mesh_bb = OpenMesh::getOrMakeProperty<void, Rect3>
			(mesh, "bounding_box");
		std::vector<VertexHandle> verts;
		for (VertexHandle vh : mesh.vertices()) {
			verts.push_back(vh);
		}
		*mesh_bb = hf.estimate_BB(verts);
	}
}


MyMesh fillHoleImplicit(MyMesh &mesh, Hole_Filling &hf,
						std::vector<HalfedgeHandle> &hole) {
	computeMeshBoundingBox(mesh, hf);
	Rect3 mesh_bb = *OpenMesh::getProperty<void, Rect3>(mesh, "bounding_box");
	double diag = mesh_bb.size.Norm();

	std::vector<VertexHandle> verts;
	for (HalfedgeHandle hh : hole) {
		verts.push_back(mesh.to_vertex_handle(hh));
	}
	auto bb = hf.estimate_BB(verts) ;
	verts = hf.next_neighbors(verts);
	auto [system, pts_list] = hf.compute_approx_mat(verts, diag * .1);
	auto [alpha, beta] = hf.solve_approx(system, pts_list.size(), 10);
	Implicit_RBF rbf(alpha, beta, pts_list);

	Mesh filling = 	hf.poly_n_out(rbf, bb);
	MyMesh ret;
	for (const Vec3 &v : filling.vertices) {
		VertexHandle vh = ret.new_vertex({v.x, v.y, v.z});
		ret.set_color(vh, ret.default_color);
	}
	for (const Triangle &t : filling.triangles) {
		ret.add_face(ret.vertex_handle(t[0]),
					 ret.vertex_handle(t[1]),
					 ret.vertex_handle(t[2]));
	}
	return ret;
}


std::vector<MyMesh> fillHolesImplicit(MyMesh &mesh,
									  float scale, float discr) {
	mesh.holes = findHoles(mesh);
	std::vector<MyMesh> fillings;
	for (auto hole : mesh.holes) {
		Hole_Filling hf(mesh, scale, discr);
		fillings.push_back(fillHoleImplicit(mesh, hf, hole));
	}
	return fillings;
}