// Copyright (C) 2005 Johan Hoffman.
// Licensed under the GNU GPL Version 2.
//
// Modified by Anders Logg 2006.
//
// First added:  2005
// Last changed: 2006-05-07


// PETSc friendly version

#ifndef __FSI_SOLVER_H
#define __FSI_SOLVER_H

#include <float.h>

//#include <unicorn/Solver.h>

#include <dolfin.h>

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

// #include <dolfin/BoundaryMesh.h>

namespace dolfin { 
  namespace unicorn
  {
    /// Fluid-Structure Interaction Solver
    class FSISolver 
    {
    public:
      
      /// Create the FSI solver
      FSISolver(Mesh& mesh, 
		Function& f, 
		Function& ic,
		DirichletBC& bc_mom, 
		DirichletBC& bc_con, 
		Function& bcf_mom, 
		Function& bcf_con, 
		Function& bisect, 
		real& rhof, 
		real& rhos, 
		real& E, 
		real& elnu, 
		real& nu,
		real& T,
		real& k,
		real& C_cfl,
		real& C_k);
      
      /// FSI solver destructor
      ~FSISolver();
      
      /// Solve FSI equations
      void solve();
      
      /// Solve FSI equations (static version)
      static void solve(Mesh& mesh, 
			Function& f, 
			Function& ic,
			DirichletBC& bc_mom, 
			DirichletBC& bc_con, 
			Function& bcf_mom, 
			Function& bcf_con, 
			Function& bisect, 
			real& rhof, 
			real& rhos, 
			real& E, 
			real& elnu,
			real& nu,
			real& T,
			real& k,
			real& C_cfl,
			real& C_k);
      
      /// Compute cell diameter
      void ComputeCellSize(Mesh& mesh, Vector& hvector);
      
      /// Get minimum cell diameter
      void GetMinimumCellSize(Mesh& mesh, real& hmin);
      
      /// Compute stabilization 
      void ComputeStabilization(Mesh& mesh, 
				Function& w, 
				real nu, 
				real k, 
				Vector& d1vector, 
				Vector& d2vector,
				real& t);
      
      /// Set initial velocity 
      void SetInitialVelocity(Vector& xvel, Function& ic);

      /// Compute cell mean of v
      void ComputeMean(Mesh& mesh, Function& vmean, Function& v);
      
      
    private:
      
      /// Compute the marker vector phi
      void ComputePhi(Vector& phi);
      
      /// Compute the partial stress alpha
      void ComputeAlpha(Vector& alpha, Vector& phi, real& nu);
      
      /// Compute the density vector rho
      void ComputeDensity(Vector& rho, Vector& phi);
      
      /// Reassemble Dot Sigma
      void UpdateDotSigma(Vector& dsig, 
			  Vector& icell_vol,
			  Form& L);
      
      /// Returns true if vertex neighboring cells are all fluid
      bool isFluid(Vertex& vertex, const Vector& phi_vec) const;
      
      /// Returns true if vertex is not in boundary
      bool isInterior(Vertex& vertex) const;
      
      /// Returns vector component (i,j)
      unsigned int getij(unsigned int j, unsigned int i) const;
      
      /// Move the structure part of the mesh according to velocity u
      void UpdateStructure(Vector& mvel_vec, const Vector& phi_vec, Function& u);
      /// Smooth the fluid mesh
      void UpdateFluid(Vector& mvel_vec, const Vector& phi_vec);

      /// Smooth the fluid mesh
      void SmoothMesh();
      
      /// Compute cell volume
      void ComputeCellVolume(Vector& icell_vol);
      
      /// Compute step size k
      void ComputeStepSize(real& k, const Vector& vel, int& style);
      void CFLStepSize(real& k, const Vector& vel);
      void CellStepSize(real& k);
      
      
      Mesh& mesh;
      Function& f;
      Function& ic;
      DirichletBC& bc_mom;
      DirichletBC& bc_con;
      Function& bcf_mom;
      Function& bcf_con;
      Function& bisect;
      real& rhof;
      real& rhos;
      real& E;
      real& elnu;
      real& nu;
      real& T;
      real& k;
      real& C_cfl;
      real& C_k;
      
      MeshFunction<unsigned int> boundary_vertex_map;
      MeshFunction<unsigned int> boundary_cell_map;
      MeshFunction<bool>         vertex_is_interior;
      BoundaryMesh*              boundary;
      Assembler		   assembler;

      MeshFunction<bool> masked_cells;
      MeshFunction<bool> smoothed_cells;
      MeshFunction<real> h0;

      MeshQuality* mqual;
      ElasticSmoother smoother;

