// elastic_genTerm_dfloat - A solver for non linear elastic problems using FEM
// Copyright (C) 2010-2026 Eric Bechet, Frederic Duboeuf
//
// See the LICENSE file for license information and contributions.
// Please report all bugs and problems to <bechet@cadxfem.org>.

#include <cstring>
#include <cstdio>
#include <fstream>

#include "elasticSolverDfloat.h"
#include "genAlgorithms.h"
#include "genSField.h"
#include "genDeriveOp.h"
#include "linearSystemCSR.h"
#include "linearSystemPETSc.h"
#include "linearSystemId.h"


#if defined(HAVE_POST)
#include "PView.h"
#include "PViewData.h"
#endif


ElasticSolverDfloat::~ElasticSolverDfloat()
{
  if (pAssembler) delete pAssembler;
  for (int i=0; i<MDomains.size(); ++i) delete MDomains[i];
}

void ElasticSolverDfloat::CheckProperties()
{
  for (int i=0; i<Data->Domains.size(); ++i)
  {
    if (Data->Domains[i]->type=="Elastic")
    {
      ElasticDomainDfloat* field = new ElasticDomainDfloat(*(Data->Domains[i]));
      MDomains.push_back (field);
    }
    else if (Data->Domains[i]->type=="IsotropicElastic")
    {
      IsotropicElasticDomainDfloat* field = new IsotropicElasticDomainDfloat(*(Data->Domains[i]));
      MDomains.push_back (field);
    }
    else if (Data->Domains[i]->type=="IsotropicPerfectPlastic")
    {
      IsotropicPerfectPlasticDomainDfloat* field = new IsotropicPerfectPlasticDomainDfloat(*(Data->Domains[i]));
      MDomains.push_back (field);
    }
  }
}


void ElasticSolverDfloat::readInputFile ( const std::string &fileName )
{
  modelname = fileName;
  std::cout<<"model name : " << modelname << std::endl ;
  Data->readInputFile(fileName);
  CheckProperties();
  printf("--> %d Domains\n", (int)MDomains.size());
  printf("--> %d BCs \n", (int)Data->BCs.size());
}




void ElasticSolverDfloat::BuildFunctionSpaces()
{
  if ( Data->dim==3 )
  {
    FSpaceDisp = genFSpace<genTensor1<double> >::Handle(new VectorLagrangeFSpace<double>());
  }

  if ( Data->dim==2 )
  {
    FSpaceDisp = genFSpace<genTensor1<double> >::Handle(new VectorLagrangeFSpace<double>(VectorLagrangeFSpace<double>::VECTOR_X, VectorLagrangeFSpace<double>::VECTOR_Y));
  }

  std::cout << "Displacement Field : iField=" << FSpaceDisp->getIncidentSpaceTag() << std::endl;

  for ( unsigned int i = 0; i < Data->BCs.size(); i++ )
  {
    if ((Data->BCs[i]->kind=="Dirichlet")||(Data->BCs[i]->kind==""))
    {
      if (Data->BCs[i]->What=="Displacement")
      {
        std::cout <<  "Dirichlet BC Disp" << std::endl;
        genFilterDofComponent filter ( Data->BCs[i]->comp1 );
        FixNodalDofs ( FSpaceDisp, Data->BCs[i]->begin(), Data->BCs[i]->end(), *pAssembler, *(Data->BCs[i]->fscalar), filter );
      }
    }
  }

  // we number the dofs : when a dof is numbered, it cannot be numbered
  // again with another number.
  for ( unsigned int i = 0; i < MDomains.size(); ++i )
  {
    MechanicsDomain* E = MDomains[i];
    std::cout << "Domain " << i << " : dim=" << E->dim << ", tag=" << E->tag << ", " << E->group()->size() << " elements" << std::endl;
    NumberDofs ( FSpaceDisp, E->begin(), E->end(), *pAssembler );
  }
}




