#ifndef __ELASTIC_SMOOTHER_H
#define __ELASTIC_SMOOTHER_H

#include <dolfin.h>
#include <boost/numeric/ublas/matrix.hpp>
#include <unicorn/EquiAffineMap.h>
#include <unicorn/MeshQuality.h>
#include <unicorn/TimeDependentPDE.h>
#include "unicorn/Project.h"

namespace dolfin { namespace unicorn
{
//   void meshQuality(Mesh& mesh,
// 		   real& mu_min, real& mu_max,
// 		   real& h_min, real& h_max);

  class ElasticSmoother
  {
  public:
    
    /// Constructor
    ElasticSmoother(Mesh& mesh);
    ~ElasticSmoother()
    {
    };

    void smooth(MeshFunction<bool>& smoothed_cells,
		MeshFunction<real>& h0);

    void maph0(Mesh& mesh, Mesh& sub, MeshFunction<int>& cell_map,
	       MeshFunction<real>& h0, MeshFunction<real>& subh0);

    static bool onBoundary(Cell& cell);
    
    static void worstElement(Mesh& mesh, int& index,
			     MeshFunction<bool>& masked_cells);
    static void elementNhood(Mesh& mesh, Cell& element,
			     MeshFunction<bool>& elements,
			     int depth);
    static void submesh(Mesh& mesh, Mesh& sub,
			MeshFunction<bool>& smoothed_cells,
			MeshFunction<int>& old2new_vertex,
			MeshFunction<int>& old2new_cell);
    
    Mesh& mesh;
    
    // Sub domain for Dirichlet boundary condition
    class DirichletBoundary : public SubDomain
    {
    public:
      bool inside(const real* x, bool on_boundary) const
      {
	if(on_boundary)
	  return true;
	else
	  return false;
      }
    };

    class MySource : public Function
    {
    public:
      MySource(Mesh& mesh) : Function(mesh)
      {
      }
      
      void eval(real* values, const real* x) const
      {
	int d = cell().dim();

	for(int i = 0; i < d; i++)
	{
	  values[i] = 0.0;
	}
      }
    };
    
    class Density : public Function
    {
    public:
      Density(Mesh& mesh) : Function(mesh)
      {
      }
      
      void eval(real* values, const real* x) const
      {
	values[0] = 1.0;
      }
    };
    
    class Quality : public Function
    {
    public:
    Quality(Mesh& mesh) : Function(mesh), q(mesh), maxqual(0.0),
	minqual(1.0e12)
      {
	// Compute h
	q.init(mesh.topology().dim());
      
	MeshQuality mqual(mesh);

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

	  // 	EquiAffineMap map;
	  // 	map.update(cc);
	  //real h = cc.diameter();
	  //real qual = 0.25 * (h * h) / cc.volume();
	  real qq = mqual.cellQuality(cc);
	  //qual = 1.0 / pow(qual, 2.0);
	  //qual = 1.0;
	  real qual = 1.0;

	  qual = 1.0 / pow(qq, 2.0);

	  //  	if(qq < 0.4)
	  //  	  qual = 100.0;

/* 	  cout << "mp: " << cc.midpoint() << endl; */
	  // 	cout << "qual: " << qual << endl;

	  maxqual = std::max(maxqual, qual);
	  minqual = std::min(minqual, qual);

	  q.set(c->index(), qual);
	}
      }
    
      void eval(real* values, const real* x) const
      {
	//int d = cell().dim();

	real val = 1.0;

	real qual = q.get(cell().index());
	//real scale = 1.0 * maxqual / minqual;
	real scale = 1.0;
	val = qual / scale;
// 	val = 1.0;

	/*       cout << "qual: " << val << endl; */

	values[0] = val;
      }

      MeshFunction<real> q;
      real maxqual;
      real minqual;
    };
    
    static real initial_B(Cell& cell, MeshFunction<real>& h0, int i)
    {
      EquiAffineMap map;
      map.update(cell);
      
      real h = h0.get(cell.index());
      real scale = h;
      
      uBlasDenseMatrix Finv = map.C * scale;
      
//       cout << "C: " << endl;
//       map.C.disp();
//       cout << "scale: " << scale << endl;
      
      uBlasDenseMatrix B = ublas::prod(ublas::trans(Finv), Finv);

      int d = cell.dim();

      int ii = i % d;
      int jj = i / d;

      return B(ii, jj);
    }

    class InitialValue : public Function
    {
    public:
      InitialValue(Mesh& mesh, MeshFunction<real>& h0) :
	Function(mesh), h0(h0)
      {
      }
      
