// Copyright (C) 2007 Murtazo Nazarov.
// Licensed under the GNU GPL Version 2.
// 
// Based on existing NSESolver of Johan Hoffman


#include <iostream>
#include <dolfin/timing.h>
#include "unicorn/CNSSolver.h"
#include "unicorn/CNSmix2D.h"
#include "unicorn/CNSmix3D.h"
#include "unicorn/CNSResDensity2D.h"
#include "unicorn/CNSResDensity3D.h"
#include "unicorn/CNSResMomentum2D.h"
#include "unicorn/CNSResMomentum3D.h"
#include "unicorn/CNSResEnergy2D.h"
#include "unicorn/CNSResEnergy3D.h"


using namespace dolfin;
using namespace unicorn;

//-----------------------------------------------------------------------------
CNSSolver::CNSSolver(Mesh& mesh, Function& f, Function& initial, Array<BoundaryCondition*>& bcs)
  : mesh(mesh), f(f), initial(initial), bcs(bcs) 
{
  // Declare parameters
  add("velocity file name", "velocity.pvd");
  add("pressure file name", "pressure.pvd");
  add("density file name", "density.pvd");
  add("momentum file name", "momentum.pvd");
  add("energy file name", "energy.pvd");
  add("res_den file name", "res_den.pvd");
  add("res_m file name", "res_m.pvd");

  add("nu_den file name", "nu_den.pvd");
  add("mach_num file name", "mach_num.pvd");
}
//-----------------------------------------------------------------------------
void CNSSolver::solve()
{
  real T0 = 0.0;        // start time 
  real t  = 0.0;        // current time
  real T  = 0.3;        // final time
  real gamma = 1.4;

  // Set time step (proportional to the minimum cell diameter) 
  real hmin;
  GetMinimumCellSize(mesh, hmin);  
  real k = 0.5*0.25*hmin; 

  //  Create matrix and vector
  Matrix A;
  Vector b;

  //  Get the number of space dimensions of the problem 
  
  int nsdim = mesh.topology().dim();     // nsdim: is number of space dimensions
  int nsd = nsdim + 2;                   // nsd:       is number of variables

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

  //  Initialize stabilization parameters 
  Function delta;

  //  Create functions for the solution
  //  (needed for the initialization of the forms)
  Function w;               // current solution
  Function w0;              // solution from previous time step
  Function wold;            // solution from previous time step needed for fix-point iteration

  Function P;               // pressure 
  Function u;               // velocity
  Function um;              // cell mean value velocity

  Function rho;             // density 
  Function rho0;            // density from previous timestep

  Function m;               // momentum
  Function m0;              // momentum from previous timeste


  Function e;               // energy
  Function e0;              // energy from previous timestep

  Function vol_inv;         // 1 / volume of element
  Function res_rho;         // residual for density
  Function res_m;           // residual for momentum
  Function res_e;           // residual for energy

  Function nu_rho;          // \nu for density, \nu is shock-capturing term
  Function nu_m;            // \nu for momentum
  Function nu_e;            // \nu for energy

  Function mach_num;        // Mach number 

  Function fk(mesh, k);


  //  Initialize algebraic solvers 
  KrylovSolver solver_con(gmres, amg);
  KrylovSolver solver(gmres);

  // Initialize the bilinear and linear forms for the system
  Form* a = 0;
  Form* L = 0;

  //  Initialize the linear forms for the strong residual of density, momentum and energy
  Form* Lres_rho = 0;
  Form* Lres_m = 0;
  Form* Lres_e = 0;

  FormFactory(mesh, a, L, Lres_rho, Lres_m, Lres_e,
	      w0, u, um, P, rho, rho0, m, m0, e, e0, nu_rho, nu_m, nu_e, delta, vol_inv, fk);
  
  Vector x0;             // x0:   solution from previous timestep
  Vector x;              // x:    current solution
  Vector xold;//(nsd*mesh.numVertices());
                         // xold: save current solution 

  Vector px;             // px:   pressure
  Vector ux;             // ux:   velocity
  Vector umx;            // umx:  cell mean value velocity 
  
  Vector rhox;           // rhox:  density
  Vector rho0x;          // rho0x: density from previous timestep

  Vector mx;             // mx:    momentum
  Vector m0x;            // m0x:   momentum from previous timestep

  Vector ex;             // ex:    energy
  Vector e0x;            // e0x:   energy from previous timestep

  Vector vol_invx;       // vol_invx: needed for the computing strong residual
  Vector res_rhox;       // res_rhox: needed for storing the density residual
  Vector res_mx;         // res_mx: needed for storing the momentum residual
  Vector res_ex;         // res_ex: needed for storing the energy residual

  Vector nu_rhox;        // res_rhox: needed for storing the density residual
  Vector nu_mx;          // res_mx: needed for storing the momentum residual
  Vector nu_ex;          // res_ex: needed for storing the energy residual

  Vector dvector;        // dvector: needed for the stabilazation term
  Vector mach_numx; //(mesh.numVertices());      // machnx: Mach number

  w.init(mesh, x, *a, 1);                      // current solution
  w0.init(mesh, x0, *L, 1);                    // solution from previous time step
  wold.init(mesh, xold, *L, 1);                // solution from previous time step
  
  P.init(mesh, px, *L, 2);                     // pressure 
  u.init(mesh, ux, *Lres_rho, 1);              // ceall mean velocity
  um.init(mesh, umx, *a, 2);                   // velocity
  
  rho.init(mesh, rhox, *Lres_rho, 2);          // density 
  rho0.init(mesh, rho0x, *Lres_rho, 3);        // density from previous timestep
  
  m.init(mesh, mx, *Lres_m, 3);                // momentum
  m0.init(mesh, m0x, *Lres_m, 4);              // momentum from previous timeste
  
  e.init(mesh, ex, *Lres_e, 3);                // energy
  e0.init(mesh, e0x, *Lres_e, 4);              // energy from previous timestep
  
  vol_inv.init(mesh, vol_invx, *Lres_rho, 4);  // 1 / volume of element
  res_rho.init(mesh, res_rhox, *Lres_rho, 4);  // residual for density
  res_m.init(mesh, res_mx, *Lres_m, 5);        // residual for momentum
  res_e.init(mesh, res_ex, *Lres_e, 5);        // residual for energy
  
  nu_rho.init(mesh, nu_rhox, *a, 4);    // residual for density
  nu_m.init(mesh, nu_mx, *a, 5);        // residual for momentum
  nu_e.init(mesh, nu_ex, *a, 6);        // residual for energy
  
  mach_num.init(mesh, mach_numx, *a, 4);      // Mach number 
  delta.init(mesh, dvector, *a, 3);     // Stabilization term

  ComputeRME(mesh, w, rho, m, e, nsdim);

  // Initialize the vectors:


  x0.zero();
  x.zero();
  xold.zero();     
  res_rhox.zero(); 
  res_mx.zero();   
  res_ex.zero();   
  nu_rhox.zero();  
  nu_mx.zero();    
  nu_ex.zero();    
  mach_numx.zero();

  //  Initialize vectors for the residuals of 
  //  the equations in fix point iteration
  Vector residual(nsd*mesh.numVertices());

  // Set initial data
  cout << "Set initial data :"  << endl;
  SetInitialData(w, initial);
  x0 = x;  

  // Initialize output files 
  File file_u(get("velocity file name"));      // file for saving velocity 
  File file_p(get("pressure file name"));      // file for saving pressure
  File file_rho(get("density file name"));     // file for saving density
  File file_m(get("momentum file name"));      // file for saving momentum
  File file_e(get("energy file name"));        // file for saving energy
  File file_res(get("res_m file name"));       // file for saving residual
  File file_nu(get("nu_den file name"));       // file for saving shock-capturing term form velosity
  File file_mach_num(get("mach_num file name"));  // file for saving Mach number
  
  // Compute velocity and pressure
  ComputeUP(mesh, w, P, u, gamma, nsdim);

  // Compute cell mean velocity
  ComputeMean(mesh,um,u);

  // Pick up the components from w0 and w
  ComputeRME(mesh, w0, rho0, m0, e0, nsdim);
  ComputeRME(mesh, w, rho, m, e, nsdim);
  
  // Compute the volume inverse of an element
  ComputeVolInv(mesh, vol_invx);

  // Compute Mach number
  ComputeMachNum(mesh, rho, P, u, mach_num, gamma, nsdim);

  // Compute stabilization parameters
  ComputeStabilization(mesh,u,k,dvector);

  message("Assembling matrix: compressible");
  Assembler assembler(mesh);
  // Assembling matrices 
  assembler.assemble(res_rhox, *Lres_rho);
  assembler.assemble(res_mx, *Lres_m);
  assembler.assemble(res_ex, *Lres_e);

  ComputeNu(mesh, res_rho, res_m, res_e, nu_rho, nu_m, nu_e, nsdim);
  
  // Save the initial condition to file
  file_p << P;
  file_u << u;
  file_rho << rho;
  file_mach_num << mach_num;
  
  //file_m << m;
  //file_e << e;
  //file_res << res_m;
  file_nu << nu_rho;
  
  assembler.assemble(A, *a);

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

  // Residual, tolerance and maxmimum number of fixed point iterations
  real scalar_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 solution to solution at previous time step 
    x0 = x;

    // Compute velocity and pressure
    ComputeUP(mesh, w0, P, u, gamma, nsdim);

    // Compute cell mean velocity
    ComputeMean(mesh,um,u);

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

    // Fix-point iteration for non-linear problem 
    while (scalar_residual > rtol && iteration < max_iteration){

      // Compute velocity and pressure
      ComputeUP(mesh, w, P, u, gamma, nsdim);

      // Compute cell mean velocity
      ComputeMean(mesh,um,u);

      // Compute stabilization parameters
      ComputeStabilization(mesh,u,k,dvector);

      // Assemble compressible vector 
      message("Assemble vector: compressible");
      tic();
      assembler.assemble(A, *a);
      assembler.assemble(b, *L);
      message("Assemble took %g seconds",toc());

      // Set boundary conditions 
      for (uint i = 0; i < bcs.size(); i++)
	bcs[i]->apply(A, b, *a);

      // Store solution of previous iteration
      xold = x;
            
      // Solve the linear system  
      message("Solve linear system: compressible");
      tic();
      solver.solve(A, x, b);
      message("Linear solve took %g seconds",toc());
      
      // Compute shock capturing parameters
      if(iteration < 3)
	{
	  ComputeRME(mesh, w0, rho0, m0, e0, nsdim);
	  ComputeRME(mesh, w, rho, m, e, nsdim);
	  assembler.assemble(res_rhox, *Lres_rho);
	  assembler.assemble(res_mx, *Lres_m);
	  assembler.assemble(res_ex, *Lres_e);
	  ComputeNu(mesh, res_rho, res_m, res_e, nu_rho, nu_m, nu_e, nsdim);
	}

      // Compute residual  
      residual = xold;
      residual.vec() -= x.vec();
    
      message("Residual  : l2 norm = %e",residual.norm());
      scalar_residual = residual.norm();
      iteration++;
      
      message("iteration: %d", iteration);
    }
    
    // Compute shock capturing parameters
    ///*
    ComputeRME(mesh, w0, rho0, m0, e0, nsdim);
    ComputeRME(mesh, w, rho, m, e, nsdim);
    assembler.assemble(res_rhox, *Lres_rho);
    assembler.assemble(res_mx, *Lres_m);
    assembler.assemble(res_ex, *Lres_e);
    ComputeNu(mesh, res_rho, res_m, res_e, nu_rho, nu_m, nu_e, nsdim);

    // Compute Mach number
    ComputeMachNum(mesh, rho, P, u, mach_num, gamma, nsdim);
    
    if(scalar_residual > rtol)
      warning("CNS 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_rho << rho;
      file_mach_num << mach_num;
      //file_m << m;
      //file_e << e;
      //file_res << res_m;
      //file_nu << nu_rho;
      
      sample++;
    }

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

    // Update progress
    prog = t / T;
  }

  ComputeRME(mesh, w0, rho0, m0, e0, nsdim);
  ComputeRME(mesh, w, rho, m, e, nsdim);
  message("save solution to file");
  file_p << P;
  file_u << u;
  file_rho << rho;
  file_mach_num << mach_num;
  //file_m << m;
  //file_e << e;
  //file_res << res_rho;

  delete a;
  delete L;
  delete Lres_rho;
  delete Lres_m;
  delete Lres_e;
}
//-----------------------------------------------------------------------------
void CNSSolver::solve(Mesh& mesh, Function& f, Function& initial, Array<BoundaryCondition*>& bcs)
{
  CNSSolver solver(mesh, f, initial, bcs);
  solver.solve();
}
//-----------------------------------------------------------------------------
void CNSSolver::SetInitialData(Function& w, Function& initial)
{
  real* indata;
  uint N = w.vector().vec().size();

  indata = new real[N];
  initial.interpolate(indata);
  real* warr = w.vector().vec().array();
  w.vector().vec().set(indata);

  delete indata;
}

