#include <petscdm.h>
#include <petscdmlabel.h>
#include <petscds.h>
#include <petscdmplex.h>
#include <petscksp.h>
#include <petscsnes.h>
#include <petscts.h>
#include <math.h>
#include <petscsys.h>

/* define the pointwise functions we'd like to project */

void projectstress(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[],
                         const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[],
                         const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[],
                         const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[],
                         PetscReal t, const PetscScalar x[], PetscInt numConstants,
                         const PetscScalar constants[], PetscScalar v[]){
  const PetscReal mu =76.923076923, lbda=115.384615385;
  PetscInt Ncomp = dim;
  PetscInt comp,d;
  PetscReal sigma[dim*dim];

   for(comp=0;comp<Ncomp;comp++){
    for(d=0;d<dim;d++){
      sigma[comp*dim+d]=mu*(u_x[comp*dim+d]+u_x[d*dim+comp]);
    }
    for(d=0;d<dim;d++){
      sigma[comp*dim+comp]+=lbda*u_x[d*dim+d];
    }
  }

   for(d=0;d<dim;d++){
     v[d] = sigma[d*dim+d];
   }
   v[3] = sigma[0*dim+1];
   v[4] = sigma[0*dim+2];
   v[5] = sigma[1*dim+2];

    
}

void projectdisplacement(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[],
                         const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[],
                         const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[],
                         const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[],
                         PetscReal t, const PetscScalar x[], PetscInt numConstants,
                         const PetscScalar constants[], PetscScalar v[]){
  PetscInt d;
  
  for(d=0;d<dim;d++){
     v[d] =u[d];
   }
}

static PetscErrorCode initial_displacement_vector(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *ctx)
{
  u[0]=0.0;
  u[1]=0.1*pow(x[2],2);
  u[2]=0.1*x[2];
  return 0;
}

static PetscErrorCode expected_stress_vector(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *ctx)
{
  const PetscReal mu =76.923076923, lbda=115.384615385;
  PetscReal strain[dim*dim];
  PetscReal gradu[dim*dim];
  PetscReal stress[dim*dim];
  PetscInt i,j;

  /* gradient of the displacement field: */
  for(i=0;i<dim;i++){
    for(j=0;j<dim;j++){
      gradu[i*dim+j]=0.0;
    }
  }

  gradu[1*dim+2]=0.2*x[2];
  gradu[2*dim+2]=0.1;

  for(i=0;i<dim;i++){
    for(j=0;j<dim;j++){
      strain[i*dim+j]=0.5*(gradu[i*dim+j]+gradu[j*dim+i]);
    }
  }

  for(i=0;i<dim;i++){
    for(j=0;j<dim;j++){
      stress[i*dim+j]=2.0*mu*strain[i*dim+j];
    }
    for(j=0;j<dim;j++){
      stress[i*dim+i]+=lbda*strain[j*dim+j];
    }
  }
  
  for(i=0;i<dim;i++){
    u[i] = stress[i*dim+i];
  }
  u[3] = stress[0*dim+1];
  u[4] = stress[0*dim+2];
  u[5] = stress[1*dim+2];
  
  
  return 0;
}

static PetscErrorCode zero_stress(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *ctx)
{
  const PetscInt Ncomp = 2*dim;
  PetscInt comp;
  for (comp=0;comp<Ncomp;comp++) u[comp]=0.0;
  return 0;
}

static PetscErrorCode zero_vector(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *ctx)
{
  PetscInt comp;
  for (comp=0;comp<dim;comp++) u[comp]=0.0;
  return 0;
}