      void eval(real* values, const real* x) const
      {
	real val = 0.0;

	//FIXME: should we have to cast?
	Cell& c = const_cast<Cell&>(cell());

	int d = c.dim();
	int N = 2 * d + d * d;

/* 	cout << "d: " << d << endl; */

	EquiAffineMap map;
	map.update(c);

	real h = h0.get(c.index());
	real scale = h;

	uBlasDenseMatrix Finv = map.C * scale;

/*  	cout << "C: " << endl; */
/*  	map.C.disp(); */
/*  	cout << "scale: " << scale << endl; */

        uBlasDenseMatrix B = ublas::prod(ublas::trans(Finv), Finv);

	for(int i = 0; i < N; i++)
	{
	  if(i < d)
	    val = x[i];
	  else if(i >= 2 * d)
	  {
/* 	    cout << "i: " << i << endl; */
	    
	    int ii = (i - 2 * d) % d;
	    int jj = (i - 2 * d) / d;
	    
/* 	    cout << "ii: " << ii << endl; */
/* 	    cout << "jj: " << jj << endl; */
	    
	    val = B(ii, jj);
	    
/* 	    cout << "F(" << i << "): " << val << endl; */
	  }
	  values[i] = val;
	}
      }

      real mu;
      real lmbda;
      MeshFunction<real>& h0;
    };
    
    class MyBC : public Function
    {
    public:
    MyBC(Mesh& mesh) : Function(mesh)
      {
      }
      
      void eval(real* values, const real* x) const
      {
	int d = mesh().topology().dim();
	
	for(int i = 0; i < d; i++)
	{
	  values[i] = 0.0;
	}
      }
    };
    
    class ElasticityPDE : public TimeDependentPDE
    {
    public:
      ElasticityPDE(Form* a, Form* L, Form* aJac,
		    Form* aP, Form* LP,
		    Form* aS, Form* LS,
		    Mesh& mesh,
		    BoundaryCondition& bc, real T,
		    Function& U,
		    Function& B,
		    Function& B0,
		    real kk,
		    MeshFunction<real>& h0) :
      TimeDependentPDE(*a, *L, mesh, bc, T, aJac),
	aS(aS), LS(LS),
	aP(aP), LP(LP),
	U(U), B(B), B0(B0), hmin(0.0), kk(kk), h0(h0),
	qual(mesh)
      {
	sampleperiod = T / 10.0;
	B.init(mesh, Bx, *aS, 0);
/* 	cout << "Bx.size: " << Bx.size() << endl; */
	B0.init(mesh, B0x, *aS, 0);
	
	//     B.init(mesh, Bx, *L, 2);
	//     B0.init(mesh, B0x, *L, 2);
	
	// X is compatible with U
	X.init(mesh, Xx, *a, 0);
	//X0x.init(X.vector().size());
	X0.init(mesh, X0x, *a, 0);

	int d = mesh.topology().dim();
	
	vxvalues = new real[d * mesh.numVertices()];
      }

      ~ElasticityPDE()
      {
	delete [] vxvalues;
      }
      
      void computeX()
      {
	X.vector().vec().copy(U.vector().vec(), 0, 0, X.vector().size());
	
	X.vector().vec() *= kk;
	X.vector().vec() += X0x.vec();
      }
      
      void computeB()
      {
	tic();
	
	set("output destination", "silent");
	assembler.assemble(BM, *aS);
	BM.mat().lump(Bm.vec());

	assembler.assemble(Btmp, *LS);
	set("output destination", "terminal");
	
	//     cout << "dotB: " << endl;
	//     tmp.disp();
	
	B.vector().vec().copy(Btmp.vec(), 0, 0, B.vector().size());
	
	B.vector().vec().div(Bm.vec());
	
	B.vector().vec() *= kk;
	B.vector().vec() += B0.vector().vec();
	message("computeB took %g seconds",toc());
      }
      
      void shift()
      {
	X0.vector().vec().copy(X.vector().vec(), 0, 0,
			       X0.vector().vec().size());
	B0.vector().vec().copy(B.vector().vec(), 0, 0,
			       B0.vector().vec().size());
      }
      
      void preparestep()
      {
	shift();
      }
      
      void prepareiteration()
      {
	// FIXME: BC doesn't seem to always be satisfied, why not?
// 	cout << "Ubefore: " << endl;
// 	U.vector().vec().disp();
	set("output destination", "silent");
	bc().apply(M, U.vector(), a());
	set("output destination", "terminal");
// 	cout << "Uafter: " << endl;
// 	U.vector().vec().disp();
	computeX();
	computeB();
	deform(mesh(), X);
      }
      
      void fu(const uBlasVector& x, uBlasVector& dotx, real T)
      {
// 	cout << "X: " << endl;
// 	X.vector().vec().disp();
// 	cout << "U: " << endl;
// 	U.vector().vec().disp();
// 	cout << "B: " << endl;
// 	B.vector().vec().disp();
	
	set("output destination", "silent");
	assembler.assemble(dotx, L());
	bc().apply(M, dotx, a());
	set("output destination", "terminal");
	
// 	cout << "dtU: " << endl;
// 	dotx.disp();
      }
      
