#include "../include/solve.hpp"

namespace ngsolve
{
  using namespace ngsolve;
  using namespace ngbla;
  using namespace ngla;




  class NonlinMechIntegrator : public BilinearFormIntegrator
  {
    enum { D = 3 };
    CoefficientFunction * coeff_e;
    CoefficientFunction * coeff_nu;
  public:
    NonlinMechIntegrator (CoefficientFunction * acoeff_e,
			  CoefficientFunction * acoeff_nu)
      : BilinearFormIntegrator(), coeff_e(acoeff_e),
	coeff_nu(acoeff_nu)
    { ; }
    virtual ~NonlinMechIntegrator ()
    { ; }

    static Integrator * Create (ARRAY<CoefficientFunction*> & coeffs)
    {
      return new NonlinMechIntegrator (coeffs[0], coeffs[1]);
    }

    virtual bool BoundaryForm () const { return 0; }
    virtual int DimFlux () const { return 6; }

    /**
       Assembles element matrix.
       Result is in elmat, memory is allocated by functions on LocalHeap.
    */
    virtual void
    AssembleElementMatrix (const FiniteElement & fel,
			   const ElementTransformation & eltrans,
			   FlatMatrix<double> & elmat,
			   LocalHeap & locheap) const
    {
      Vector<double> vec0 (D*fel.GetNDof());
      vec0 = 0;
      AssembleLinearizedElementMatrix (fel, eltrans, vec0, elmat, locheap);
    }

