// Copyrightx (C) 2005 Johan Hoffman.
// Licensed under the GNU GPL Version 2.
//
// Modified by Garth N. Wells 2005.
// Modified by Anders Logg 2005-2006.
//
// First added:  2005
// Last changed: 2006-10-19

#include <sstream>
#include <vector>
#include <algorithm>


#include <dolfin.h>
#include <dolfin/timing.h>
#include <dolfin/MeshFunction.h>
#include <dolfin/Assembler.h>
#include <dolfin/PETScMatrix.h>
#include <dolfin/PETScVector.h>

#include "unicorn/NSESolver.h"
#include "unicorn/NSEMomentum3D.h"
#include "unicorn/NSEMomentum2D.h"
#include "unicorn/NSEContinuity3D.h"
#include "unicorn/NSEContinuity2D.h"
#include "unicorn/NSEResMomentum2D.h"
#include "unicorn/NSEResMomentum3D.h"
#include "unicorn/NSEResContinuity2D.h"
#include "unicorn/NSEResContinuity3D.h"

using namespace dolfin;
using namespace unicorn;

//-----------------------------------------------------------------------------
NSESolver::NSESolver(Mesh& mesh, Function& f,
		     Array<BoundaryCondition*>& bc_mom, 
		     BoundaryCondition& bc_con)
  : mesh(mesh), f(f), bc_mom(bc_mom), bc_con(bc_con) 
{
  // Declare parameters
//   add("velocity file name", "velocity.pvd");
//   add("pressure file name", "pressure.pvd");
//   add("res_m file name", "res_m.pvd");
//   add("res_c file name", "res_c.pvd");
//   add("e_ind file name", "e_ind.pvd");
}
//-----------------------------------------------------------------------------
void NSESolver::solve()
{
  // Get the number of space dimensions of the problem 
  int nsd = mesh.topology().dim();

  message("Number of space dimensions: %d",nsd);

  real T0 = 0.0;        // start time 
  real t  = 0.0;        // current time
  real T  = 8.0;        // final time
  real nu = 1.0/3900.0; // viscosity 

  // Set time step (proportional to the minimum cell diameter) 
  real hmin;
  GetMinimumCellSize(mesh, hmin);  
  real k = 0.25*hmin; 
  //real k = 0.05*hmin; 
  message("k: %f",k);


  Function u; // velocity from previous time step 
  Function u0; // velocity from previous time step 
  Function uc; // velocity linearized convection 
  Function um; // cell mean velocity linearized convection 
  Function p;   // current pressure
  Function vol_inv;  // 1 / volume of element
  Function res_m;  // 1 / volume of element
  Function res_c;  // 1 / volume of element
  Function e_ind;  // 1 / volume of element
  Function delta1, delta2;

  Function fk(mesh, k);
  Function fnu(mesh, nu);

  // Initialize the bilinear and linear forms
  Form* amom = 0;
  Form* acon = 0;
  Form* Lmom = 0;
  Form* Lcon = 0;
  Form* Lres_m = 0;
  Form* Lres_c = 0;

  if ( nsd == 3 )
  {
    amom = new NSEMomentum3DBilinearForm(um,delta1,delta2,fk,fnu);
    Lmom = new NSEMomentum3DLinearForm(um,u0,f,p,delta1,delta2,fk,fnu);
    acon = new NSEContinuity3DBilinearForm(delta1);
    Lcon = new NSEContinuity3DLinearForm(uc);

    Lres_m   = new NSEResMomentum3DLinearForm(p, u, u0, vol_inv, fk);
    Lres_c   = new NSEResContinuity3DLinearForm(u, u0, vol_inv);
  } 
  else if ( nsd == 2 )
  {
    amom = new NSEMomentum2DBilinearForm(um,delta1,delta2,fk,fnu);
    Lmom = new NSEMomentum2DLinearForm(um,u0,f,p,delta1,delta2,fk,fnu);
    acon = new NSEContinuity2DBilinearForm(delta1);
    Lcon = new NSEContinuity2DLinearForm(uc);

    Lres_m   = new NSEResMomentum2DLinearForm(p, u, u0, vol_inv, fk);
    Lres_c   = new NSEResContinuity2DLinearForm(u, u0, vol_inv);
  }
  else
  {
    error("Navier-Stokes solver only implemented for 2 and 3 space dimensions.");
  }





  // Create matrices and vectors 
  Matrix Amom, Acon;
  Vector bmom, bcon;

  // Initialize vectors for velocity and pressure 
  // x0vel: velocity from previous time step 
  // xcvel: linearized velocity 
  // xvel:  current velocity 
  // pvel:  current pressure 
  Vector x0vel;
  Vector xcvel;
  Vector xmvel;
  Vector xvel;
  Vector xpre;

  Vector vol_invx;         // vol_invx: needed for the computing strong residual                       
  Vector res_mx;     // res_mx: needed for storing the momentum residual                            
  Vector res_cx;     // res_mx: needed for storing the momentum residual                            
  Vector e_indx;     // Error indicator (on each cell)

  Vector d1vector;
  Vector d2vector;

  // Initialize vectors for the time step residuals of 
  // the momentum and continuity equations  
  Vector residual_mom(nsd*mesh.numVertices());
  Vector residual_con(mesh.numVertices());

  // Initialize algebraic solvers 
  KrylovSolver solver_con(gmres, amg);
  KrylovSolver solver_mom(gmres);
 
  u.init(mesh, xvel, *amom, 1);
  u0.init(mesh, x0vel, *Lmom, 2);
  uc.init(mesh, xcvel, *Lcon, 1);
  um.init(mesh, xmvel, *amom, 2);
  p.init(mesh, xpre, *Lmom, 4);
  vol_inv.init(mesh, vol_invx, *Lres_m, 4);
  res_m.init(mesh, res_mx, *Lres_m, 4);
  res_c.init(mesh, res_cx, *Lres_c, 3);
  e_ind.init(mesh, e_indx, *Lres_c, 3);
  delta1.init(mesh, d1vector, *amom, 3);
  delta2.init(mesh, d2vector, *amom, 4);

  // Initialize values
  x0vel.zero();
  xcvel.zero();
  xmvel.zero();
  xvel.zero();
  xpre.zero();
  e_indx.zero();

  cout << "uc.size(): " << uc.vector().size() << endl;
  cout << "um.size(): " << um.vector().size() << endl;

  // Set initial velocity
  SetInitialVelocity(xvel);

  Assembler assembler(mesh);

  // Synchronise functions and boundary conditions with time
//   u.sync(t);
//   p.sync(t);
//   bc_con.sync(t);
//   bc_mom.sync(t);
    
  // Initialize output files 
  File file_u("velocity.pvd");  // file for saving velocity 
  File file_p("pressure.pvd");  // file for saving pressure
  File file_resm("res_m.pvd");          // file for saving pressure
  File file_resc("res_c.pvd");          // file for saving pressure
  File file_eind("e_ind.pvd");          // file for saving pressure

  // Compute stabilization parameters
  ComputeStabilization(mesh,u0,nu,k,d1vector,d2vector);

  // Compute the volume inverse of an element
  ComputeVolInv(mesh, vol_invx);

  message("Assembling matrix: continuity");

  // Assembling matrices 
  assembler.assemble(Acon, *acon);
  assembler.assemble(Amom, *amom);

  assembler.assemble(res_mx, *Lres_m);
  assembler.assemble(res_cx, *Lres_c);

  file_resm << res_m;
  file_resc << res_c;


  // Initialize time-stepping parameters
  int time_step = 0;
  int sample = 0;
  int no_samples = 50;

  // Residual, tolerance and maxmimum number of fixed point iterations
  real residual;
  real rtol = 1.0e-2;
  int iteration;
  int max_iteration = 50;

  // Start time-stepping
  Progress prog("Time-stepping");
  while (t<T) 
  {

    time_step++;
    message("Time step %d",time_step);

    // Set current velocity to velocity at previous time step 
    x0vel = xvel;

    // Initialize residual 
    residual = 2*rtol;
    iteration = 0;

    // Fix-point iteration for non-linear problem 
    while (residual > rtol && iteration < max_iteration){
      
      // Compute stabilization parameters
      tic();
      ComputeStabilization(mesh,u0,nu,k,d1vector,d2vector);
      message("Compute stab took %g seconds",toc());


      message("Assemble vector: continuity");

      // Assemble continuity vector 
      assembler.assemble(bcon, *Lcon);

      // Set boundary conditions for continuity equation 
      //assembler.applyBC(Acon, bcon, mesh, acon->trial(), bc_con);
      bc_con.apply(Acon, bcon, *acon);

      message("Solve linear system: continuity");

      // Solve the linear system for the continuity equation 
      tic();
      solver_con.solve(Acon, xpre, bcon);
      message("Linear solve took %g seconds",toc());

      message("Assemble vector: momentum");

      // Assemble momentum vector 
      tic();
      assembler.assemble(bmom, *Lmom);
      message("Assemble took %g seconds",toc());

      // Set boundary conditions for the momentum equation 
      //assembler.applyBC(Amom, bmom, mesh, amom->trial(),bc_mom);
     
      // Set boundary conditions 
      for (uint i = 0; i < bc_mom.size(); i++)
	bc_mom[i]->apply(Amom, bmom, *amom);


      message("Solve linear system: momentum");

      // Solve the linear system for the momentum equation 
      tic();
      solver_mom.solve(Amom, xvel, bmom);
      message("Linear solve took %g seconds",toc());
      
      // Set linearized velocity to current velocity 
      xcvel = xvel;
      // Compute cell mean
      ComputeMean(mesh,um,uc);
    
      message("Assemble matrix and vector: momentum");
      tic();
      assembler.assemble(Amom, *amom);
      assembler.assemble(bmom, *Lmom);
      //assembler.applyBC(Amom, bmom, mesh, amom->trial(),bc_mom);

      // Set boundary conditions 
      for (uint i = 0; i < bc_mom.size(); i++)
	bc_mom[i]->apply(Amom, bmom, *amom);

      bc_con.apply(Acon, bcon, *acon);
      message("Assemble mom took %g seconds",toc());

      // Compute residual for momentum equation 
      Amom.mat().mult(xvel.vec(),residual_mom.vec());
      residual_mom.vec() -= bmom.vec();
      
      message("Assemble vector: continuity");
      assembler.assemble(bcon, *Lcon);
      //assembler.applyBC(Acon, bcon, mesh, acon->trial(),bc_con);
      bc_con.apply(Acon, bcon, *acon);

      // Compute residual for continuity equation 
      Acon.mat().mult(xpre.vec(),residual_con.vec());
      residual_con.vec() -= bcon.vec();
      
      message("Momentum residual  : l2 norm = %e",residual_mom.norm());
      message("continuity residual: l2 norm = %e",residual_con.norm());
      message("Total NSE residual : l2 norm = %e",sqrt(sqr(residual_mom.norm()) + sqr(residual_con.norm())));

      residual = sqrt(sqr(residual_mom.norm()) + sqr(residual_con.norm()));
      iteration++;
    }
    assembler.assemble(res_mx, *Lres_m);
    assembler.assemble(res_cx, *Lres_c);

    ComputeErrorIndicator(e_indx, res_mx, res_cx, k, T);

    if(residual > rtol)
      warning("NSE fixed point iteration did not converge"); 



    if ( (time_step == 1) || (t > (T-T0)*(real(sample)/real(no_samples))) ){
      message("save solution to file");
      file_p << p;
      file_u << u;

      ////file_resm << res_m;
      //file_resc << res_c;
      // Note: just for debugging
      //file_eind << e_ind;

      std::stringstream filename, number;
    
      number << std::setfill('0') << std::setw(6) << sample;
      filename << "solution" << number.str() << ".xml.gz" << std::ends;

      File xmlfile(filename.str());
      
      xmlfile << u.vector();
      

      sample++;
    }

    // Increase time with timestep
    t = t + k;

    // Update progress
    prog = t / T;
  }

  std::vector<int> cells;
  //ComputeLargestIndicators(e_indx, cells, 1.0);
  //AdaptiveRefinement(cells, acon->trial());

  message("save solution to file");
  file_p << p;
  file_u << u;
  file_eind << e_ind;

  delete amom;
  delete Lmom;
  delete acon;
  delete Lcon;
}
//-----------------------------------------------------------------------------
void NSESolver::solve(Mesh& mesh, Function& f,
		      Array<BoundaryCondition*>& bc_mom, 
		      BoundaryCondition& bc_con)
{
  NSESolver solver(mesh, f, bc_mom, bc_con);
  solver.solve();
}
//-----------------------------------------------------------------------------
void NSESolver::ComputeCellSize(Mesh& mesh, Vector& hvector)
{
  real* harr = hvector.vec().array();

  // Compute cell size h
  hvector.init(mesh.numCells());	
  for (CellIterator cell(mesh); !cell.end(); ++cell)
    {
      harr[(*cell).index()] = (*cell).diameter();
    }

  hvector.vec().restore(harr);
}
//-----------------------------------------------------------------------------
void NSESolver::GetMinimumCellSize(Mesh& mesh, real& hmin)
{
  // Get minimum cell diameter
  hmin = 1.0e6;
  for (CellIterator cell(mesh); !cell.end(); ++cell)
    {
      if ((*cell).diameter() < hmin) hmin = (*cell).diameter();
    }
}
//-----------------------------------------------------------------------------
void NSESolver::ComputeStabilization(Mesh& mesh, Function& w, real nu, real k, 
				     Vector& d1vector, Vector& d2vector)
{
  // Compute least-squares stabilizing terms: 
  //
  // if  h/nu > 1 or ny < 10^-10
  //   d1 = C1 * ( 0.5 / sqrt( 1/k^2 + |U|^2/h^2 ) )   
  //   d2 = C2 * h 
  // else 
  //   d1 = C1 * h^2  
  //   d2 = C2 * h^2  

  real C1 = 4.0;   
  real C2 = 2.0;   

//   d1vector.init(mesh.numCells());	
//   d2vector.init(mesh.numCells());	

  real normw; 

  int nsd = mesh.topology().dim(); 

  real* d1arr = d1vector.vec().array();
  real* d2arr = d2vector.vec().array();
  real* warr = w.vector().vec().array();

  for (CellIterator cell(mesh); !cell.end(); ++cell)
    {
      normw = 0.0;
      for (VertexIterator n(*cell); !n.end(); ++n)
	for ( int i = 0; i < nsd; i++ )
	  normw += sqr(warr[i*mesh.numVertices() + (*n).index()]); 
      
      normw = sqrt(normw);
  
      normw /= (*cell).numEntities(0);
  
      real h = (*cell).diameter();
      //real h = MeshQuality::myDiameter(*cell);

      uint cid = (*cell).index();
      
      if (((h/nu) > 1.0) || (nu < 1.0e-10) ){
	d1arr[cid] = C1 * (0.5 / sqrt( 1.0/sqr(k) + sqr(normw/h)));
	d2arr[cid] = C2 * h;
      } else{
	d1arr[cid] = C1 * sqr(h);
	d2arr[cid] = C2 * sqr(h);
      }	
    }

  d1vector.vec().restore(d1arr);
  d2vector.vec().restore(d2arr);
}
//-----------------------------------------------------------------------------
void NSESolver::ComputeMean(Mesh& mesh, Function& vmean, Function& v)
{
  int nsd = mesh.topology().dim(); 

  real* varr = v.vector().vec().array();
  real* vmeanarr = vmean.vector().vec().array();

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

      for ( int i = 0; i < nsd; i++ )
	{
	  real cellmean = 0.0;
	  for (VertexIterator n(cell); !n.end(); ++n)
	    cellmean += varr[i*mesh.numVertices() + (*n).index()];
	  
	  cellmean /= cell.numEntities(0);
      
	  vmeanarr[i*mesh.numCells() + cell.index()] = cellmean;
	}
    }
}
//-----------------------------------------------------------------------------
void NSESolver::SetInitialVelocity(Vector& xvel)
{
  // Function for setting initial velocity, 
  // given as a function of the coordinates (x,y,z).
  //
  // This function is only temporary: initial velocity 
  // should be possible to set in the main-file. 


//  real x,y,z;

  real* xvelarr = xvel.vec().array();

  for (VertexIterator vertex(mesh); !vertex.end(); ++vertex)
  {
    // Get coordinates of the vertex 
//    real x = (*vertex).coord().x;
//    real y = (*vertex).coord().y;
//    real z = (*vertex).coord().z;
    
    // Specify the initial velocity using (x,y,z) 
    //
    // Example: 
    // 
    // xvel((*vertex).index()) = sin(x)*cos(y)*sqrt(z);
    
    xvelarr[(*vertex).index()] = 0.0;
  }

  xvel.vec().restore(xvelarr);
}
//-----------------------------------------------------------------------------                                              
void NSESolver::ComputeVolInv(Mesh& mesh, Vector& vol_inv)
{
  real* varr = vol_inv.vec().array();

  // Compute cell size h                                                                                                     
  vol_inv.init(mesh.numCells());
  for (CellIterator cell(mesh); !cell.end(); ++cell)
    {
      varr[(*cell).index()] = 1.0 / (*cell).volume();  // (*cell).diameter();                                             
    }

  vol_inv.vec().restore(varr);
}