double ElasticSolverDfloat::AssembleRHS(int loadstep)
{
  GaussQuadrature Integ_Boundary ( GaussQuadrature::Val );
  double volume=0;
  genTerm<genTensor0<double>,0>::Handle Volume(new ConstantField<genTensor0<double> >(1.0));
  genSField<genTensor1<double> >::Handle Sol(new genSField<genTensor1<double> >(pAssembler,FSpaceDisp));
  // build the external force vector
  double ext_energy=0;
  double int_energy=0;

  genTerm<genTensor2<double>,1>::Handle G = Gradient(FSpaceDisp);
  genTerm<genTensor2<double>,1>::Handle B =  (G+Transpose(G))*0.5;
  genTerm<genTensor2<double>,1>::Handle CB ( new genCache<genTensor2<double>,1>(B) ); // Strain tensor (test functions)
  genTerm<genTensor2<double>,0>::Handle GSol = Gradient(Sol);
  genTerm<genTensor2<double>,0>::Handle BSol =  (GSol+Transpose(GSol))*0.5;
  genTerm<genTensor2<double>,0>::ConstHandle CBSol ( new genCache<genTensor2<double>,0>(BSol) ); // Strain tensor (previous solution)

  
  for ( unsigned int i = 0; i < Data->BCs.size(); i++ )
  {
    if ((Data->BCs[i]->kind=="Neumann")||(Data->BCs[i]->kind==""))
    {
      if (Data->BCs[i]->What=="Force")
      {
        volume=0;
        std::cout << "Neumann BC Force" << std::endl;
        genTerm<genTensor1<double>,0>::Handle Func(new FunctionField<genTensor1<double> >(*Data->BCs[i]->fvector));
        genTerm<genTensor0<double>,1>::Handle loadterm = 
        Compose<FullContrProdOp>(FSpaceDisp,Func*(double(loadstep)/totalsteps));//        FullContractedProductCompose(FSpaceDisp,Func);
        genTerm<genTensor0<double>,0>::Handle energy1 = Compose<FullContrProdOp>(Sol,Func*(double(loadstep)/totalsteps));
        Assemble ( loadterm, Data->BCs[i]->begin(), Data->BCs[i]->end(), Integ_Boundary, *pAssembler );
        Assemble ( Volume, Data->BCs[i]->begin(), Data->BCs[i]->end(), Integ_Boundary, volume);
        Assemble ( energy1, Data->BCs[i]->begin(), Data->BCs[i]->end(), Integ_Boundary, ext_energy);
        printf("Support(%dD,%d) = %f [L^%d]\n", Data->BCs[i]->dim, Data->BCs[i]->tag, volume, Data->BCs[i]->dim);
      }
    }
  }
  printf("Energy Fext = %e \n", ext_energy);
  for (int i=0;i<MDomains.size();++i)
  {
    MechanicsDomain* E = MDomains[0];
#if 1
    genTerm<genTensor2<double>,0>::ConstHandle S = E->Stress(CBSol); // Stress tensor (based on previous solution)
    genTerm<genTensor0<double>,1>::Handle R=Compose<FullContrProdOp>(CB,S*-1.0);
    genTerm<genTensor0<double>,0>::Handle energy2=Compose<FullContrProdOp>(CBSol,S*-1.0);
#else 
    genTerm<genTensor4<double>,0>::ConstHandle S = E->Stiffness<double,0>(CBSol); // Stiffness tensor (based on previous solution)
    genTerm<genTensor0<double>,1>::Handle R=Compose<FullContrProdOp>(Compose<FullContrProdOp>(CB,S),CBSol*-1);
    genTerm<genTensor0<double>,0>::Handle energy2=Compose<FullContrProdOp>(Compose<FullContrProdOp>(CBSol,S),CBSol*-1);
#endif
    GaussQuadrature Integ_Bulk ( GaussQuadrature::GradGrad );
    Assemble(R, E->begin(), E->end(), Integ_Bulk, *pAssembler );
    Assemble(energy2, E->begin(), E->end(), Integ_Bulk, int_energy);
  }
  printf("Energy Fint = %e \n", int_energy);
  printf("Residual = %e \n", ext_energy+int_energy);
  double fac=(fabs(ext_energy)+fabs(int_energy));
  if (fac==0.) return 0.;
  return (fabs(ext_energy+int_energy)/fac);
}


