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

#include <unicorn/TimeDependentPDE.h>
#include <unicorn/Project.h>

//#include <dolfin.h>
#include <dolfin/Matrix.h>
#include <dolfin/Vector.h>
#include <dolfin/Assembler.h>
#include <dolfin/Form.h>
#include <dolfin/Mesh.h>
#include <dolfin/BoundaryCondition.h>
#include <dolfin/Function.h>
#include <dolfin/TimeStepper.h>
#include <dolfin/UFC.h>
#include <dolfin/SparsityPatternBuilder.h>
#include <dolfin/SparsityPattern.h>
#include <dolfin/timing.h>
#include <dolfin/parameters.h>


using namespace dolfin;
using namespace dolfin::unicorn;

TimeDependentPDE::TimeDependentPDE(Form& a, Form& L, Mesh& mesh, 
				   BoundaryCondition& bc, real T,
				   Form* aJac, Form* aP, Form* LP)
  : _a(&a), _Lf(&L),
    _mesh(&mesh), _bc(&bc),
    ode(0), ts(0),
    t(0), T(T),
    k(get("ODE initial time step")),
    solutionfile("solution.pvd"),
    x(0), dotx(0),
    fevals(0), 
    sampleperiod(T / 100.0), lastsample(0.0),
    assembler(mesh),
    aJac(aJac), aP(aP), LP(LP)
{
}
//-----------------------------------------------------------------------------
TimeDependentPDE::~TimeDependentPDE()
{
  if(x != 0)
    delete x;
  if(dotx != 0)
    delete dotx;
  if(ts != 0)
    delete ts;
  if(ode != 0)
    delete ode;
}
//-----------------------------------------------------------------------------
dolfin::uint TimeDependentPDE::solve(Function& U)
{
  dolfin_assert(_a);
  dolfin_assert(_Lf);
  dolfin_assert(_mesh);

  init(U);

  cout << "N: " << N << endl;
  cout << "x size: " << x->size() << endl;

  // Initialize ODE (requires x)
  ode = new TimeDependentODE(*this, N, T);
  ts = new TimeStepper(*ode);

  // Write a message
  message("Solving time dependent PDE.");

  save(U, t);

  // Start time-stepping
  while(t < T && !ts->stopped()) {
    
    preparestep();
    t = ts->step();

    save(U, t);
  }

  return 0;
}
//-----------------------------------------------------------------------------
void TimeDependentPDE::u0(uBlasVector& u)
{
  cout << "TimeDependentPDE::u0" << endl;

  // Compute initial dof value by L2 projection
//   Form* aP = new ProjectionBilinearForm();
//   Form* LP = new ProjectionLinearForm(u0val);
  
  Function Pf = pproject(*aP, *LP, mesh());
  
  cout << "u0: " << endl;
  Pf.vector().disp();
  
  //U.vector().vec().copy(Pf.vector().vec(), 0, 0, U.vector().size());
  
  x->vec().copy(Pf.vector().vec(), 0, 0, x->size());

  u.copy(x->vec(), 0, 0, u.size());
}
//-----------------------------------------------------------------------------
bool TimeDependentPDE::update(const uBlasVector& u, real t, bool end)
{
  return true;
}
//-----------------------------------------------------------------------------
real TimeDependentPDE::timestep(real t, real k0) const
{
  return k;
}
//-----------------------------------------------------------------------------
void TimeDependentPDE::fu(const uBlasVector& x, uBlasVector& dotx, real t)
{
  set("output destination", "silent");
  assembler.assemble(dotx, L());
  bc().apply(M, dotx, a());
  set("output destination", "terminal");

  fevals = fevals + 1;
}
//-----------------------------------------------------------------------------
Form& TimeDependentPDE::a()
{
  dolfin_assert(_a);
  return *_a;
}
//-----------------------------------------------------------------------------
Form& TimeDependentPDE::L()
{
  dolfin_assert(_Lf);
  return *_Lf;
}
//-----------------------------------------------------------------------------
Mesh& TimeDependentPDE::mesh()
{
  dolfin_assert(_mesh);
  return *_mesh;
}
//-----------------------------------------------------------------------------
BoundaryCondition& TimeDependentPDE::bc()
{
  dolfin_assert(_bc);
  return *_bc;
}
//-----------------------------------------------------------------------------
void TimeDependentPDE::init(Function& U)
{
  // Storage for dof maps                                                     
  DofMapSet dof_map_set;

  // Update dof maps
  dof_map_set.update(a().form(), mesh());

  // Create data structure for local assembly data
  UFC ufc(a().form(), mesh(), dof_map_set);

  SparsityPattern sparsity_pattern;
  SparsityPatternBuilder::build(sparsity_pattern, mesh(), ufc, dof_map_set);

  // Assume bilinear tensor is square
  N = sparsity_pattern.size(0);

  x = new Vector(N);
  dotx = new uBlasVector(N);
  dummy.init(N);

  U.init(mesh(), *x, a(), 0);

  // FIXME: How should we handle time-dependent user functions?
//   for(unsigned int i = 0; i < this->L().num_functions; i++)
//   {
//     Function* f = this->L().function(i);
//     f->sync(t);
//   }

//   bc().sync(t);

  //J.init(N, N);

  cout << "Initial computation of M and J" << endl;

  computeM();
  computeJ();

//   cout << "J:" << endl;
//   J.disp(10);

  implicit = get("ODE implicit");
  constantM = get("ODE M matrix constant");
  constantJ = false;

  cout << "constantM = " << constantM << endl;
}
//-----------------------------------------------------------------------------
void TimeDependentPDE::save(Function& U, real t)
{
  cout << "Saving" << endl;
  
  if(t == 0.0)
  {
    solutionfile << U;
  }
  
  while(lastsample + sampleperiod < t)
  {
    lastsample = std::min(t, lastsample + sampleperiod);
    solutionfile << U;
  }
}
void TimeDependentPDE::preparestep()
{
}
//-----------------------------------------------------------------------------
void TimeDependentPDE::prepareiteration()
{
}
//-----------------------------------------------------------------------------
void TimeDependentPDE::computeJ()
{
  cout << "Computing J" << endl;

  tic();
  set("output destination", "silent");
  assembler.assemble(J, *aJac);
  set("output destination", "terminal");
  message("Assemble J took %g seconds",toc());

  tic();
  // FIXME: The Jacobian should have a zero row for strong Dirichlet BC
  // The following solution works as long as k < 0.9 (1.0)
  set("output destination", "silent");
  bc().apply(J, dummy, *aJac);
  set("output destination", "terminal");
  message("BC J took %g seconds",toc());

//   cout << "J: " << endl;
//   J.disp();

  cout << "Computing J done" << endl;
}
//-----------------------------------------------------------------------------
void TimeDependentPDE::computeM()
{
  cout << "Computing M" << endl;
  //std::cout << "M.M2:" << M.mat().mat() << std::endl;
  tic();
  set("output destination", "silent");
  assembler.assemble(M, a());
  bc().apply(M, dummy, a());
  set("output destination", "terminal");
  message("Assemble M took %g seconds",toc());

//   cout << "M: " << endl;
//   M.disp();

  tic();
  cout << "Lumping M" << endl;
  M.mat().lump(m.vec());
  message("Lump M took %g seconds",toc());
  cout << "Computing M done" << endl;
}
//-----------------------------------------------------------------------------
TimeDependentODE::TimeDependentODE(TimeDependentPDE& pde, int N, real T) :
  ODE(N, T), pde(&pde)
{
}
//-----------------------------------------------------------------------------
void TimeDependentODE::u0(uBlasVector& u)
{
  cout << "TimeDependentODE::u0" << endl;
  pde->u0(u);
}
//-----------------------------------------------------------------------------
void TimeDependentODE::f(const uBlasVector& u, real t, uBlasVector& y)
{
  pde->x->vec().copy(u, 0, 0, u.size());

  pde->prepareiteration();

  //pde->fu(*(pde->x), *(pde->dotx), t);
  pde->fu(u, *(pde->dotx), t);

  //Vector tmp(*(pde->dotx));

  if(!pde->implicit)
  {
    uBlasVector m2(pde->m.size());
    m2.copy(pde->m.vec(), 0, 0, m2.size());

    pde->dotx->div(m2);
  }
  //pde->M.mult(tmp, *(pde->dotx));

  y.copy(*(pde->dotx), 0, 0, u.size());
}
//-----------------------------------------------------------------------------
bool TimeDependentODE::update(const uBlasVector& u, real t, bool end)
{
  return pde->update(u, t, end);
}
//-----------------------------------------------------------------------------
real TimeDependentODE::timestep(real t, real k0) const
{
  return pde->timestep(t, k0);
}
//-----------------------------------------------------------------------------
void TimeDependentODE::J(const uBlasVector& x, uBlasVector& y,
			 const uBlasVector& u, real t)
{
  cout << "TDODE::J" << endl;

  //const bool linearPDE = false;

  if(!pde->constantJ)
  {
    pde->computeJ();
    //FEM::assemble(*(pde->aJac), pde->J, pde->mesh());
    //FEM::applyBC(pde->J, pde->mesh(), pde->aJac->trial(), pde->bc());
  }

  //cout << "J mult" << endl;
  PETScVector x2(x.size());
  PETScVector y2(y.size());

  x2.copy(x, 0, 0, x2.size());

  pde->J.mat().mult(x2, y2);

  y.copy(y2, 0, 0, y.size());
}
//-----------------------------------------------------------------------------
void TimeDependentODE::M(const uBlasVector& x, uBlasVector& y,
			 const uBlasVector& u, real t)
{
  //cout << "TDODE::M" << endl;

  if(!pde->constantM)
  {
    pde->computeM();
    //FEM::assemble(*(pde->aJac), pde->J, pde->mesh());
    //FEM::applyBC(pde->J, pde->mesh(), pde->aJac->trial(), pde->bc());
  }

  PETScVector x2(x.size());
  PETScVector y2(y.size());

  x2.copy(x, 0, 0, x2.size());

  pde->M.mat().mult(x2, y2);

  y.copy(y2, 0, 0, y.size());
}
//-----------------------------------------------------------------------------
Matrix& TimeDependentODE::Jmatrix(const uBlasVector& u, real t)
{
  cout << "Jmatrix at t = " << t << endl;

  if(!pde->constantJ)
  {
    pde->computeJ();
  }

//   cout << "J =" << endl;
//   pde->J.disp();

  return pde->J;
}
//-----------------------------------------------------------------------------
Matrix& TimeDependentODE::Mmatrix(real t)
{
  // FIXME: Adaptively find out if M is constant
  // FIXME: Define parameter for constant M
  if(!pde->constantM)
  {
    pde->computeM();
  }

  return pde->M;
}
//-----------------------------------------------------------------------------