//-----------------------------------------------------------------------------
void NSESolver::ComputeErrorIndicator(Vector& e_indx, Vector& res_mx, Vector& res_cx, real k, real T)
{
//   int N = mesh.numCells();



//   // Compute error indicator using end-point quadrature in time
//   for (CellIterator cell(mesh); !cell.end(); ++cell)
//     {
//       // FIXME: Should we use the Euclidean norm for the momentum residual?
//       // FIXME: We assume the stability factor is 1, need to compute dual problem
//       // FIXME: for computing it correctly.

//       int id = (*cell).index();
//       real h = (*cell).diameter();

//       real ei = e_indx(id);
//       ei += 1.0 / T * k * h * (sqrt(sqr(res_mx(id + 0 * N)) + sqr(fabs(res_mx(id + 1 * N)))) +
// 		 fabs(res_cx(id)));
//       e_indx(id) = ei;
//     }
}
//-----------------------------------------------------------------------------
void NSESolver::ComputeLargestIndicators(Vector& e_ind, std::vector<int>& cells, real percentage)
{
//   int N = e_ind.size();
//   int M = (int)((real)N * percentage * 0.01);

//   cells.resize(M);

//   std::vector<std::pair<int, real> > indicators(N);

//   for (CellIterator cell(mesh); !cell.end(); ++cell)
//     {
//       int id = (*cell).index();
//       std::pair<int, real> p;
//       p.first = id;
//       p.second = e_ind(id);

//       indicators[id] = p;
//     }

//   less_pair comp;

//   std::sort(indicators.begin(), indicators.end(), comp);

//   for(int i = 0; i < M; i++)
//     {
//       std::pair<int, real> p = indicators[N - 1 - i];
//       cells[i] = p.first;

//       cout << "cell id: " << p.first << endl;
//       cout << "error: " << p.second << endl;
//     }
}
//-----------------------------------------------------------------------------
void NSESolver::AdaptiveRefinement(std::vector<int>& cells, FiniteElement& element)
{
//   Mesh refined_mesh(mesh);
  
//   MeshFunction<bool> cell_refinement_marker(refined_mesh);
//   cell_refinement_marker.init(refined_mesh.topology().dim());

//   int M = cells.size();

//   for (CellIterator c(refined_mesh); !c.end(); ++c)
//     {
//       cell_refinement_marker.set(c->index(), false);
//     }

//   for(int i = 0; i < M; i++)
//     {
//       cell_refinement_marker.set(cells[i], true);
//     }

//   refined_mesh.refine(cell_refinement_marker);

//   File meshfile("refined.xml");
//   meshfile << refined_mesh;

// //   Function zero(refined_mesh, element);
// //   File vtkfile("refined.pvd");
// //   vtkfile << zero;
}
//-----------------------------------------------------------------------------