void ElasticSolverDfloat::AssembleLHS(int loadstep)
{
  double volume=0;
  genTerm<genTensor0<double>,0>::Handle Volume(new ConstantField<genTensor0<double> >(1.0));
  genSField<genTensor1<double> >::Handle Sol(new genSField<genTensor1<double> >(pAssembler,FSpaceDisp));
  for (int i=0;i<MDomains.size();++i)
  {
    MechanicsDomain* E = MDomains[0];
    genTerm<genTensor2<double>,1>::Handle G = Gradient(FSpaceDisp);
    genTerm<genTensor2<double>,1>::Handle B =  (G+Transpose(G))*0.5;
    genTerm<genTensor2<double>,1>::ConstHandle CB ( new genCache<genTensor2<double>,1>(B) ); // Strain tensor (test functions)
    genTerm<genTensor2<dfloat<double> >,1>::Handle dCB=Build<genTensor2<dfloat<double> > >(CB);

    genTerm<genTensor2<double>,0>::Handle GSol = Gradient(Sol);
    genTerm<genTensor2<double>,0>::Handle BSol =  (GSol+Transpose(GSol))*0.5;
    genTerm<genTensor2<double>,0>::Handle CBSol ( new genCache<genTensor2<double>,0>(BSol) ); // Strain tensor (previous solution)
    genTerm<genTensor2<dfloat<double> >,0>::Handle dCBSol=Build<genTensor2<dfloat<double> > >(CBSol);
    genTerm<genTensor2<double>,1>::Handle CBz(CB*0.0);

    genTerm<genTensor2<dfloat<double> >,1>::Handle CB1=Build<genTensor2<dfloat<double> > >(CBz,CB);
    genTerm<genTensor2<dfloat<double> >,1>::Handle CB2=Build<genTensor2<dfloat<double> > >(CB);
    genTerm<genTensor2<dfloat<double> >,2>::ConstHandle CBB=Compose<PlusOp>(CB2,CB1);
    genTerm<genTensor2<double>,1>::ConstHandle S = E->Stress(CB); // Stress term (based on previous solution)
    genTerm<genTensor0<double>,2>::Handle K=Compose<FullContrProdOp>(S,CB); // actual stiffness matrix
    genTerm<genTensor2<dfloat<double> >,2>::ConstHandle Sd = E->Stress(CBB); // stress tensor
    genTerm<genTensor0<dfloat<double> >,2>::Handle RR=Compose<FullContrProdOp>(Sd,dCBSol);

    GaussQuadrature Integ_Bulk ( GaussQuadrature::GradGrad );

    genTerm<genTensor0<double>,2>::Handle T=K+Derive(RR); // newton raphson
//  genTerm<genTensor0<double>,2>::Handle T=K;            // modified newton raphson
    Assemble(T, E->begin(), E->end(), Integ_Bulk, *pAssembler);

    volume=0;
    Assemble(Volume, E->begin(), E->end(), Integ_Bulk, volume);
    printf("Support(%dD,%d) = %f [L^%d]\n", E->dim, E->tag, volume, E->dim);
  }
}

double ElasticSolverDfloat::BuildLinearSystem(int loadstep)
{
  double res=AssembleRHS(loadstep);
  AssembleLHS(loadstep);
  printf ( "nDofs=%d\n",pAssembler->sizeOfR() );
  printf ( "-- done assembling!\n" );

  if ( Data->solvertype == 1 )
  {
    exportKb();
    getchar();
  }
  return res;
}

void ElasticSolverDfloat::solve()
{
  linearSystem<double>* lsys=NULL;
  if ( Data->solvertype == 2 )
  {
#if defined(HAVE_TAUCS)
    lsys = new linearSystemCSRTaucs<double>;
#else
    printf ( "Taucs is not installed : Gmm is chosen to solve\n" );
    linearSystemCSRGmm<double>* buf= new linearSystemCSRGmm<double>;
    buf->setNoisy ( 2 );
    lsys=buf;
    Data->solvertype = 1;
#endif
  }
  else if ( Data->solvertype == 3 )
  {
#if defined(HAVE_PETSC)
    lsys = new linearSystemPETSc<double>;
#else
    printf ( "Petsc is not installed : Gmm is chosen to solve\n" );
    linearSystemCSRGmm<double>* buf= new linearSystemCSRGmm<double>;
    buf->setNoisy ( 2 );
    lsys=buf;
    Data->solvertype = 1;
#endif
  }
  else if ( Data->solvertype == 1 )
  {
    linearSystemCSRGmm<double>* buf= new linearSystemCSRGmm<double>;
    buf->setNoisy ( 2 );
    lsys=buf;
    Data->solvertype = 1;
  }


  if ( pAssembler ) delete pAssembler;
  pAssembler = new dofManager<double> ( lsys );

  printf ( "-- start solving\n" );
  BuildFunctionSpaces();
  for (int j=1;j<=totalsteps;++j)
  {
    double res;
    for (int i=0;i<100;++i)
    {
      std::cout << "### Start step " << j << " iteration " << i << std::endl;
      if ((i!=0)||(j!=1)) pAssembler->systemClear();
      res=BuildLinearSystem(j);
      SavedUnknownTerm<genTensor1<double> > saveDof(FSpaceDisp);
      for ( unsigned int d = 0; d < MDomains.size(); ++d )
      {
        SaveUnknown(*pAssembler,MDomains[d]->begin(), MDomains[d]->end(),saveDof);
      }
      pAssembler->systemSolve();
      SavedUnknownAddToSolution(saveDof,*pAssembler);
      if ((res<1e-10) && ((i>0)||(j>1))) break;
      for ( unsigned int d = 0; d < MDomains.size(); ++d )
      {
        genSField<genTensor1<double> >::Handle Sol(new genSField<genTensor1<double> >(pAssembler,FSpaceDisp));
        genTerm<genTensor2<double>,0>::Handle Epsilon = (Gradient(Sol)+Transpose(Gradient(Sol)))*0.5;
        MDomains[d]->UpdateHistory(Epsilon);
      }
    }
    if (res>1e-7)
    {
      std::cout << "not converged step " << j << std::endl;
      break;
    }
    for ( unsigned int d = 0; d < MDomains.size(); ++d )
    {
      MDomains[d]->CopyHistory();
    }
    std::stringstream vname1; vname1 << "displacement";
    std::stringstream vname2; vname2 << "Stress";
    std::stringstream vname3; vname3 << "VM";
    std::stringstream vname4; vname4 << "Strain";
    pvDisplacement = buildDisplacementView(vname1.str(),j);
    pvStress = buildStressView(vname2.str(),j);
    pvVonMises = buildVonMisesView(vname3.str(),j);
    pvStrain = buildStrainView(vname4.str(),j);
}
  printf ( "-- done solving! \n" );

  pvDisplacement->getData()->writeMSH("disp.msh", 3);
  pvStress->getData()->writeMSH("dstress.msh", 3);
  pvVonMises->getData()->writeMSH("dvm.msh", 3);
  pvStrain->getData()->writeMSH("dstrain.msh", 3);
  delete pvDisplacement;
  delete pvStress;
  delete pvVonMises;
  delete pvStrain;
}



