// Copyright (C) 2006 Johan Jansson.
// Licensed under the GNU GPL Version 2.
//
// Modified by Garth N. Wells 2006.
//
// First added:  2005
// Last changed: 2006-08-21

#ifndef __TIME_DEPENDENT_PDE_H
#define __TIME_DEPENDENT_PDE_H

#include <dolfin.h>
#include <dolfin/TimeStepper.h>

namespace dolfin { namespace unicorn
{

//   class dolfin::Form;
//   class dolfin::Mesh;
//   class dolfin::BoundaryCondition;
//   class dolfin::Function;
//   class dolfin::TimeStepper;
  class TimeDependentODE;

  /// This class implements the solution functionality for time dependent PDEs.

  class TimeDependentPDE
  {
  public:

    /// Define a time dependent PDE with boundary conditions
    TimeDependentPDE(Form& a, Form& L, Mesh& mesh,
		     BoundaryCondition& bc, real T,
		     Form* aJac = 0,
		     Form* aP = 0,
		     Form* LP = 0);

    /// Destructor
    virtual ~TimeDependentPDE();

    /// Solve PDE (in general a mixed system)
    virtual uint solve(Function& u);

    /// Compute initial value
    virtual void u0(uBlasVector& u);

    /// Update ODE
    virtual bool update(const uBlasVector& u, real t, bool end);

    /// ODE timestep
    virtual real timestep(real t, real k0) const;

    /// Compute right hand side dotu = f(u)
    virtual void fu(const uBlasVector& x, uBlasVector& dotx, real t);

    virtual void init(Function& U);

    virtual void save(Function& U, real t);

    virtual void preparestep();
    virtual void prepareiteration();

    virtual void computeJ();
    virtual void computeM();

    /// Return the bilinear form a(.,.)
    Form& a();

    /// Return the linear form L(.,.)
    Form& L();

    /// Return the mesh
    Mesh& mesh();

    /// Return the boundary condition
    BoundaryCondition& bc();

  protected:

    Form* _a;
    Form* _Lf;
    Mesh* _mesh;
    BoundaryCondition* _bc;

    TimeDependentODE* ode;
    TimeStepper* ts;

    int N;
    real t;
    real T;

    real k;

    File solutionfile;

  public:
    Vector* x;
    uBlasVector* dotx;

    bool implicit;
    bool constantM;
    bool constantJ;

    int fevals;
    real sampleperiod;
    real lastsample;

    Assembler assembler;
    Form* aJac;
    Form* aP;
    Form* LP;

    Vector dummy;

    Matrix J;
    Matrix M;
    Vector m;
  };

  class TimeDependentODE : public ODE
  {
  public:
    TimeDependentODE(TimeDependentPDE& pde, int N, real T);
    void u0(uBlasVector& u);

    // Evaluate right-hand side (mono-adaptive version)
    using ODE::f;
    virtual void f(const uBlasVector& u, real t, uBlasVector& y);
    virtual bool update(const uBlasVector& u, real t, bool end);
    virtual real timestep(real t, real k0) const;
    virtual void J(const uBlasVector& x, uBlasVector& y,
		   const uBlasVector& u, real t);
    virtual void M(const uBlasVector& x, uBlasVector& y,
		   const uBlasVector& u, real t);
    virtual Matrix& Jmatrix(const uBlasVector& u, real t);
    virtual Matrix& Mmatrix(real t);
    
  protected:
    TimeDependentPDE* pde;
  };
}}
#endif