//-----------------------------------------------------------------------------
void CNSSolver::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 CNSSolver::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 CNSSolver::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;
        }
    }

  v.vector().vec().restore(varr);
  vmean.vector().vec().restore(vmeanarr);
}
//-----------------------------------------------------------------------------
void CNSSolver::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(); 
  
  vol_inv.vec().restore(varr);
}


//-----------------------------------------------------------------------------
void CNSSolver::ComputeStabilization(Mesh& mesh, Function& u, real k, 
				     Vector& dvector)
{
  real C1 = 2.0;   

  int N = mesh.numVertices();

  real* dvecarr = dvector.vec().array();
  real* uarr = u.vector().vec().array();

  dvector.init(mesh.numCells());	

  real normu; 

  for (CellIterator cell(mesh); !cell.end(); ++cell)
  {
    normu = 0.0;
    for (VertexIterator n(*cell); !n.end(); ++n)
      for (uint i = 0; i < mesh.geometry().dim(); i++)
	normu += sqr(uarr[i * N + (*n).index()]);
    
    normu += sqrt(normu);
    normu /= (*cell).numEntities(0);
    
    real h = (*cell).diameter();
    uint id = (*cell).index();
    
    dvecarr[id] =  C1 / sqrt( 1.0/sqr(k) + sqr(normu/h ));
  }
      dvector.vec().restore(dvecarr);
}