      real* phi_vxvalues;
      real* u_vxvalues;
    };
    //-----------------------------------------------------------------------------
    // INLINE: Function definitions - these functions are here because
    //                                they are used in every time step
    //-----------------------------------------------------------------------------
    inline bool FSISolver::isFluid(Vertex& vertex, const Vector& phi_vec) const
    {
      const real* z = phi_vec.vec().array();
      
      for (CellIterator cell(vertex); !cell.end(); ++cell) 
	if (z[cell->index()] == 0) return false;
      
      return true;
    }
    //-----------------------------------------------------------------------------
    inline bool FSISolver::isInterior(Vertex& vertex) const
    {  
      return (vertex_is_interior(vertex));
    }
    //-----------------------------------------------------------------------------
    inline unsigned int FSISolver::getij(unsigned int i, unsigned int j) const
    {
      return (i + mesh.topology().dim()*j);
    }
    //-----------------------------------------------------------------------------
    inline void FSISolver::UpdateStructure(Vector& mvel_vec, const Vector& phi_vec, Function& u) 
    {
      unsigned int d = mesh.topology().dim();
      unsigned int N = mesh.numVertices();

      real* w = mvel_vec.vec().array();
      real value;

      u.interpolate(u_vxvalues);

      for (VertexIterator v(mesh); !v.end(); ++v) {

	for (unsigned int ndx = 0; ndx < d; ndx++) {

	  unsigned int ai = ndx * N + v->index();
					
	  value = u_vxvalues[ai];

	  // CHECK u might be really small, or not
//  	  cout << "FSIc: u(" << ndx <<") of vertex(" << v->index() << ") is " << value << endl;
	  // END CHECK

	  if (isFluid(*v, phi_vec)) continue;

	  v->x()[ndx] += k * value;
		
	  w[ai] =  w[ai] + k * value; //u(*v, ndx);
	}
      }
    }
    //-----------------------------------------------------------------------------
    
    //Deprecated: use elastic_smooth

    inline void FSISolver::UpdateFluid(Vector& mvel_vec, const Vector& phi_vec)
    {
      // 		real old_coord[3];
      // 			unsigned int ndx;

      // 			for (VertexIterator v(mesh); ! v.end(); ++v) {  
      
      // 				if (!isFluid(*v, phi_vec)) continue;
      // 				if (!isInterior(*v))       continue;
      
      // 				for (ndx = 0; ndx < mesh.topology().dim(); ndx++) {
	
      // 					real mass = 1;
	
      // 					// move boundary point
      // 					old_coord[ndx] = v->coordinates()[ndx]; 
	
      // 					// iterate over the neighboring vertices
      // 					VertexIterator vn(*v); 

      // 					v->coordinates()[ndx] = vn->coordinates()[ndx];
	
      // 					for (++vn; !vn.end(); ++vn) {
      // 						v->coordinates()[ndx] += vn->coordinates()[ndx];
      // 						mass += 1;
      // 					}
	
      // 					// divide by the number of neighbors
      // 					v->coordinates()[ndx] /= mass;
	
      // 					// mesh velocity contribution
      // 					mvel_vec.set(getij(v->index(), ndx), mvel_vec.get(getij(v->index(), ndx)) + (v->coordinates()[ndx] - old_coord[ndx]));	
      // 				}	
      // 			}
    }
    //-----------------------------------------------------------------------------
    inline void FSISolver::UpdateDotSigma(Vector& dsig, Vector& cell_vol, Form& L)
    {
      ComputeCellVolume(cell_vol);    // remember that we use cell volume in L-form
      assembler.assemble(dsig, L);
    }
    //-----------------------------------------------------------------------------
    inline void FSISolver::ComputeCellVolume(Vector& icell_vol)
    {
      //icell_vol.init(mesh.numCells());

      real* icv = icell_vol.vec().array();

      for (CellIterator cell(mesh); !cell.end(); ++cell) 
	icv[cell->index()] = 1.0/cell->volume();
    }
    //-----------------------------------------------------------------------------
    inline void FSISolver::ComputeStepSize(real& k, const Vector& vel, int& style)
    {
      switch (style) {
      case -1: 
	CFLStepSize(k, vel);
	cout << "FSIc: CFL stepping ";
	break;
      case 0:
	CellStepSize(k);
	cout << "FSIc: Cell stepping ";
	break;
      case 1:
	// Do nothing. k = k
	cout << "FSIc: Custom stepping ";
	break;
      default:
	error("Unknown time stepping style");
      }

      cout << "k = " << k << endl;
    }
    //-----------------------------------------------------------------------------
//     inline void FSISolver::CFLStepSize(real& k, const Vector& u) 
//     { 
//       real global_min;
//       real local_min;
//       real u_mag;
//       unsigned int ndx;

//       real* value;

//       global_min = FLT_MAX;

//       for (CellIterator c(mesh); !c.end(); ++c) {
// 	for (VertexIterator v(*c); !v.end(); ++v) {
		
// 	  u_mag = 0;
		
// 	  u(value, v->x());

// 	  for (ndx = 0; ndx < mesh.topology().dim(); ndx++) {
// 	    u_mag += sqr(value[ndx]);
// 	  }

// 	  u_mag = sqrt(u_mag);
// 	  local_min  = c->diameter() / u_mag;
// 	  global_min = (local_min < global_min)? local_min : global_min;
// 	}
//       }

//       k = 0.25 * global_min;
//     }
    //-----------------------------------------------------------------------------
    inline void FSISolver::CellStepSize(real& k)
    {
      real hmin;
	
      GetMinimumCellSize(mesh, hmin);
      k = 0.25 * hmin;
    }
    
	


  }
}

#endif
