static char help[] = "EXOTest11 \n\n";

#include <petsc.h>
#include "exodusII.h"
#include <petsc/private/dmimpl.h>
#include "exodusIO.h"

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc,char **argv) {  
  char              ifilename[PETSC_MAX_PATH_LEN];
  PetscBool         interpolate=PETSC_TRUE,flg1=PETSC_FALSE,flg2=PETSC_FALSE;
  PetscInt          order=1,p,pStart,pEnd;
  PetscInt          d,sdim;


  DM                dmUAS,dmU,dmAlpha,dmSigma,dmUAlpha,dmDist;
  PetscSection      sectionUAS;
  PetscInt          fieldU        = 0;
  PetscInt          fieldAlpha    = 2;
  PetscInt          fieldSigma    = 1;
  PetscInt          fieldUAlpha[] = {0,2};

  IS                ISU,ISAlpha,ISSigma,ISUAlpha;
  PetscSF           migrationSF;
  IS                csIS,cellIS;
  const PetscInt    *csID,*cellID;
  PetscInt          numCells,set,cell;

  /* 
    dof layout: cell,face,edge,vertex 
    ordered by increasing height in the DAG
  */
  PetscInt          dofUP1Tri[]     = {2,0,0};
  PetscInt          dofAlphaP1Tri[] = {1,0,0};
  PetscInt          dofUP2Tri[]     = {2,2,0};
  PetscInt          dofAlphaP2Tri[] = {1,1,0};
  PetscInt          dofUP1Quad[]     = {2,0,0};
  PetscInt          dofAlphaP1Quad[] = {1,0,0};
  PetscInt          dofUP2Quad[]     = {2,2,2};
  PetscInt          dofAlphaP2Quad[] = {1,1,1};

  PetscInt          dofSigma2D[]      = {0,0,3};

  PetscInt          dofUP1Tet[]     = {3,0,0,0};
  PetscInt          dofAlphaP1Tet[] = {1,0,0,0};
  PetscInt          dofUP2Tet[]     = {3,3,0,0};
  PetscInt          dofAlphaP2Tet[] = {1,1,0,0};
  PetscInt          dofUP1Hex[]     = {3,0,0,0};
  PetscInt          dofAlphaP1Hex[] = {1,0,0,0};
  PetscInt          dofUP2Hex[]     = {3,3,3,3};
  PetscInt          dofAlphaP2Hex[] = {1,1,1,1};

  PetscInt          dofSigma3D[]     = {0,0,0,6};

  PetscInt          *dofU,*dofAlpha,*dofSigma;

  PetscMPIInt       rank,numproc;
  PetscErrorCode    ierr;

  PetscSection      defaultSection;
  PetscInt          pStartDepth[4],pEndDepth[4];
  PetscInt          closureSize,*closure;
  PetscInt          numCS;

  PetscInitialize(&argc,&argv,NULL,help);
  ex_opts(EX_VERBOSE+EX_DEBUG);

  PetscOptionsBegin(PETSC_COMM_WORLD,"","options","none");
  ierr = PetscOptionsString("-i","filename to read","",ifilename,ifilename,sizeof(ifilename),&flg1);CHKERRQ(ierr);
  ierr = PetscOptionsInt("-order","Polynomial order","",order,&order,NULL);CHKERRQ(ierr);
  PetscOptionsEnd();
  
  MPI_Comm_rank(PETSC_COMM_WORLD,&rank);
  MPI_Comm_size(PETSC_COMM_WORLD,&numproc);

  /* 
    Read mesh from file
  */
  ex_opts(EX_VERBOSE+EX_DEBUG);
  ierr = DMPlexCreateFromFile(PETSC_COMM_WORLD,ifilename,interpolate,&dmUAS);CHKERRQ(ierr);

  /* 
    Create the main section containning all fields
  */
  ierr = PetscSectionCreate(PETSC_COMM_WORLD,&sectionUAS);CHKERRQ(ierr);
  ierr = PetscSectionSetNumFields(sectionUAS,3);CHKERRQ(ierr);
  ierr = PetscSectionSetFieldName(sectionUAS,fieldU,"U");CHKERRQ(ierr);
  ierr = PetscSectionSetFieldName(sectionUAS,fieldAlpha,"Alpha");CHKERRQ(ierr);
  ierr = PetscSectionSetFieldName(sectionUAS,fieldSigma,"Sigma");CHKERRQ(ierr);
  ierr = DMPlexGetChart(dmUAS,&pStart,&pEnd);CHKERRQ(ierr);
  /*
    For some reason, PetscSectionSetChart needs to be called AFTER
    PetscSectionSetNumFields and friends or calls to PetscSectionSetDof
    will fail
  */
  ierr = PetscSectionSetChart(sectionUAS,pStart,pEnd);CHKERRQ(ierr);
  ierr = DMGetDimension(dmUAS,&sdim);CHKERRQ(ierr);
  for (d = 0; d < sdim+1; d++) {
    ierr = DMPlexGetDepthStratum(dmUAS,d,&pStartDepth[d],&pEndDepth[d]);CHKERRQ(ierr);
  }

  ierr = PetscSectionSetFieldComponents(sectionUAS,0,sdim);CHKERRQ(ierr);
  ierr = PetscSectionSetFieldComponents(sectionUAS,1,1);CHKERRQ(ierr);
  ierr = PetscSectionSetFieldComponents(sectionUAS,2,sdim*(sdim+1)/2);CHKERRQ(ierr);

  /*
    Going through cell sets then cells, and setting up storage for the sections
  */
  ierr = DMGetLabelSize(dmUAS,"Cell Sets",&numCS);
  ierr = DMGetLabelIdIS(dmUAS,"Cell Sets",&csIS);CHKERRQ(ierr);
  ierr = ISGetIndices(csIS,&csID);CHKERRQ(ierr);
  for (set = 0; set < numCS; set++) {
    ierr = DMGetStratumSize(dmUAS,"Cell Sets",csID[set],&numCells);CHKERRQ(ierr);
    if (numCells > 0) {
      if (sdim == 2) {
        dofSigma = dofSigma2D;
      } else {
        dofSigma = dofSigma3D;
      }

      /*
        Identify cell type based on closure size only. This works for Tri/Tet/Quad/Hex meshes
        Not sure if it would also be enough to identify more exotic elements like pyramid or prisms...
      */

      ierr = DMGetStratumIS(dmUAS,"Cell Sets",csID[set],&cellIS);CHKERRQ(ierr);
      ierr = DMGetStratumSize(dmUAS,"Cell Sets",csID[set],&numCells);CHKERRQ(ierr);
      ierr = ISGetIndices(cellIS,&cellID);CHKERRQ(ierr);
      ierr = DMPlexGetTransitiveClosure(dmUAS,cellID[0],PETSC_TRUE,&closureSize,&closure);CHKERRQ(ierr);
      switch(closureSize) {
        case 7 : /* Tri */
        if (order == 1) {
          dofU     = dofUP1Tri;
          dofAlpha = dofAlphaP1Tri;
        } else {
          dofU     = dofUP2Tri;
          dofAlpha = dofAlphaP2Tri;          
        }
        break;
        case 9 : /* Quad */
        if (order == 1) {
          dofU     = dofUP1Quad;
          dofAlpha = dofAlphaP1Quad;
        } else {
          dofU     = dofUP2Quad;
          dofAlpha = dofAlphaP2Quad;          
        }
        break;
        case 15 : /* Tet */
        if (order == 1) {
          dofU     = dofUP1Tet;
          dofAlpha = dofAlphaP1Tet;
        } else {
          dofU     = dofUP2Tet;
          dofAlpha = dofAlphaP2Tet;          
        }
        break;
        case 27 : /* Hex */
        if (order == 1) {
          dofU     = dofUP1Hex;
          dofAlpha = dofAlphaP1Hex;
        } else {
          dofU     = dofUP2Hex;
          dofAlpha = dofAlphaP2Hex;          
        }
        break;
        default :
          SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Unkown element with closure size %d\n",closureSize);
      }
      ierr = DMPlexRestoreTransitiveClosure(dmUAS,cellID[0],PETSC_TRUE,&closureSize,&closure);CHKERRQ(ierr);

      for (cell = 0; cell < numCells; cell++){
        ierr = DMPlexGetTransitiveClosure(dmUAS,cellID[cell],PETSC_TRUE,&closureSize,&closure);CHKERRQ(ierr);

        for (p = 0; p < closureSize; p++) {
          /* 
            Find depth of p
          */
          for (d = 0; d < sdim+1; d++) {
            if ((closure[2*p] >= pStartDepth[d]) && (closure[2*p] < pEndDepth[d])) {
              ierr = PetscSectionSetDof(sectionUAS,closure[2*p],dofU[d]+dofAlpha[d]+dofSigma[d]);CHKERRQ(ierr);
              ierr = PetscSectionSetFieldDof(sectionUAS,closure[2*p],fieldU,dofU[d]);CHKERRQ(ierr);
              ierr = PetscSectionSetFieldDof(sectionUAS,closure[2*p],fieldAlpha,dofAlpha[d]);CHKERRQ(ierr);
              ierr = PetscSectionSetFieldDof(sectionUAS,closure[2*p],fieldSigma,dofSigma[d]);CHKERRQ(ierr);        
            }
          }
        }
        ierr = DMPlexRestoreTransitiveClosure(dmUAS,cellID[cell],PETSC_TRUE,&closureSize,&closure);CHKERRQ(ierr);
      }
    }
  }
  ierr = ISRestoreIndices(csIS,&csID);CHKERRQ(ierr);
  ierr = ISDestroy(&csIS);CHKERRQ(ierr);

  ierr = PetscSectionSetUp(sectionUAS);CHKERRQ(ierr);
  ierr = DMSetDefaultSection(dmUAS,sectionUAS);CHKERRQ(ierr);
  ierr = DMSetFromOptions(dmUAS);CHKERRQ(ierr);

  /* 
    Distribute main DM
  */
  ierr = DMSetUseNatural(dmUAS,PETSC_TRUE);CHKERRQ(ierr);
  ierr = DMPlexDistribute(dmUAS,0,&migrationSF,&dmDist);CHKERRQ(ierr);
  if (dmDist) {
    ierr = DMDestroy(&dmUAS);CHKERRQ(ierr);
    dmUAS   = dmDist;
  }

  ierr = DMGetDefaultSection(dmUAS,&defaultSection);CHKERRQ(ierr);CHKERRQ(ierr);
  ierr = PetscSectionView(defaultSection,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);

  /* 
    get DM and IS for each field of dmUAS
  */
  ierr = DMCreateSubDM(dmUAS,1,&fieldU,&ISU,&dmU);CHKERRQ(ierr);
  ierr = DMCreateSubDM(dmUAS,1,&fieldAlpha,&ISAlpha,&dmAlpha);CHKERRQ(ierr);
  ierr = DMCreateSubDM(dmUAS,1,&fieldSigma,&ISSigma,&dmSigma);CHKERRQ(ierr);
  ierr = DMCreateSubDM(dmUAS,2,fieldUAlpha,&ISUAlpha,&dmUAlpha);CHKERRQ(ierr);


  ierr = DMDestroy(&dmUAlpha);CHKERRQ(ierr);
  ierr = DMDestroy(&dmSigma);CHKERRQ(ierr);
  ierr = DMDestroy(&dmAlpha);CHKERRQ(ierr);
  ierr = DMDestroy(&dmU);CHKERRQ(ierr);
  ierr = DMDestroy(&dmUAS);CHKERRQ(ierr);
  PetscFinalize();
  return 0;
}