void ElasticSolverDfloat::exportKb()
{
  FILE* f = fopen("K.txt", "w");
  double valeur;
  std::string sysname = "A";
  for (int i=0; i < pAssembler->sizeOfR(); ++i)
  {
    for (int j=0; j < pAssembler->sizeOfR(); ++j)
    {
      pAssembler->getLinearSystem(sysname)->getFromMatrix(i,j,valeur);
      fprintf(f, "%+e ", valeur);
    }
    fprintf(f,"\n");
  }
  fclose(f);

  f = fopen("b.txt", "w");
  for ( int i=0 ; i < pAssembler->sizeOfR(); ++i)
  {
    pAssembler->getLinearSystem(sysname)->getFromRightHandSide(i,valeur);
    fprintf(f, "%+e\n", valeur) ;
  }
  fclose(f);
}





PView* ElasticSolverDfloat::buildStrainView(const std::string &postFileName, double step)
{
  genSField<genTensor1<double> >::Handle Sol(new genSField<genTensor1<double> >(pAssembler,FSpaceDisp));
  genTerm<genTensor2<double>,0>::ConstHandle Epsilon = (Gradient(Sol)+Transpose(Gradient(Sol)))*0.5;
  return buildViewElementStep(Epsilon,MDomains,postFileName,pvStrain,step);
}

PView* ElasticSolverDfloat::buildStressView(const std::string &postFileName, double step)
{
  MechanicsDomain* E = MDomains[0];
  genSField<genTensor1<double> >::Handle Sol(new genSField<genTensor1<double> >(pAssembler,FSpaceDisp));
  genTerm<genTensor2<double>,0>::ConstHandle Epsilon = (Gradient(Sol)+Transpose(Gradient(Sol)))*0.5;
  genTerm<genTensor2<double>,0>::ConstHandle Stress = E->Stress(Epsilon);
  return buildViewElementStep(Stress,MDomains,postFileName,pvStress,step);
}

PView* ElasticSolverDfloat::buildVonMisesView(const std::string &postFileName, double step)
{
  MechanicsDomain* E = MDomains[0];
  genSField<genTensor1<double> >::Handle Sol(new genSField<genTensor1<double> >(pAssembler,FSpaceDisp));
  genTerm<genTensor2<double>,0>::ConstHandle Epsilon = (Gradient(Sol)+Transpose(Gradient(Sol)))*0.5;
  genTerm<genTensor2<double>,0>::ConstHandle Stress = E->Stress(Epsilon);
  genTerm<genTensor0<double>,0>::Handle  VM2= Compose<FullContrProdOp>(Stress,Stress)*1.5;
  genTerm<genTensor0<double>,0>::ConstHandle  VM= Apply<SqrtOp>(VM2);
  return buildViewElementStep(VM,MDomains,postFileName,pvVonMises,step);
}

PView* ElasticSolverDfloat::buildDisplacementView(const std::string &postFileName, double step)
{
  genTerm<genTensor1<double>,0>::ConstHandle Field(new genSField<genTensor1<double> >( pAssembler, FSpaceDisp ));
  return buildViewNodalStep(Field,MDomains,postFileName,pvDisplacement,step);
}