#include <petscdmda.h>                 /*I  "petscdmda.h"  I*/
#include <petsc/private/dmswarmimpl.h> /*I  "petscdmswarm.h"  I*/
#include "../src/dm/impls/swarm/data_bucket.h"

static PetscErrorCode private_PetscViewerCreate_XDMF(MPI_Comm comm, const char filename[], PetscViewer *v)
{
  long int      *bytes;
  PetscContainer container;
  PetscViewer    viewer;

  PetscFunctionBegin;
  PetscCall(PetscViewerCreate(comm, &viewer));
  PetscCall(PetscViewerSetType(viewer, PETSCVIEWERASCII));
  PetscCall(PetscViewerFileSetMode(viewer, FILE_MODE_WRITE));
  PetscCall(PetscViewerFileSetName(viewer, filename));

  PetscCall(PetscMalloc1(1, &bytes));
  bytes[0] = 0;
  PetscCall(PetscContainerCreate(comm, &container));
  PetscCall(PetscContainerSetPointer(container, (void *)bytes));
  PetscCall(PetscObjectCompose((PetscObject)viewer, "XDMFViewerContext", (PetscObject)container));

  /* write xdmf header */
  PetscCall(PetscViewerASCIIPrintf(viewer, "<?xml version=\"1.0\" encoding=\"utf-8\"?>\n"));
  PetscCall(PetscViewerASCIIPrintf(viewer, "<Xdmf xmlns:xi=\"http://www.w3.org/2001/XInclude/\" Version=\"2.99\">\n"));
  /* write xdmf domain */
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "<Domain>\n"));
  *v = viewer;
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode private_PetscViewerDestroy_XDMF(PetscViewer *v)
{
  PetscViewer    viewer;
  DM             dm = NULL;
  long int      *bytes;
  PetscContainer container = NULL;

  PetscFunctionBegin;
  if (!v) PetscFunctionReturn(PETSC_SUCCESS);
  viewer = *v;

  PetscCall(PetscObjectQuery((PetscObject)viewer, "DMSwarm", (PetscObject *)&dm));
  if (dm) {
    PetscCall(PetscViewerASCIIPrintf(viewer, "</Grid>\n"));
    PetscCall(PetscViewerASCIIPopTab(viewer));
  }

  /* close xdmf header */
  PetscCall(PetscViewerASCIIPrintf(viewer, "</Domain>\n"));
  PetscCall(PetscViewerASCIIPopTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "</Xdmf>\n"));

  PetscCall(PetscObjectQuery((PetscObject)viewer, "XDMFViewerContext", (PetscObject *)&container));
  if (container) {
    PetscCall(PetscContainerGetPointer(container, &bytes));
    PetscCall(PetscFree(bytes));
    PetscCall(PetscContainerDestroy(&container));
  }
  PetscCall(PetscViewerDestroy(&viewer));
  *v = NULL;
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode private_CreateDataFileNameXDMF(const char filename[], char dfilename[])
{
  const char dot_xmf[] = ".xmf";
  size_t     len;
  char       viewername_minus_ext[PETSC_MAX_PATH_LEN];
  PetscBool  flg;

  PetscFunctionBegin;
  PetscCall(PetscStrendswith(filename, dot_xmf, &flg));
  PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "File extension must be %s", dot_xmf);
  PetscCall(PetscStrncpy(viewername_minus_ext, filename, sizeof(viewername_minus_ext)));
  PetscCall(PetscStrlen(filename, &len));
  len -= sizeof(dot_xmf) - 1;
  if (sizeof(viewername_minus_ext) > len) viewername_minus_ext[len] = '\0';
  PetscCall(PetscSNPrintf(dfilename, PETSC_MAX_PATH_LEN - 1, "%s_swarm_fields.pbin", viewername_minus_ext));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode private_DMSwarmView_XDMF(DM dm, PetscViewer viewer)
{
  DMSwarmCellDM  celldm;
  PetscBool      isswarm = PETSC_FALSE;
  const char    *viewername;
  char           datafile[PETSC_MAX_PATH_LEN];
  char          *datafilename;
  PetscViewer    fviewer;
  PetscInt       k, ng, dim, Nfc;
  Vec            dvec;
  long int      *bytes     = NULL;
  PetscContainer container = NULL;
  const char    *dmname, **coordFields;

  PetscFunctionBegin;
  PetscCall(PetscObjectQuery((PetscObject)viewer, "XDMFViewerContext", (PetscObject *)&container));
  PetscCheck(container, PetscObjectComm((PetscObject)viewer), PETSC_ERR_SUP, "Valid to find attached data XDMFViewerContext");
  PetscCall(PetscContainerGetPointer(container, &bytes));

  PetscCall(PetscObjectTypeCompare((PetscObject)dm, DMSWARM, &isswarm));
  PetscCheck(isswarm, PetscObjectComm((PetscObject)viewer), PETSC_ERR_SUP, "Only valid for DMSwarm");

  PetscCall(PetscObjectCompose((PetscObject)viewer, "DMSwarm", (PetscObject)dm));

  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscObjectGetName((PetscObject)dm, &dmname));
  if (!dmname) PetscCall(DMGetOptionsPrefix(dm, &dmname));
  if (!dmname) {
    PetscCall(PetscViewerASCIIPrintf(viewer, "<Grid Name=\"DMSwarm\" GridType=\"Uniform\">\n"));
  } else {
    PetscCall(PetscViewerASCIIPrintf(viewer, "<Grid Name=\"DMSwarm[%s]\" GridType=\"Uniform\">\n", dmname));
  }

  /* create a sub-viewer for topology, geometry and all data fields */
  /* name is viewer.name + "_swarm_fields.pbin" */
  PetscCall(PetscViewerCreate(PetscObjectComm((PetscObject)viewer), &fviewer));
  PetscCall(PetscViewerSetType(fviewer, PETSCVIEWERBINARY));
  PetscCall(PetscViewerBinarySetSkipHeader(fviewer, PETSC_TRUE));
  PetscCall(PetscViewerBinarySetSkipInfo(fviewer, PETSC_TRUE));
  PetscCall(PetscViewerFileSetMode(fviewer, FILE_MODE_WRITE));

  PetscCall(PetscViewerFileGetName(viewer, &viewername));
  PetscCall(private_CreateDataFileNameXDMF(viewername, datafile));
  PetscCall(PetscViewerFileSetName(fviewer, datafile));
  PetscCall(PetscStrrchr(datafile, '/', &datafilename));

  PetscCall(DMSwarmGetSize(dm, &ng));

  /* write topology header */
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "<Topology Dimensions=\"%" PetscInt_FMT "\" TopologyType=\"Mixed\">\n", ng));
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "<DataItem Format=\"Binary\" Endian=\"Big\" DataType=\"Int\" Dimensions=\"%" PetscInt_FMT "\" Seek=\"%ld\">\n", ng * 3, bytes[0]));
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "%s\n", datafilename));
  PetscCall(PetscViewerASCIIPopTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "</DataItem>\n"));
  PetscCall(PetscViewerASCIIPopTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "</Topology>\n"));
  PetscCall(PetscViewerASCIIPopTab(viewer));

  /* write topology data */
  for (k = 0; k < ng; k++) {
    PetscInt pvertex[3];

    pvertex[0] = 1;
    pvertex[1] = 1;
    pvertex[2] = k;
    PetscCall(PetscViewerBinaryWrite(fviewer, pvertex, 3, PETSC_INT));
  }
  bytes[0] += sizeof(PetscInt) * ng * 3;

  /* write geometry header */
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(DMGetDimension(dm, &dim));
  switch (dim) {
  case 1:
    SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "No support for 1D");
  case 2:
    PetscCall(PetscViewerASCIIPrintf(viewer, "<Geometry Type=\"XY\">\n"));
    break;
  case 3:
    PetscCall(PetscViewerASCIIPrintf(viewer, "<Geometry Type=\"XYZ\">\n"));
    break;
  }
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "<DataItem Format=\"Binary\" Endian=\"Big\" DataType=\"Float\" Precision=\"8\" Dimensions=\"%" PetscInt_FMT " %" PetscInt_FMT "\" Seek=\"%ld\">\n", ng, dim, bytes[0]));
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "%s\n", datafilename));
  PetscCall(PetscViewerASCIIPopTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "</DataItem>\n"));
  PetscCall(PetscViewerASCIIPopTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "</Geometry>\n"));
  PetscCall(PetscViewerASCIIPopTab(viewer));

  PetscCall(DMSwarmGetCellDMActive(dm, &celldm));
  PetscCall(DMSwarmCellDMGetCoordinateFields(celldm, &Nfc, &coordFields));
  PetscCheck(Nfc == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "We only support a single coordinate field right now, not %" PetscInt_FMT, Nfc);

  /* write geometry data */
  PetscCall(DMSwarmCreateGlobalVectorFromField(dm, coordFields[0], &dvec));
  PetscCall(VecView(dvec, fviewer));
  PetscCall(DMSwarmDestroyGlobalVectorFromField(dm, coordFields[0], &dvec));
  bytes[0] += sizeof(PetscReal) * ng * dim;

  PetscCall(PetscViewerDestroy(&fviewer));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode private_VecView_Swarm_XDMF(Vec x, PetscViewer viewer)
{
  long int      *bytes     = NULL;
  PetscContainer container = NULL;
  const char    *viewername;
  char           datafile[PETSC_MAX_PATH_LEN];
  char          *datafilename;
  PetscViewer    fviewer;
  PetscInt       N, bs;
  const char    *vecname;
  char           fieldname[PETSC_MAX_PATH_LEN];

  PetscFunctionBegin;
  PetscCall(PetscObjectQuery((PetscObject)viewer, "XDMFViewerContext", (PetscObject *)&container));
  PetscCheck(container, PetscObjectComm((PetscObject)viewer), PETSC_ERR_SUP, "Unable to find attached data XDMFViewerContext");
  PetscCall(PetscContainerGetPointer(container, &bytes));
  PetscCall(PetscViewerFileGetName(viewer, &viewername));
  PetscCall(private_CreateDataFileNameXDMF(viewername, datafile));

  /* re-open a sub-viewer for all data fields */
  /* name is viewer.name + "_swarm_fields.pbin" */
  PetscCall(PetscViewerCreate(PetscObjectComm((PetscObject)viewer), &fviewer));
  PetscCall(PetscViewerSetType(fviewer, PETSCVIEWERBINARY));
  PetscCall(PetscViewerBinarySetSkipHeader(fviewer, PETSC_TRUE));
  PetscCall(PetscViewerBinarySetSkipInfo(fviewer, PETSC_TRUE));
  PetscCall(PetscViewerFileSetMode(fviewer, FILE_MODE_APPEND));
  PetscCall(PetscViewerFileSetName(fviewer, datafile));
  PetscCall(PetscStrrchr(datafile, '/', &datafilename));

  PetscCall(VecGetSize(x, &N));
  PetscCall(VecGetBlockSize(x, &bs));
  N = N / bs;
  PetscCall(PetscObjectGetName((PetscObject)x, &vecname));
  if (!vecname) {
    PetscCall(PetscSNPrintf(fieldname, PETSC_MAX_PATH_LEN - 1, "swarmfield_%d", ((PetscObject)x)->tag));
  } else {
    PetscCall(PetscSNPrintf(fieldname, PETSC_MAX_PATH_LEN - 1, "%s", vecname));
  }

  /* write data header */
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "<Attribute Center=\"Node\" Name=\"%s\" Type=\"None\">\n", fieldname));
  PetscCall(PetscViewerASCIIPushTab(viewer));
  if (bs == 1) {
    PetscCall(PetscViewerASCIIPrintf(viewer, "<DataItem Format=\"Binary\" Endian=\"Big\" DataType=\"Float\" Precision=\"8\" Dimensions=\"%" PetscInt_FMT "\" Seek=\"%ld\">\n", N, bytes[0]));
  } else {
    PetscCall(PetscViewerASCIIPrintf(viewer, "<DataItem Format=\"Binary\" Endian=\"Big\" DataType=\"Float\" Precision=\"8\" Dimensions=\"%" PetscInt_FMT " %" PetscInt_FMT "\" Seek=\"%ld\">\n", N, bs, bytes[0]));
  }
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "%s\n", datafilename));
  PetscCall(PetscViewerASCIIPopTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "</DataItem>\n"));
  PetscCall(PetscViewerASCIIPopTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "</Attribute>\n"));
  PetscCall(PetscViewerASCIIPopTab(viewer));

  /* write data */
  PetscCall(VecView(x, fviewer));
  bytes[0] += sizeof(PetscReal) * N * bs;

  PetscCall(PetscViewerDestroy(&fviewer));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode private_ISView_Swarm_XDMF(IS is, PetscViewer viewer)
{
  long int      *bytes     = NULL;
  PetscContainer container = NULL;
  const char    *viewername;
  char           datafile[PETSC_MAX_PATH_LEN];
  char          *datafilename;
  PetscViewer    fviewer;
  PetscInt       N, bs;
  const char    *vecname;
  char           fieldname[PETSC_MAX_PATH_LEN];

  PetscFunctionBegin;
  PetscCall(PetscObjectQuery((PetscObject)viewer, "XDMFViewerContext", (PetscObject *)&container));
  PetscCheck(container, PetscObjectComm((PetscObject)viewer), PETSC_ERR_SUP, "Unable to find attached data XDMFViewerContext");
  PetscCall(PetscContainerGetPointer(container, &bytes));
  PetscCall(PetscViewerFileGetName(viewer, &viewername));
  PetscCall(private_CreateDataFileNameXDMF(viewername, datafile));

  /* re-open a sub-viewer for all data fields */
  /* name is viewer.name + "_swarm_fields.pbin" */
  PetscCall(PetscViewerCreate(PetscObjectComm((PetscObject)viewer), &fviewer));
  PetscCall(PetscViewerSetType(fviewer, PETSCVIEWERBINARY));
  PetscCall(PetscViewerBinarySetSkipHeader(fviewer, PETSC_TRUE));
  PetscCall(PetscViewerBinarySetSkipInfo(fviewer, PETSC_TRUE));
  PetscCall(PetscViewerFileSetMode(fviewer, FILE_MODE_APPEND));
  PetscCall(PetscViewerFileSetName(fviewer, datafile));
  PetscCall(PetscStrrchr(datafile, '/', &datafilename));

  PetscCall(ISGetSize(is, &N));
  PetscCall(ISGetBlockSize(is, &bs));
  N = N / bs;
  PetscCall(PetscObjectGetName((PetscObject)is, &vecname));
  if (!vecname) {
    PetscCall(PetscSNPrintf(fieldname, PETSC_MAX_PATH_LEN - 1, "swarmfield_%d", ((PetscObject)is)->tag));
  } else {
    PetscCall(PetscSNPrintf(fieldname, PETSC_MAX_PATH_LEN - 1, "%s", vecname));
  }

  /* write data header */
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "<Attribute Center=\"Node\" Name=\"%s\" Type=\"None\">\n", fieldname));
  PetscCall(PetscViewerASCIIPushTab(viewer));
  if (bs == 1) {
    PetscCall(PetscViewerASCIIPrintf(viewer, "<DataItem Format=\"Binary\" Endian=\"Big\" DataType=\"Int\" Precision=\"4\" Dimensions=\"%" PetscInt_FMT "\" Seek=\"%ld\">\n", N, bytes[0]));
  } else {
    PetscCall(PetscViewerASCIIPrintf(viewer, "<DataItem Format=\"Binary\" Endian=\"Big\" DataType=\"Int\" Precision=\"4\" Dimensions=\"%" PetscInt_FMT " %" PetscInt_FMT "\" Seek=\"%ld\">\n", N, bs, bytes[0]));
  }
  PetscCall(PetscViewerASCIIPushTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "%s\n", datafilename));
  PetscCall(PetscViewerASCIIPopTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "</DataItem>\n"));
  PetscCall(PetscViewerASCIIPopTab(viewer));
  PetscCall(PetscViewerASCIIPrintf(viewer, "</Attribute>\n"));
  PetscCall(PetscViewerASCIIPopTab(viewer));

  /* write data */
  PetscCall(ISView(is, fviewer));
  bytes[0] += sizeof(PetscInt) * N * bs;

  PetscCall(PetscViewerDestroy(&fviewer));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@C
  DMSwarmViewFieldsXDMF - Write a selection of DMSwarm fields to an XDMF3 file

  Collective

  Input Parameters:
+ dm              - the `DMSWARM`
. filename        - the file name of the XDMF file (must have the extension .xmf)
. nfields         - the number of fields to write into the XDMF file
- field_name_list - array of length nfields containing the textual name of fields to write

  Level: beginner

  Note:
  Only fields registered with data type `PETSC_DOUBLE` or `PETSC_INT` can be written into the file

.seealso: `DM`, `DMSWARM`, `DMSwarmViewXDMF()`
@*/
PetscErrorCode DMSwarmViewFieldsXDMF(DM dm, const char filename[], PetscInt nfields, const char *field_name_list[])
{
  Vec         dvec;
  PetscInt    f, N;
  PetscViewer viewer;

  PetscFunctionBegin;
  PetscCall(private_PetscViewerCreate_XDMF(PetscObjectComm((PetscObject)dm), filename, &viewer));
  PetscCall(private_DMSwarmView_XDMF(dm, viewer));
  PetscCall(DMSwarmGetLocalSize(dm, &N));
  for (f = 0; f < nfields; f++) {
    void         *data;
    PetscDataType type;

    PetscCall(DMSwarmGetField(dm, field_name_list[f], NULL, &type, &data));
    PetscCall(DMSwarmRestoreField(dm, field_name_list[f], NULL, &type, &data));
    if (type == PETSC_DOUBLE) {
      PetscCall(DMSwarmCreateGlobalVectorFromField(dm, field_name_list[f], &dvec));
      PetscCall(PetscObjectSetName((PetscObject)dvec, field_name_list[f]));
      PetscCall(private_VecView_Swarm_XDMF(dvec, viewer));
      PetscCall(DMSwarmDestroyGlobalVectorFromField(dm, field_name_list[f], &dvec));
    } else if (type == PETSC_INT) {
      IS              is;
      const PetscInt *idx;

      PetscCall(DMSwarmGetField(dm, field_name_list[f], NULL, &type, &data));
      idx = (const PetscInt *)data;

      PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)dm), N, idx, PETSC_USE_POINTER, &is));
      PetscCall(PetscObjectSetName((PetscObject)is, field_name_list[f]));
      PetscCall(private_ISView_Swarm_XDMF(is, viewer));
      PetscCall(ISDestroy(&is));
      PetscCall(DMSwarmRestoreField(dm, field_name_list[f], NULL, &type, &data));
    } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Can only write PETSC_INT and PETSC_DOUBLE");
  }
  PetscCall(private_PetscViewerDestroy_XDMF(&viewer));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
  DMSwarmViewXDMF - Write `DMSWARM` fields to an XDMF3 file

  Collective

  Input Parameters:
+ dm       - the `DMSWARM`
- filename - the file name of the XDMF file (must have the extension .xmf)

  Level: beginner

  Note:
  Only fields user registered with data type `PETSC_DOUBLE` or `PETSC_INT` will be written into the file

  Developer Note:
  This should be removed and replaced with the standard use of `PetscViewer`

.seealso: `DM`, `DMSWARM`, `DMSwarmViewFieldsXDMF()`
@*/
PetscErrorCode DMSwarmViewXDMF(DM dm, const char filename[])
{
  DM_Swarm   *swarm = (DM_Swarm *)dm->data;
  Vec         dvec;
  PetscInt    f;
  PetscViewer viewer;

  PetscFunctionBegin;
  PetscCall(private_PetscViewerCreate_XDMF(PetscObjectComm((PetscObject)dm), filename, &viewer));
  PetscCall(private_DMSwarmView_XDMF(dm, viewer));
  for (f = 4; f < swarm->db->nfields; f++) { /* only examine user defined fields - the first 4 are internally created by DMSwarmPIC */
    DMSwarmDataField field;

    /* query field type - accept all those of type PETSC_DOUBLE */
    field = swarm->db->field[f];
    if (field->petsc_type == PETSC_DOUBLE) {
      PetscCall(DMSwarmCreateGlobalVectorFromField(dm, field->name, &dvec));
      PetscCall(PetscObjectSetName((PetscObject)dvec, field->name));
      PetscCall(private_VecView_Swarm_XDMF(dvec, viewer));
      PetscCall(DMSwarmDestroyGlobalVectorFromField(dm, field->name, &dvec));
    } else if (field->petsc_type == PETSC_INT) {
      IS              is;
      PetscInt        N;
      const PetscInt *idx;
      void           *data;

      PetscCall(DMSwarmGetLocalSize(dm, &N));
      PetscCall(DMSwarmGetField(dm, field->name, NULL, NULL, &data));
      idx = (const PetscInt *)data;

      PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)dm), N, idx, PETSC_USE_POINTER, &is));
      PetscCall(PetscObjectSetName((PetscObject)is, field->name));
      PetscCall(private_ISView_Swarm_XDMF(is, viewer));
      PetscCall(ISDestroy(&is));
      PetscCall(DMSwarmRestoreField(dm, field->name, NULL, NULL, &data));
    }
  }
  PetscCall(private_PetscViewerDestroy_XDMF(&viewer));
  PetscFunctionReturn(PETSC_SUCCESS);
}