// 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 ELASTICITY_PDE_H
#define ELASTICITY_PDE_H

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

namespace dolfin { 
  namespace unicorn {

class ElasticityPDE : public TimeDependentPDE
{
public:
  ElasticityPDE(Form* a, Form* L,
		Mesh& mesh,
		BoundaryCondition& bc, real T,
		Form* aJac,
		Form* aP, Form* LP,
		Form* aS, Form* LS,
		Function& U,
		Function& B,
		Function& B0,
		real kk) :
    TimeDependentPDE(*a, *L, mesh, bc, T, aJac, aP, LP),
    U(U), B(B), B0(B0), aS(aS), LS(LS), kk(kk)
  {
    B.init(mesh, Bx, *aS, 0);
    cout << "Bx.size: " << Bx.size() << endl;
    B0.init(mesh, B0x, *aS, 0);

//     B.init(mesh, Bx, *L, 2);
//     B0.init(mesh, B0x, *L, 2);

    // X is compatible with U
    X.init(mesh, Xx, *a, 0);
    //X0x.init(X.vector().size());
    X0.init(mesh, X0x, *a, 0);
    
    int d = mesh.topology().dim();

    vxvalues = new real[d * mesh.numVertices()];
  }

  void deform(Mesh& mesh, Function& u)
  {
    MeshGeometry& geometry = mesh.geometry();
    
    int N = mesh.numVertices();

    u.interpolate(vxvalues);
    
    // Update the mesh
    for (VertexIterator n(mesh); !n.end(); ++n)
    {
      Vertex& vertex = *n;
      
      for(unsigned int i = 0; i < mesh.topology().dim(); i++)
      {
	geometry.x(vertex.index(), i) = vxvalues[i * N + vertex.index()];
// 	cout << "u(" << i << ", " << vertex.index() << ": " <<
// 	  vxvalues[i * N + vertex.index()] << endl;
      }
    }
  }

  void computeX()
  {
    X.vector().vec().copy(U.vector().vec(), 0, 0, X.vector().size());

    X.vector().vec() *= kk;
    X.vector().vec() += X0x.vec();
  }

  void computeB()
  {
    tic();

    assembler.assemble(BM, *aS);
    BM.mat().lump(Bm.vec());

    log(false);
    assembler.assemble(Btmp, *LS);
    log(true);

//     cout << "dotB: " << endl;
//     tmp.disp();

    Btmp.vec().div(Bm.vec());

    B.vector().vec() = Btmp.vec();

    B.vector().vec() *= kk;
    B.vector().vec() += B0.vector().vec();
    message("computeB took %g seconds",toc());
  }

  void shift()
  {
    X0.vector().vec().copy(X.vector().vec(), 0, 0, X0.vector().vec().size());
    B0.vector().vec().copy(B.vector().vec(), 0, 0, B0.vector().vec().size());
  }

  void preparestep()
  {
    shift();
  }

  void prepareiteration()
  {
    computeX();
    computeB();
    deform(mesh(), X);
  }

  void fu(const uBlasVector& x, uBlasVector& dotx, real T)
  {
//     cout << "U: " << endl;
//     U.vector().vec().disp();
//     cout << "B: " << endl;
//     B.vector().vec().disp();

    log(false);
    assembler.assemble(dotx, L());
    bc().apply(M, dotx, a());
    log(true);

//     cout << "dtU: " << endl;
//     dotx.disp();
  }

  void u0(uBlasVector& x)
  {
    // Compute initial dof value by L2 projection
    Function Pf = pproject(*aP, *LP, mesh());
  
    cout << "u0: " << endl;
    Pf.vector().disp();
  
    // Copy into components
    X.vector().vec().copy(Pf.vector().vec(), 0, 0, X.vector().size());
    U.vector().vec().copy(Pf.vector().vec(), 0, U.vector().size(),
			  U.vector().size());
    B.vector().vec().copy(Pf.vector().vec(), 0, 2*U.vector().size(),
			  B.vector().size());
    
    //FIXME: shift() seems to be needed for cG(1) but not dG(0), why?
    shift();
    
    cout << "X0: " << endl;
    X.vector().disp();
    cout << "x0: " << endl;
    x.disp();
    cout << "B0: " << endl;
    B.vector().disp();
  }

  Function& U;
  Function& B;
  Function& B0;
  Function X;
  Function X0;
  Vector B0x, Bx, Btmp, Utmp;
  Matrix BM;
  Vector Bm;

  Vector X0x, Xx;

  Form* aS;
  Form* LS;

  real* vxvalues;

  real kk;
};


  }}
#endif