    virtual void
    AssembleLinearizedElementMatrix (const FiniteElement & bfel,
				     const ElementTransformation & eltrans,
				     FlatVector<double> & elveclin,
				     FlatMatrix<double> & elmat,
				     LocalHeap & lh) const
    {
      int fullynonlin=1;
      if(!fullynonlin)
      {
      const NodalFiniteElement & fel =
	dynamic_cast<const NodalFiniteElement&> (bfel);

      int ndof = fel.GetNDof();

      elmat.AssignMemory (ndof*D, ndof*D, lh);
      elmat = 0;

      FlatMatrixFixHeight<6, double> bmat (ndof * D, lh);
      FlatMatrixFixHeight<6, double> dbmat (ndof * D, lh);
      Mat<6,6> dmat;


      FlatMatrix<double> mlin(ndof, D,
			      const_cast<double*>
			      (static_cast<const double*>
			       (elveclin.Data())));

      int order = 2 * fel.Order();
      const IntegrationRule & ir =
	GetIntegrationRules().SelectIntegrationRule (fel.ElementType(), order);

      void * heapp = lh.GetPointer();
      for (int i = 0; i < ir.GetNIP(); i++)
	{
	  SpecificIntegrationPoint<D,D>
	    sip(ir[i], eltrans, lh);

	  int nd = fel.GetNDof();

	  Mat<D,D,double> gradref = Trans (mlin) * fel.GetDShape(sip.IP(), lh);
	  Mat<D> gradu = gradref * sip.GetJacobianInverse();
	  Mat<D> matf = gradu + Id<D>();

	  FlatMatrix<> grad (3, nd, lh);
	  Mat<3,3> jacobi = matf * sip.GetJacobian();

	  grad =  Trans (Inv (jacobi)) *
	    Trans (fel.GetDShape(sip.IP(),lh));

	  bmat = 0;
	  for (int i = 0; i < nd; i++)
	    {
	      bmat(0, 3*i  ) = grad(0, i);
	      bmat(1, 3*i+1) = grad(1, i);
	      bmat(2, 3*i+2) = grad(2, i);

	      bmat(3, 3*i  ) = grad(1, i);
	      bmat(3, 3*i+1) = grad(0, i);

	      bmat(4, 3*i  ) = grad(2, i);
	      bmat(4, 3*i+2) = grad(0, i);

	      bmat(5, 3*i+1) = grad(2, i);
	      bmat(5, 3*i+2) = grad(1, i);
	    }

	  dmat = 0;
	  double nu = Evaluate (*coeff_nu, sip);
	  double e = Evaluate (*coeff_e, sip);
	  for (int i = 0; i < 3; i++)
	    {
	      dmat(i,i) = 1-nu;
	      for (int j = 0; j < i; j++)
		dmat(i,j) = dmat(j,i) = nu;
	    }
	  for (int i = 3; i < 6; i++)
	    dmat(i,i) = 0.5 * (1-2*nu);

	  dmat *= (e / ((1 + nu) * (1 - 2 * nu)));

	  double fac =
	    fabs (Det (jacobi) ) * sip.IP().Weight();

	  dbmat = dmat * bmat;
	  elmat += fac * (Trans (dbmat) * bmat);

	  lh.CleanUp (heapp);
	}
     }
     else
     {
      const NodalFiniteElement & fel =
	dynamic_cast<const NodalFiniteElement&> (bfel);

      int ndof = fel.GetNDof();

      elmat.AssignMemory (ndof*D, ndof*D, lh);
      elmat = 0;

      FlatMatrixFixHeight<9, double> bmat (ndof * D, lh);
      FlatMatrixFixHeight<9, double> dbmat (ndof * D, lh);
      Mat<9,9> dmat;


      FlatMatrix<double> mlin(ndof, D,
			      const_cast<double*>
			      (static_cast<const double*>
			       (elveclin.Data())));

      int order = 2 * fel.Order();
      const IntegrationRule & ir =
	GetIntegrationRules().SelectIntegrationRule (fel.ElementType(), order);

      void * heapp = lh.GetPointer();
      for (int i = 0; i < ir.GetNIP(); i++)
	{
	  SpecificIntegrationPoint<D,D>
	    sip(ir[i], eltrans, lh);

	  int nd = fel.GetNDof();

	  //Mat<D,D,double> gradref = Trans (mlin) * fel.GetDShape(sip.IP(), lh);
	  //Mat<D> gradu = gradref * sip.GetJacobianInverse();
	  //Mat<D> matf = gradu + Id<D>();

	  FlatMatrix<> grad (3, nd, lh);
	  Mat<3,3> jacobi = sip.GetJacobian();

	  grad =  Trans (Inv (jacobi)) *
	    Trans (fel.GetDShape(sip.IP(),lh));

	  bmat = 0;
	  for (int k = 0; k < nd; k++)
	    {
	      for (int j=0; j < 3; j++)
	      {
 	        bmat(0+3*j, 3*k+j) = grad(0, k);
	        bmat(1+3*j, 3*k+j) = grad(1, k);
	        bmat(2+3*j, 3*k+j) = grad(2, k);
	      }
            }


	  // gradient in reference coordinates:
	  Mat<D,D,double> gradref = Trans (mlin) * fel.GetDShape(sip.IP(), lh);

	  // gradient in real coordinates:
	  Mat<D> grad1 = gradref * sip.GetJacobianInverse();

	  // deformation gradient:
	  Mat<D> matf = grad1 + Id<D>();

	  double eps=1e-8;

	  int i1, i2, ii;
	  for (i1=0; i1 < 3; i1++)
  	    for (i2=0; i2 < 3; i2++)
	     {
	       ii=i1*3+i2;

	       matf(i1,i2)+=eps;

	       // Green-Lagrange strain tensor
	       Mat<D> strain;
	       strain = 0.5 * (Trans ( matf) * matf - Id<D>());

	       Mat<D> piola2;
	       double e = Evaluate (*coeff_e, sip);
	       double nu = Evaluate (*coeff_nu, sip);
	       double mu = e / 2 / (1+nu);
	       double lam = e * nu / ((1+nu)*(1-2*nu));
	       piola2 = (2*mu) * strain + (lam * Trace(strain)) * Id<D>();

	       Mat<D> piola1a = matf * piola2;

	       matf(i1,i2)-=2*eps;

	       // Green-Lagrange strain tensor
	       strain = 0.5 * (Trans ( matf) * matf - Id<D>());

	       piola2 = (2*mu) * strain + (lam * Trace(strain)) * Id<D>();

	       Mat<D> dpiola1 = (1./2./eps)*(piola1a-(matf * piola2));

	       int j1, j2, jj;
	       for (j1=0; j1 < 3; j1++)
  	         for (j2=0; j2 < 3; j2++)
		 {
	           jj=j1*3+j2;
		   dmat(ii,jj) = dpiola1(j1,j2);
		 }

	       matf(i1,i2)+=eps;

	     }


	  //(*testout) << "dmat=" << dmat << endl;


	  double fac =
	    fabs (Det (jacobi) ) * sip.IP().Weight();

	  dbmat = dmat * bmat;
	  elmat += fac * (Trans (dbmat) * bmat);

	  lh.CleanUp (heapp);
	}
      }
    }