static PetscErrorCode SetupDiscretization(DM dm){

  PetscInt dim = 3;
  PetscFE fe_displacement, fe_stress;
  PetscDS prob;
  PetscErrorCode ierr;
  PetscBool simplex = PETSC_TRUE;
  PetscQuadrature q;
  PetscInt order;

  /* get the dimension of the problem from the DM */

  ierr = DMGetDimension(dm, &dim); CHKERRQ(ierr);

  /* Creating the FE for the displacement */
  ierr = PetscFECreateDefault(dm, dim, dim, simplex,"disp_",PETSC_DEFAULT,&fe_displacement);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) fe_displacement, "displacement");CHKERRQ(ierr);
  ierr = PetscFEGetQuadrature(fe_displacement,&q);CHKERRQ(ierr);
  ierr = PetscQuadratureGetOrder(q,&order);CHKERRQ(ierr);
  /* Creating the FE for the stress */
  ierr = PetscFECreateDefault(dm,dim,2*dim,simplex,"stress_",PETSC_DEFAULT,&fe_stress);CHKERRQ(ierr);
  ierr = PetscFESetQuadrature(fe_stress,q);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) fe_stress, "cauchy_stress");CHKERRQ(ierr);

 
    /* Discretization and boundary conditons: */
  ierr = DMGetDS(dm, &prob);CHKERRQ(ierr);
  ierr = PetscDSSetDiscretization(prob, 0, (PetscObject) fe_displacement); CHKERRQ(ierr);
  ierr = PetscDSSetDiscretization(prob, 1, (PetscObject) fe_stress);CHKERRQ(ierr);
  ierr = DMSetDS(dm, prob); CHKERRQ(ierr); 

  /* Define the boundaries */
    
  const PetscInt Ncomp = dim;
  const PetscInt Nfid = 1;
  PetscInt fid[Nfid];  /* fixed faces [numer of fixed faces] */
  
  PetscInt restrictall[3] = {0, 1, 2};  /* restricting all movement */

    
 
    fid[0] = 3;   /* fixed face */
    
    ierr = DMAddBoundary(dm, PETSC_TRUE, "fixed", "Face Sets",0,Ncomp,restrictall,(void (*)()) zero_vector, Nfid,fid,NULL);CHKERRQ(ierr);

    
  ierr = PetscFEDestroy(&fe_displacement); CHKERRQ(ierr);
  
  ierr = PetscFEDestroy(&fe_stress); CHKERRQ(ierr);
  
  return(0);
}
int main(int argc, char *argv[])
{
  DM dm, distributeddm;			/* problem definition */
  Vec u,expected_solution,projected_solution;			
  PetscViewer viewer;
  int dim;                        /* dimension of the anlysis */
  PetscErrorCode ierr;
  PetscMPIInt rank, numProcs;
  PetscPartitioner part;
  

  
  ierr = PetscInitialize(&argc,&argv,NULL,NULL);
  ierr = DMPlexCreateFromFile(PETSC_COMM_WORLD,"testcube.msh", PETSC_TRUE,&(dm));
  ierr = DMGetDimension(dm,&(dim));


  /* distribute the mesh */
  MPI_Comm_rank(PETSC_COMM_WORLD, &rank);
  MPI_Comm_size(PETSC_COMM_WORLD, &numProcs);
  DMPlexGetPartitioner(dm, &part);
  PetscPartitionerSetType(part, PETSCPARTITIONERPARMETIS);

  ierr = DMPlexDistribute(dm,0,NULL,&distributeddm); 

  if (distributeddm) {
    ierr=DMDestroy(&(dm));
    dm = distributeddm;
  }


  ierr = DMView(dm,PETSC_VIEWER_STDOUT_WORLD);
  ierr = DMSetMatType(dm, MATAIJ);CHKERRQ(ierr);


 
  ierr = SetupDiscretization(dm);CHKERRQ(ierr);
  
  ierr = DMPlexCreateClosureIndex(dm,NULL);CHKERRQ(ierr); 
  ierr = DMCreateGlobalVector(dm, &(u)); CHKERRQ(ierr);
  ierr = VecDuplicate(u,&expected_solution); CHKERRQ(ierr);
  ierr = VecDuplicate(u,&projected_solution); CHKERRQ(ierr);

  /* intitialize the fields:  */

  PetscErrorCode (*initial[2])(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar u[], void* ctx) = {initial_displacement_vector,zero_stress};
  ierr = DMProjectFunction(dm,0.0,initial,NULL, INSERT_ALL_VALUES, u); CHKERRQ(ierr);

  PetscErrorCode (*expected_sol[2])(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar u[], void* ctx) = {initial_displacement_vector,expected_stress_vector};
  ierr = DMProjectFunction(dm,0.0,expected_sol,NULL, INSERT_ALL_VALUES, expected_solution); CHKERRQ(ierr);

  ierr = PetscViewerVTKOpen(PETSC_COMM_WORLD,"expected_solution.vtu",FILE_MODE_WRITE,&viewer);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) expected_solution, "expected_fields"); CHKERRQ(ierr);
  ierr = VecView(expected_solution,viewer);CHKERRQ(ierr);
  ierr= PetscViewerDestroy(&viewer); CHKERRQ(ierr);
  
  

  /* project the fields: */

   void (*projection[2])(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[],
                         const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[],
                         const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[],
                         const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[],
                         PetscReal t, const PetscScalar x[], PetscInt numConstants,
                         const PetscScalar constants[], PetscScalar v[]) = {projectdisplacement, projectstress};


   ierr = DMProjectField(dm, 0.0, u, projection,INSERT_ALL_VALUES,projected_solution); CHKERRQ(ierr);
   ierr = PetscViewerVTKOpen(PETSC_COMM_WORLD,"projected_solution.vtu",FILE_MODE_WRITE,&viewer);CHKERRQ(ierr);
   ierr = PetscObjectSetName((PetscObject) projected_solution, "projected_fields"); CHKERRQ(ierr);
   ierr = VecView(projected_solution,viewer);CHKERRQ(ierr);
   ierr= PetscViewerDestroy(&viewer); CHKERRQ(ierr);


   VecDestroy(&u);
   VecDestroy(&expected_solution);
   VecDestroy(&projected_solution);
   DMDestroy(&dm);

   PetscFinalize();

  return 0;
}