//-----------------------------------------------------------------------------
void CNSSolver::ComputeNu(Mesh& mesh, Function& res_rho, Function& res_m, Function& res_e, 
			  Function& nu_rho, Function& nu_m, Function& nu_e, int nsdim)
{

  int N = mesh.numVertices();
  int M = mesh.numCells();

  real normres = 0.0; 
  real h; 
 
  real C = 1.0;
  
  real* res_rhoarr = res_rho.vector().vec().array();
  real* res_marr   = res_m.vector().vec().array();
  real* res_earr   = res_e.vector().vec().array();

  real* nu_rhoarr = nu_rho.vector().vec().array();
  real* nu_marr   = nu_m.vector().vec().array();
  real* nu_earr   = nu_e.vector().vec().array();

  // nu_rho
  for (CellIterator cell(mesh); !cell.end(); ++cell)
    {
      uint id = (*cell).index();

      normres = sqr(res_rhoarr[id]);		    
      normres = sqrt(normres);
      h = (*cell).diameter();
      nu_rhoarr[id] = C * normres * h * h ;
    }
  normres = 0.0; 

  // nu_m
  for (CellIterator cell(mesh); !cell.end(); ++cell)
    {
      uint id = (*cell).index();
      
      if (nsdim == 3)
	normres = sqr(res_marr[0 * M + id]) + sqr(res_marr[1 * M + id]) + sqr(res_marr[2 * M + id]);
      else
	normres = sqr(res_marr[0 * M + id]) + sqr(res_marr[1 * M + id]);

      normres = sqrt(normres);
      h = (*cell).diameter();
      nu_marr[id] = C * normres * h * h ;
    }

  normres = 0.0; 
  // nu_e
  for (CellIterator cell(mesh); !cell.end(); ++cell)
    {
      uint id = (*cell).index();

      normres = sqr(res_earr[id]);
      normres = sqrt(normres);
      h = (*cell).diameter();
      nu_earr[id] = C * normres * h * h ;
    }
  
  res_rho.vector().vec().restore(res_rhoarr);
  res_m.vector().vec().restore(res_marr);
  res_e.vector().vec().restore(res_earr);

  nu_rho.vector().vec().restore(nu_rhoarr);
  nu_m.vector().vec().restore(nu_marr);
  nu_e.vector().vec().restore(nu_earr);

}