    virtual void
    ApplyElementMatrix (const FiniteElement & bfel,
			const ElementTransformation & eltrans,
			const FlatVector<double> & elx,
			FlatVector<double> & ely,
			LocalHeap & lh) const
    {
      int i;

      const NodalFiniteElement & fel =
	dynamic_cast<const NodalFiniteElement&> (bfel);
      int ndof = fel.GetNDof ();

      ely = 0;

      FlatMatrix<double> melx(ndof, D,
			      const_cast<double*>
			      (static_cast<const double*>
			       (elx.Data())));

      FlatMatrix<double> mely(ndof, D,
			      static_cast<double*>
			       (ely.Data()));

      int order = 2 * fel.Order();
      const IntegrationRule & ir =
	GetIntegrationRules().SelectIntegrationRule (fel.ElementType(), order);

      for (i = 0; i < ir.GetNIP(); i++)
	{
	  const IntegrationPoint & ip = ir.GetIP(i);

	  SpecificIntegrationPoint<D,D> sip (ir[i], eltrans, lh);

	  // gradient in reference coordinates:
	  Mat<D,D,double> gradref = Trans (melx) * fel.GetDShape(ip, lh);

	  // gradient in real coordinates:
	  Mat<D> grad = gradref * sip.GetJacobianInverse();

	  // deformation gradient:
	  Mat<D> matf = grad + Id<D>();

	  // Green-Lagrange strain tensor
	  Mat<D> strain = 0.5 * (Trans ( matf) * matf - Id<D>());

	  Mat<D> piola2;
	  double e = Evaluate (*coeff_e, sip);
	  double nu = Evaluate (*coeff_nu, sip);
	  double mu = e / 2 / (1+nu);
	  double lam = e * nu / ((1+nu)*(1-2*nu));
	  piola2 = (2*mu) * strain + (lam * Trace(strain)) * Id<D>();

	  Mat<D> piola1 = matf * piola2;

	  double fac = fabs (sip.GetJacobiDet()) * ip.Weight();

	  // test with nabla v:
	  Mat<D,D,double> trans = sip.GetJacobianInverse() * Trans (piola1);
	  mely += fac * (fel.GetDShape(sip.IP(),lh) * trans);
	}
    }

    virtual void
    ApplyLinearizedElementMatrix (const FiniteElement & fel,
				  const ElementTransformation & eltrans,
				  const FlatVector<double> & ellin,
				  const FlatVector<double> & elx,
				  FlatVector<double> & ely,
				  LocalHeap & locheap) const
    {
      ApplyElementMatrix (fel, eltrans, elx, ely, locheap);
    }

    virtual double Energy (const FiniteElement & bfel,
			   const ElementTransformation & eltrans,
			   const FlatVector<double> & elx,
			   LocalHeap & lh) const
    {
      int i;

      const NodalFiniteElement & fel =
	dynamic_cast<const NodalFiniteElement&> (bfel);
      int ndof = fel.GetNDof ();

      FlatMatrix<double> melx(ndof, D,
			      const_cast<double*>
			      (static_cast<const double*>
			       (elx.Data())));

      int order = 2 * fel.Order();
      const IntegrationRule & ir =
	GetIntegrationRules().SelectIntegrationRule (fel.ElementType(), order);

      double energy = 0;

      for (i = 0; i < ir.GetNIP(); i++)
	{
	  const IntegrationPoint & ip = ir.GetIP(i);

	  SpecificIntegrationPoint<D,D> sip (ir[i], eltrans, lh);

	  // gradient in reference coordinates:
	  Mat<D,D,double> gradref = Trans (melx) * fel.GetDShape(ip, lh);

	  // gradient in real coordinates:
	  Mat<D> grad = gradref * sip.GetJacobianInverse();

	  // deformation gradient:
	  Mat<D> matf = grad + Id<D>();

	  // Green-Lagrange strain tensor
	  Mat<D> strain = 0.5 * (Trans ( matf) * matf - Id<D>());

	  Mat<D> piola2;
	  double e = Evaluate (*coeff_e, sip);
	  double nu = Evaluate (*coeff_nu, sip);
	  double mu = e / 2 / (1+nu);
	  double lam = e * nu / ((1+nu)*(1-2*nu));

	  energy += fabs (sip.GetJacobiDet()) * ip.Weight() *
	    (mu*Trace (strain * strain) + lam/2*sqr (Trace(strain)));
	}
      return energy;
    }

