#include <unicorn/ElasticSmoother.h>
#include <unicorn/ElasticFormFactory.h>

using namespace dolfin::unicorn;
using namespace dolfin;

//-----------------------------------------------------------------------------
ElasticSmoother::ElasticSmoother(Mesh& mesh) : mesh(mesh)
{
}
//-----------------------------------------------------------------------------
void ElasticSmoother::maph0(Mesh& mesh, Mesh& sub,
			    MeshFunction<int>& cell_map,
			    MeshFunction<real>& h0,
			    MeshFunction<real>& subh0)
{
  subh0.init(sub, sub.topology().dim());

  for (CellIterator c(mesh); !c.end(); ++c)
  {
    Cell& cell = *c;

    //int subid = cell_map(cell);
//     cout << "subid: " << subid << endl;
//     cout << "h0: " << h0.get(cell) << endl;

    if(cell_map.get(cell) != -1)
      subh0.set(cell_map(cell), h0.get(cell));
  }
}
//-----------------------------------------------------------------------------
void ElasticSmoother::smooth(MeshFunction<bool>& smoothed_cells,
			     MeshFunction<real>& h0)
{
  cout << "elastic smooth" << endl;
  int d = mesh.topology().dim();

  Mesh sub;

  MeshFunction<int> vertex_map;
  MeshFunction<int> cell_map;

  MeshFunction<real> subh0;
  submesh(mesh, sub, smoothed_cells, vertex_map, cell_map);

  cout << "submesh created" << endl;
  //sub = mesh;

  maph0(mesh, sub, cell_map, h0, subh0);

  MeshQuality q(sub);
  
  real mu_max, mu_min, mu_avg;
  real h_max, h_min, h_avg;
  
  q.meshQuality(mu_min, mu_max, mu_avg, h_min, h_max, h_avg);

  real E = 1.0e0; // Young's modulus
  real nu = 0.0; // Poisson's ratio

  real lmbdaval = E * nu / ((1 + nu) * (1 - 2 * nu));
  real muval = E / (2 * (1 + nu));

  //real betaval = E; // Damping (sqrt(E))
  real betaval = 0.2 * E; // Damping (sqrt(E))

  //real kk = 1.0 / 2.0 * h_min;
  //real kk = 1.0e-1;
  real kk = 1.0 / 2.0 * h_min * 1.0e0;


  cout << "lmbda: " << lmbdaval << endl;
  cout << "mu: " << muval << endl;

  MySource f(sub);
  Density rho(sub);
  Quality qual(sub);
  MyBC bcf(sub);
  InitialValue u0(sub, subh0);

  DirichletBoundary dboundary;
  DirichletBC bc(bcf, sub, dboundary);
  
  Function lmbda(sub, lmbdaval);
  Function mu(sub, muval);
  Function beta(sub, betaval);
  Function kf(sub, kk);

  real T = 100 * kk;
  //real T = kk;
  //real T = 3.0;

  real ode_tol = 1.0e+3;

  set("ODE method", "dg");
  set("ODE order", 0);
  //set("ODE implicit", false);
  set("ODE implicit", true);
  set("ODE nonlinear solver", "newton");
  //set("ODE nonlinear solver", "fixed-point");
  //set("ODE linear solver", "direct");
  set("ODE linear solver", "iterative");
  set("ODE matrix-free jacobian", false);
  set("ODE M matrix constant", true);


  set("ODE tolerance", ode_tol);
  set("ODE discrete tolerance", ode_tol);
  //set("ODE fixed-point damping", 1.0e-01);
  
  set("ODE monitor convergence", true);
  
  set("ODE fixed time step", true);
  set("ODE initial time step", kk);
  set("ODE maximum time step", kk);
  //set("ODE initial time step", h_min / 100.0);
  //set("ODE maximum time step", h_min / 100.0);
  
  set("ODE save solution", false);
  set("ODE solution file name", "primal.py");
  set("ODE number of samples", 100);

  Function U;
  Function B;
  Function B0;

  Form* a = 0;
  Form* L = 0;
  Form* aJac = 0;
  Form* aS = 0;
  Form* LS = 0;
  Form* aP = 0;
  Form* LP = 0;

  ElasticFormFactory(sub,
		     a, L, aJac, aS, LS, aP, LP,
		     U, B, u0, rho, mu, lmbda, beta, qual, kf, f);

  ElasticityPDE pde(a, L, aJac, aP, LP, aS, LS, sub, bc, T, U, B, B0, kk, h0);

  //Compute solution
  pde.solve(U);


  MeshGeometry& geometry = mesh.geometry();

  for (VertexIterator n(mesh); !n.end(); ++n)
  {
    Vertex& vertex = *n;

    if(vertex_map(vertex) != -1)
    {
      Vertex subv(sub, vertex_map(vertex));

      for(int i = 0; i < d; i++)
	geometry.x(vertex.index(), i) = subv.point()[i];
    }
  }  

  delete a;
  delete L ;
  delete aJac ;
  delete aS ;
  delete LS ;
  delete aP ;
  delete LP ;
}
//-----------------------------------------------------------------------------
bool ElasticSmoother::onBoundary(Cell& cell)
{
  int d = cell.mesh().topology().dim();

  for (FacetIterator f(cell); !f.end(); ++f)
  {
    if(f->numEntities(d) == 1)
    {
      return true;
    }
  }
  return false;
}
//-----------------------------------------------------------------------------
void ElasticSmoother::worstElement(Mesh& mesh, int& index,
				   MeshFunction<bool>& masked_cells)
{
  int d = mesh.topology().dim();
  mesh.init(d - 1, d);

  MeshQuality mqual(mesh);

  real mu_min = 1.0e12;
  index = -1;

  for (CellIterator c(mesh); !c.end(); ++c)
  {
    Cell& cell = *c;

    real qual = mqual.cellQuality(cell);
    if(qual < mu_min && !onBoundary(cell) && !masked_cells.get(cell))
      //if(qual < mu_min)
    {
      index = cell.index();
      mu_min = qual;
    }
    //mu_min = std::min(mu_min, qual);
  }
  cout << "index: " << index << endl;
}
//-----------------------------------------------------------------------------
void ElasticSmoother::elementNhood(Mesh& mesh, Cell& element,
				   MeshFunction<bool>& elements,
				   int depth)
{
  elements.set(element, true);

  if(depth == 0)
    return;

  for (CellIterator c(element); !c.end(); ++c)
  {
    Cell& cell = *c;

    //elements.set(cell, true);
    elementNhood(mesh, cell, elements, depth - 1);
  }
}
//-----------------------------------------------------------------------------
void ElasticSmoother::submesh(Mesh& mesh, Mesh& sub,
			      MeshFunction<bool>& smoothed_cells,
			      MeshFunction<int>& old2new_vertex,
			      MeshFunction<int>& old2new_cell)
{
  old2new_vertex.init(mesh, 0);
  old2new_cell.init(mesh, mesh.topology().dim());

  int ncells = 0;
  int nvertices = 0;
  
  // Count cells and vertices in submesh
  for (CellIterator c(mesh); !c.end(); ++c)
  {
    Cell& cell = *c;

    if(smoothed_cells.get(cell) == true)
    {
      ncells++;
    }
  }

  for (VertexIterator n(mesh); !n.end(); ++n)
  {
    Vertex& vertex = *n;

    bool included = false;

    for (CellIterator c(vertex); !c.end(); ++c)
    {
      Cell& cell = *c;
      
      if(smoothed_cells.get(cell) == true)
      {
	included = true;
      }
    }

    if(included)
    {
      nvertices++;
    }
  }

  


  // Get cell type
  const CellType& cell_type = mesh.type();

  unsigned int current_vertex = 0;
  unsigned int current_cell = 0;

  MeshEditor editor;
  editor.open(sub, cell_type.cellType(),
              mesh.topology().dim(), mesh.geometry().dim());

  // Specify number of vertices and cells
  editor.initVertices(nvertices);
  editor.initCells(ncells);

  for (VertexIterator n(mesh); !n.end(); ++n)
  {
    Vertex& vertex = *n;

    bool included = false;

    for (CellIterator c(vertex); !c.end(); ++c)
    {
      Cell& cell = *c;
      
      if(smoothed_cells.get(cell) == true)
      {
	included = true;
      }
    }

    if(included)
    {
      old2new_vertex.set(vertex.index(), current_vertex);
      editor.addVertex(current_vertex, vertex.point());
      current_vertex++;
    }
    else
    {
      old2new_vertex.set(vertex.index(), -1);
    }
  }

  Array<unsigned int> cell_vertices(cell_type.numEntities(0));
  for (CellIterator c(mesh); !c.end(); ++c)
  {
    Cell& cell = *c;

    if(smoothed_cells.get(cell) == true)
    {
      int cv_idx = 0;
      for (VertexIterator n(cell); !n.end(); ++n)
      {
	int id = old2new_vertex.get(n->index());
	if(id == -1)
	{
	  cout << "broken: " << n->index() << endl;
	}
	cell_vertices[cv_idx++] = id;
      }

      old2new_cell.set(cell.index(), current_cell);
      editor.addCell(current_cell++, cell_vertices);    
    }
    else
    {
      old2new_cell.set(cell.index(), -1);
    }
  }

  editor.close();

  File submeshfile("submesh.xml");
  submeshfile << sub;


}
//-----------------------------------------------------------------------------
void ElasticSmoother::ElasticityPDE::deform(Mesh& mesh, Function& u)
{
//   cout << "X: " << endl;
//   X.vector().vec().disp();
//   cout << "U: " << endl;
//   U.vector().vec().disp();

  MeshGeometry& geometry = mesh.geometry();
  
  int N = mesh.numVertices();
  
  u.interpolate(vxvalues);
  
  // Update the mesh
  for(VertexIterator n(mesh); !n.end(); ++n)
  {
    Vertex& vertex = *n;
    
    for(unsigned int i = 0; i < mesh.topology().dim(); i++)
    {
      geometry.x(vertex.index(), i) = vxvalues[i * N + vertex.index()];
//       cout << "u(" << i << ", " << vertex.index() << ": " <<
// 	vxvalues[i * N + vertex.index()] << endl;
    }
  }
}
//-----------------------------------------------------------------------------