//-----------------------------------------------------------------------------
void CNSSolver::ComputeUP(Mesh& mesh, Function& w, Function& p, Function& u, real gamma, int nsdim)
{
  int N = mesh.numVertices();

  real rhoval, m0val, m1val, m2val, eval;

  Vector& wx = w.vector();

  real* warr = w.vector().vec().array();
  real* parr = p.vector().vec().array();
  real* uarr = u.vector().vec().array();

  for (VertexIterator n(mesh); !n.end(); ++n)
    {
      int id = (*n).index();

      Point pp = (*n).point();
      
      rhoval = warr[0 * N + id];
      m0val = warr[1 * N + id];
      m1val = warr[2 * N + id];
      if (nsdim == 3)
	{
	  m2val = warr[3 * N + id];
	  eval = warr[4 * N + id];
	  uarr[0 * N + id] = m0val / rhoval;
	  uarr[1 * N + id] = m1val / rhoval;
	  uarr[2 * N + id] = m2val / rhoval;
	  
	  parr[id] = (gamma - 1.0) * (eval - rhoval * 0.5 * (sqr(uarr[0 * N + id]) +
							     sqr(uarr[1 * N + id]) +
							     sqr(uarr[2 * N + id])));
	}
      else
	{
	  eval = warr[3 * N + id];
	  uarr[0 * N + id] =  m0val / rhoval;
	  uarr[1 * N + id] = m1val / rhoval;
	  
	  parr[id] =  (gamma - 1.0) * (eval - rhoval * 0.5 * (sqr(uarr[0 * N + id]) +
							      sqr(uarr[1 * N + id])));
	}
    }

  w.vector().vec().restore(warr);
  u.vector().vec().restore(uarr);
  p.vector().vec().restore(parr);
}