    virtual void
    CalcFlux (const FiniteElement & bfel,
	      const ElementTransformation & eltrans,
	      const IntegrationPoint & ip,
	      const FlatVector<double> & elx,
	      FlatVector<double> & flux,
	      bool applyd,
	      LocalHeap & lh) const
    {
    /*
      cerr << "calcflux<double> called for class "
	   << typeid(*this).name()
	   << endl;
	   */

      const NodalFiniteElement & fel =
	dynamic_cast<const NodalFiniteElement&> (bfel);
      int ndof = fel.GetNDof ();

      FlatMatrix<double> melx(ndof, D,
			      const_cast<double*>
			      (static_cast<const double*>
			       (elx.Data())));

      SpecificIntegrationPoint<D,D> sip (ip, eltrans, lh);

      // gradient in reference coordinates:
	Mat<D,D,double> gradref = Trans (melx) * fel.GetDShape(ip, lh);

	// gradient in real coordinates:
	Mat<D> grad = gradref * sip.GetJacobianInverse();

	// deformation gradient:
	Mat<D> matf = grad + Id<D>();

	// Green-Lagrange strain tensor
	Mat<D> mflux = 0.5 * (Trans ( matf) * matf - Id<D>());

	if (applyd)
	{
	  double e = Evaluate (*coeff_e, sip);
	  double nu = Evaluate (*coeff_nu, sip);
	  double mu = e / 2 / (1+nu);
	  double lam = e * nu / ((1+nu)*(1-2*nu));
	  mflux = (2*mu) * mflux + (lam * Trace(mflux)) * Id<D>();
        }

	flux.AssignMemory (6, lh);

	flux(0) = mflux(0,0);
	flux(1) = mflux(1,1);
	flux(2) = mflux(2,2);
	flux(3) = mflux(1,2);
	flux(4) = mflux(0,2);
	flux(5) = mflux(0,1);

    }

  };




  class NumProcNonlinElast : public NumProc
  {
  protected:
    ///
    BilinearForm * bfa;
    BilinearForm * bfalin;
    ///
    LinearForm * lff;
    ///
    GridFunction * gfu;
    ///
    Preconditioner * pre;
    ///
    int maxsteps;
    ///
    double prec;

  public:
    ///
    NumProcNonlinElast (PDE & apde, const Flags & flags);
    ///
    virtual ~NumProcNonlinElast();

    static NumProc * Create (PDE & pde, const Flags & flags)
    {
      return new NumProcNonlinElast (pde, flags);
    }


    ///
    virtual void Do(LocalHeap & lh);
    ///
    virtual string GetClassName () const
    {
      return "NonlinElast";
    }

    virtual void PrintReport (ostream & ost)
    {
      ost << GetClassName() << endl
	  << "Bilinear-form = " << bfa->GetName() << endl
	  << "Linear-form   = " << lff->GetName() << endl
	  << "Gridfunction  = " << gfu->GetName() << endl
	  << "Preconditioner = " << ((pre) ? pre->ClassName() : "None") << endl
	  << "precision     = " << prec << endl
	  << "maxsteps      = " << maxsteps << endl;
    }

    ///
    static void PrintDoc (ostream & ost);
  };