      void save(Function& U, real t)
      {
	cout << "Saving at t: " << t << endl;
	//     Function U_0 = U[0];

	//     cout << "U: " << endl;
	//     U.vector().disp();

	if(t == 0.0)
	{
	  solutionfile << U;
	}
	  
	while(lastsample + sampleperiod < t)
	{
	  cout << "saved: " << endl;
	  lastsample = std::min(t, lastsample + sampleperiod);
	  solutionfile << U;

	  //MeshQuality qual(mesh());
      
	  real mu_max, mu_min, mu_avg;
	  real h_max, h_min, h_avg;
      
	  qual.meshQuality(mu_min, mu_max, mu_avg, h_min, h_max, h_avg);
      
	  real vol_min = 1.0e12;
	  real vol_max = 0.0;
      
	  for (CellIterator c(mesh()); !c.end(); ++c)
	  {
	    Cell& cell = *c;
	
	    vol_min = std::min(vol_min, cell.volume());
	    vol_max = std::max(vol_max, cell.volume());
	  }
      
      
	  cout << "FSIc: Mesh quality:" << endl;
	  cout << "FSIc: mu_min: " << mu_min << endl;
	  cout << "FSIc: mu_max: " << mu_max << endl;
	  cout << "FSIc: mu_avg: " << mu_avg << endl;
	  cout << "FSIc: h_min: " << h_min << endl;
	  cout << "FSIc: h_max: " << h_max << endl;
	  cout << "FSIc: h_avg: " << h_avg << endl;
	  cout << "FSIc: vol_min: " << vol_min << endl;
	  cout << "FSIc: vol_max: " << vol_max << endl;
      
	}
      }

      void u0(uBlasVector& x)
      {
	cout << "ElasticityPDE::u0" << endl;

	int d = mesh().topology().dim();
	int N = mesh().numVertices();
	int M = mesh().numCells();

	real* Xarr = X.vector().vec().array();

	for (VertexIterator n(mesh()); !n.end(); ++n)
	{
	  for(int i = 0; i < d; i++)
	  {
	    Xarr[i * N + n->index()] = n->x()[i];
	  }
	}

	U.vector().zero();

	real* Barr = B.vector().vec().array();

	for (CellIterator c(mesh()); !c.end(); ++c)
	{
	  for(int i = 0; i < d * d; i++)
	  {
	    real val = initial_B(*c, h0, i);
	    
	    Barr[i * M + c->index()] = val;
	  }
	}
// 	// Compute initial value
// 	Function Pf = pproject(*aP, *LP, mesh());

// 	U.vector().vec().copy(Pf.vector().vec(), 0, 0, U.vector().size());

// 	cout << "u0: " << endl;
// 	Pf.vector().disp();
	
// 	// Copy into components
// 	X.vector().vec().copy(Pf.vector().vec(), 0, 0, X.vector().size());
// 	U.vector().vec().copy(Pf.vector().vec(), 0, U.vector().size(),
// 			      U.vector().size());
// 	B.vector().vec().copy(Pf.vector().vec(), 0, 2*U.vector().size(),
// 			      B.vector().size());

	
	
	//FIXME: shift() seems to be needed for cG(1) but not dG(0), why?
	shift();
	
// 	cout << "X0: " << endl;
// 	X.vector().disp();
// 	cout << "x0: " << endl;
// 	x.disp();
// 	cout << "B0: " << endl;
// 	B.vector().disp();
      }

      bool update(const uBlasVector& u, real t, bool end)
      {
	real mu_max, mu_min, mu_avg;
	real h_max, h_min, h_avg;
	
	qual.meshQuality(mu_min, mu_max, mu_avg, h_min, h_max, h_avg);

	real sstate_tol = 1.0e-2;

	real norm = dotx->norm(linf);
	cout << "dotx norm: " << norm << endl;

	real xnorm = x->norm(linf);
	cout << "x norm: " << xnorm << endl;

	cout << "hmin: " << h_min << endl;

	real xmin = std::max(x->norm(linf), 1.0e-6);

	real k0 = 1.0 / 2.0 * h_min;

	// Compute k according to k < h / |U|
	if(true || t == 0.0)
	  k = k0;
	else
	  k = 0.25 * h_min / xmin;

	k = std::min(k, 0.1);
	
	cout << "k: " << k << endl;

	if(norm < sstate_tol)
	  return false;
	return true;
      }

      void deform(Mesh& mesh, Function& u);

      Form* aS;
      Form* LS;
      Form* aP;
      Form* LP;

      Function& U;
      Function& B;
      Function& B0;
      Function X;
      Function X0;
      Vector B0x, Bx, Btmp;
      Matrix BM;
      Vector Bm;
      Vector X0x, Xx;
      
      
      real* vxvalues;

      real hmin;
      real kk;
      MeshFunction<real>& h0;

      MeshQuality qual;
    };
  };




}}

#endif