//-----------------------------------------------------------------------------
void CNSSolver::ComputeRME(Mesh& mesh, Function& w, Function& rho, Function& m, Function& e, int nsdim)
{
  int N = mesh.numVertices();

  real rhoval, m0val, m1val, m2val, eval;

  Vector& wx = w.vector();

  real* warr   = w.vector().vec().array();
  real* rhoarr = rho.vector().vec().array();
  real* marr   = m.vector().vec().array();
  real* earr   = e.vector().vec().array();

  for (VertexIterator n(mesh); !n.end(); ++n)
    {
      int id = (*n).index();
      
      rhoval = warr[0 * N + id];
      m0val  = warr[1 * N + id];
      m1val  = warr[2 * N + id];

      if (nsdim == 3)
	{
	  m2val = warr[3 * N + id];
	  eval  = warr[4 * N + id];
	  marr[2 * N + id] = m2val;
	}
      else
	{
	  eval = warr[3 * N + id];
	}

      rhoarr[id]       =  rhoval;
      marr[0 * N + id] = m0val;
      marr[1 * N + id] = m1val;
      earr[id]         = eval;
    }

  w.vector().vec().restore(warr); 
  rho.vector().vec().restore(rhoarr);
  m.vector().vec().restore(marr);
  e.vector().vec().restore(earr);

}

//-----------------------------------------------------------------------------
void CNSSolver::ComputeMachNum(Mesh& mesh, Function& rho, Function& p, Function& u, Function& machn, real gamma, int nsdim)
{
  int N = mesh.numVertices();

  real normu; 
  real* rhoarr   = rho.vector().vec().array();
  real* parr     = p.vector().vec().array();
  real* uarr     = u.vector().vec().array();
  real* machnarr = machn.vector().vec().array();

  for (VertexIterator n(mesh); !n.end(); ++n)
    {
      normu = 0.0;
      int id = (*n).index();
      real csum = 0.0;
      for (uint i = 0; i < nsdim; i++)
	{
	  csum += sqr(uarr[i * N + id]);
	}
      normu += sqrt(csum);
      real p = parr[id];
      if (p < 0.0)
	{
	  cout <<"pressure =  "<< p << "  at x = " << (*n).point().x() << "  and y = " << (*n).point().y() << endl; 
	}
      machnarr[id] = normu / sqrt(gamma * p * rhoarr[id]);
    }

  rho.vector().vec().restore(rhoarr);
  p.vector().vec().restore(parr);
  u.vector().vec().restore(uarr);
  machn.vector().vec().restore(machnarr);
}