  NumProcNonlinElast :: NumProcNonlinElast (PDE & apde, const Flags & flags)
    : NumProc (apde)
  {
    bfa = pde.GetBilinearForm (flags.GetStringFlag ("bilinearform", ""));
    lff = pde.GetLinearForm (flags.GetStringFlag ("linearform", ""));
    gfu = pde.GetGridFunction (flags.GetStringFlag ("gridfunction", ""));
    pre = pde.GetPreconditioner (flags.GetStringFlag ("preconditioner", ""), 1);
    maxsteps = int(flags.GetNumFlag ("maxsteps", 200));
    prec = flags.GetNumFlag ("prec", 1e-12);
  }

  NumProcNonlinElast :: ~NumProcNonlinElast()
  {
    ;
  }

  void NumProcNonlinElast :: PrintDoc (ostream & ost)
  {
    ost <<
      "\n\nNumproc NonlinElast:\n" \
      "------------\n" \
      "Solves the nonlinear system resulting from a boundary value problem\n\n" \
      "Required flags:\n" 
      "-bilinearform=<bfname>\n" 
      "    bilinear-form providing the matrix\n" \
      "-linearform=<lfname>\n" \
      "    linear-form providing the right hand side\n" \
      "-gridfunction=<gfname>\n" \
      "    grid-function to store the solution vector\n" 
      "\nOptional flags:\n"
      "-predoncitioner=<prename>\n"
      "-maxsteps=n\n"
      "-prec=eps\n"
      "-qmr\n"
      "-print\n"
      "    write matrices and vectors into logfile\n"
	<< endl;
  }


  void NumProcNonlinElast :: Do(LocalHeap & lh)
  {
    cout << "Solve nonlinear elasticity" << endl;

    const BaseVector & vecf = lff->GetVector();
    BaseVector & vecu = gfu->GetVector();
    // const BaseMatrix & premat = pre->GetMatrix();


    BaseVector & uold = *vecu.CreateVector();
    BaseVector & d = *vecu.CreateVector();
    BaseVector & w = *vecu.CreateVector();
    BilinearFormApplication applya(bfa);

    /*
    LinearizedBilinearFormApplication applyalin(&bfa, &uk);
    CGSolver<double> inva (applyalin, c);
    inva.SetMaxSteps (500);
    */

    // const BaseMatrix & inva = premat;

    // stationary solver:

    double err, errold, err0 = 1;
    double energy, energyold;

    //    LocalHeap lh (1000000);

    for (int i = 1; i <= maxsteps; i++)
      {
	// if (i <= 5)
	bfa -> AssembleLinearization (vecu, lh);
	const BaseMatrix & mata = bfa->GetMatrix();	

	BaseMatrix & inva = 
	  *dynamic_cast<const BaseSparseMatrix&>(mata).InverseMatrix();

	d = vecf - applya * vecu;
	err = L2Norm(d);
	if (i == 1) err0 = err;

	energy = bfa->Energy(vecu) - InnerProduct (vecf, vecu);

	cout << "newton it " << i;
	cout << " err = " << err/err0;
	cout << " energy = " << energy << endl;

	errold = err;
	energyold = energy;

	w = inva * d;
	uold = vecu;
	int lin_its = 0;
	double tau = 2;
	do
	  {
	    vecu = uold + w;
	    d = vecf - applya * vecu;
	    w *= 0.5;
	    tau *= 0.5;
	    err = L2Norm(d);

	    energy = bfa->Energy(vecu) - InnerProduct (vecf, vecu);

	    cout << "tau = " << tau;
	    cout << " energy = " << energy << endl;
	    cout << " linsearch, err = " << err/err0 << endl;
	  }
	// while (err > errold);
	while (energy > energyold && lin_its++ < 30 && err > 1e-7*err0);

	if (err < 1e-7*err0) break;
      }
  }






  namespace {
    class Init
    { 
    public: 
      Init ();
    };
    
    Init::Init()
    {
      GetNumProcs().AddNumProc ("nonlinelast", NumProcNonlinElast::Create, NumProcNonlinElast::PrintDoc);
      
      GetIntegrators().AddBFIntegrator ("nonlinmech", 3, 2,
					NonlinMechIntegrator::Create);
    }

    
    Init init;
  }

}