#include <petsc-private/patchimpl.h>   /*I      "petscdmpatch.h"   I*/
#include <petscdmda.h>

/*
Solver loop to update \tau:

  DMZoom(dmc, &dmz)
  DMRefine(dmz, &dmf),
  Scatter Xcoarse -> Xzoom,
  Interpolate Xzoom -> Xfine (note that this may be on subcomms),
  Smooth Xfine using two-step smoother
    normal smoother plus Kaczmarz---moves back and forth from dmzoom to dmfine
  Compute residual Rfine
  Restrict Rfine to Rzoom_restricted
  Scatter Rzoom_restricted -> Rcoarse_restricted
  Compute global residual Rcoarse
  TauCoarse = Rcoarse - Rcoarse_restricted
*/

#undef __FUNCT__
#define __FUNCT__ "DMPatchZoom"
/*
  DMPatchZoom - Create a version of the coarse patch (identified by rank) with halo on communicator commz

  Collective on DM

  Input Parameters:
  + dm - the DM
  . rank - the rank which holds the given patch
  - commz - the new communicator for the patch

  Output Parameters:
  + dmz  - the patch DM
  . sfz  - the PetscSF mapping the patch+halo to the zoomed version
  . sfzr - the PetscSF mapping the patch to the restricted zoomed version

  Level: intermediate

  Note: All processes in commz should have the same rank (could autosplit comm)

.seealso: DMPatchSolve()
*/
PetscErrorCode DMPatchZoom(DM dm, MatStencil lower, MatStencil upper, MPI_Comm commz, DM *dmz, PetscSF *sfz, PetscSF *sfzr)
{
  DMDAStencilType st;
  MatStencil      blower, bupper, loclower, locupper;
  Vec             X;
  IS              is;
  const PetscInt *ranges, *indices;
  PetscInt       *localPoints;
  PetscSFNode    *remotePoints;
  PetscInt        dim, dof;
  PetscInt        M, N, P, rM, rN, rP, halo = 1, sxb, syb, szb, sxr, syr, szr, exr, eyr, ezr, mxb, myb, mzb, i, j, k, q;
  PetscMPIInt     size;
  PetscErrorCode  ierr;

  PetscFunctionBegin;
  ierr = MPI_Comm_size(((PetscObject) dm)->comm, &size);CHKERRQ(ierr);
  if (commz == MPI_COMM_NULL) {
    /* Split communicator */
    SETERRQ(((PetscObject) dm)->comm, PETSC_ERR_SUP, "Not implemented");
  }
  /* Create patch DM */
  ierr = DMDAGetInfo(dm, &dim, &M, &N, &P, 0,0,0, &dof, 0,0,0,0, &st);CHKERRQ(ierr);

  /* Get piece for rank r, expanded by halo */
  bupper.i = PetscMin(M, upper.i + halo); blower.i = PetscMax(lower.i - halo, 0);
  bupper.j = PetscMin(N, upper.j + halo); blower.j = PetscMax(lower.j - halo, 0);
  bupper.k = PetscMin(P, upper.k + halo); blower.k = PetscMax(lower.k - halo, 0);
  rM = bupper.i - blower.i;
  rN = bupper.j - blower.j;
  rP = bupper.k - blower.k;

  ierr = DMDACreate(commz, dmz);CHKERRQ(ierr);
  ierr = DMDASetDim(*dmz, dim);CHKERRQ(ierr);
  ierr = DMDASetSizes(*dmz, rM, rN, rP);CHKERRQ(ierr);
  ierr = DMDASetNumProcs(*dmz, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE);CHKERRQ(ierr);
  ierr = DMDASetBoundaryType(*dmz, DMDA_BOUNDARY_NONE, DMDA_BOUNDARY_NONE, DMDA_BOUNDARY_NONE);CHKERRQ(ierr);
  ierr = DMDASetDof(*dmz, dof);CHKERRQ(ierr);
  ierr = DMDASetStencilType(*dmz, st);CHKERRQ(ierr);
  ierr = DMDASetStencilWidth(*dmz, 0);CHKERRQ(ierr);
  ierr = DMDASetOwnershipRanges(*dmz, PETSC_NULL, PETSC_NULL, PETSC_NULL);CHKERRQ(ierr);
  ierr = DMSetFromOptions(*dmz);CHKERRQ(ierr);
  ierr = DMSetUp(*dmz);CHKERRQ(ierr);
  ierr = DMDAGetCorners(*dmz, &sxb, &syb, &szb, &mxb, &myb, &mzb);CHKERRQ(ierr);
  sxr  = PetscMax(sxb,     lower.i - blower.i);
  syr  = PetscMax(syb,     lower.j - blower.j);
  szr  = PetscMax(szb,     lower.k - blower.k);
  exr  = PetscMin(sxb+mxb, upper.i - blower.i);
  eyr  = PetscMin(syb+myb, upper.j - blower.j);
  ezr  = PetscMin(szb+mzb, upper.k - blower.k);

  ierr = PetscMalloc2(rM*rN*rP,PetscInt,&localPoints,rM*rN*rP,PetscSFNode,&remotePoints);CHKERRQ(ierr);

  /* Create SF for restricted map */
  ierr = DMGetGlobalVector(dm, &X);CHKERRQ(ierr);
  ierr = PetscLayoutGetRanges(X->map, &ranges);CHKERRQ(ierr);
  loclower.i = blower.i + sxr; locupper.i = blower.i + exr;
  loclower.j = blower.j + syr; locupper.j = blower.j + eyr;
  loclower.k = blower.k + szr; locupper.k = blower.k + ezr;
  ierr = DMDACreatePatchIS(dm, &loclower, &locupper, &is);CHKERRQ(ierr);
  ierr = ISGetIndices(is, &indices);CHKERRQ(ierr);
  q = 0;
  for(k = szb; k < szb+mzb; ++k) {
    if ((k < szr) || (k >= ezr)) continue;
    for(j = syb; j < syb+myb; ++j) {
      if ((j < syr) || (j >= eyr)) continue;
      for(i = sxb; i < sxb+mxb; ++i) {
        const PetscInt lp = ((k-szb)*rN + (j-syb))*rM + i-sxb;
        PetscInt r;

        if ((i < sxr) || (i >= exr)) continue;
        localPoints[q]        = lp;
        ierr = PetscFindInt(indices[q], size+1, ranges, &r);CHKERRQ(ierr);
        remotePoints[q].rank  = r < 0 ? -(r+1) - 1 : r;
        remotePoints[q].index = indices[q] - ranges[remotePoints[q].rank];
        ++q;
      }
    }
  }
  ierr = ISRestoreIndices(is, &indices);CHKERRQ(ierr);
  ierr = ISDestroy(&is);CHKERRQ(ierr);
  ierr = PetscSFCreate(((PetscObject) dm)->comm, sfzr);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) *sfzr, "Restricted Map");CHKERRQ(ierr);
  ierr = PetscSFSetGraph(*sfzr, M*N*P, q, localPoints, PETSC_COPY_VALUES, remotePoints, PETSC_COPY_VALUES);CHKERRQ(ierr);

  /* Create SF for buffered map */
  loclower.i = blower.i + sxb; locupper.i = blower.i + sxb+mxb;
  loclower.j = blower.j + syb; locupper.j = blower.j + syb+myb;
  loclower.k = blower.k + szb; locupper.k = blower.k + szb+mzb;
  ierr = DMDACreatePatchIS(dm, &loclower, &locupper, &is);CHKERRQ(ierr);
  ierr = ISGetIndices(is, &indices);CHKERRQ(ierr);
  q = 0;
  for(k = szb; k < szb+mzb; ++k) {
    for(j = syb; j < syb+myb; ++j) {
      for(i = sxb; i < sxb+mxb; ++i, ++q) {
        PetscInt r;

        localPoints[q]        = q;
        ierr = PetscFindInt(indices[q], size+1, ranges, &r);CHKERRQ(ierr);
        remotePoints[q].rank  = r < 0 ? -(r+1) - 1 : r;
        remotePoints[q].index = indices[q] - ranges[remotePoints[q].rank];
      }
    }
  }
  ierr = ISRestoreIndices(is, &indices);CHKERRQ(ierr);
  ierr = ISDestroy(&is);CHKERRQ(ierr);
  ierr = PetscSFCreate(((PetscObject) dm)->comm, sfz);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) *sfz, "Buffered Map");CHKERRQ(ierr);
  ierr = PetscSFSetGraph(*sfz, M*N*P, q, localPoints, PETSC_COPY_VALUES, remotePoints, PETSC_COPY_VALUES);CHKERRQ(ierr);

  ierr = DMRestoreGlobalVector(dm, &X);CHKERRQ(ierr);
  ierr = PetscFree2(localPoints, remotePoints);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

typedef enum {PATCH_COMM_TYPE_WORLD = 0, PATCH_COMM_TYPE_SELF = 1} PatchCommType;

#undef __FUNCT__
#define __FUNCT__ "DMPatchSolve"
PetscErrorCode DMPatchSolve(DM dm)
{
  MPI_Comm       comm = ((PetscObject) dm)->comm;
  DM_Patch      *mesh = (DM_Patch *) dm->data;
  const char    *commTypes[] = {"world", "self", "PatchCommTypes", "PATCH_COMM_TYPE_", PETSC_NULL};
  PatchCommType  commType    = PATCH_COMM_TYPE_WORLD;
  MPI_Comm       commz       = comm;
  DM             dmc, dmz;
  PetscSF        sfz, sfzr;
  Vec            XC, XZ;
  MatStencil     patchSize, lower, upper;
  PetscInt       M, N, P, i, j, k, p = 0;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = PetscOptionsGetEnum(((PetscObject) dm)->prefix, "-dm_patch_zoom_comm", commTypes, &commType, PETSC_NULL);CHKERRQ(ierr);
  if (commType == PATCH_COMM_TYPE_WORLD) {
    commz = comm;
  } else if (commType == PATCH_COMM_TYPE_SELF) {
    commz = PETSC_COMM_SELF;
  }
  ierr = DMPatchGetCoarse(dm, &dmc);CHKERRQ(ierr);
  ierr = DMPatchGetPatchSize(dm, &patchSize);CHKERRQ(ierr);
  ierr = DMGetGlobalVector(dmc, &XC);CHKERRQ(ierr);
  /* Stupid loop over whole mesh can be replaced */
  ierr = DMDAGetInfo(dmc, 0, &M, &N, &P, 0,0,0,0,0,0,0,0,0);CHKERRQ(ierr);
  M    = PetscMax(M, 1);
  N    = PetscMax(N, 1);
  P    = PetscMax(P, 1);
  for(k = 0; k < P; k += PetscMax(patchSize.k, 1)) {
    for(j = 0; j < N; j += PetscMax(patchSize.j, 1)) {
      for(i = 0; i < M; i += PetscMax(patchSize.i, 1), ++p) {
        DM  dmf;
        Mat interpz;
        PetscScalar *xcarray, *xzarray;

        /* Zoom to coarse patch */
        lower.i = i; lower.j = j; lower.k = k;
        upper.i = i + patchSize.i; upper.j = j + patchSize.j; upper.k = k + patchSize.k;
        ierr = DMPatchZoom(dmc, lower, upper, commz, &dmz, &sfz, &sfzr);CHKERRQ(ierr);
        /* Debug */
        ierr = PetscPrintf(comm, "Patch %d: (%d, %d, %d)--(%d, %d, %d)\n", p, lower.i, lower.j, lower.k, upper.i, upper.j, upper.k);CHKERRQ(ierr);
        ierr = DMView(dmz, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
        ierr = PetscSFView(sfz,  PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
        ierr = PetscSFView(sfzr, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
        /* TODO: Need coarse and zoomed state vectors here: */
        xcarray = PETSC_NULL;
        xzarray = PETSC_NULL;
        /* Scatter Xcoarse -> Xzoom */
        ierr = DMGetGlobalVector(dmz, &XZ);CHKERRQ(ierr);
        ierr = VecGetArray(XC, &xcarray);CHKERRQ(ierr);
        ierr = VecGetArray(XZ, &xzarray);CHKERRQ(ierr);
        ierr = PetscSFBcastBegin(sfz, MPIU_SCALAR, xcarray, xzarray);CHKERRQ(ierr);
        ierr = PetscSFBcastEnd(sfz, MPIU_SCALAR, xcarray, xzarray);CHKERRQ(ierr);
        ierr = VecRestoreArray(XC, &xcarray);CHKERRQ(ierr);
        ierr = VecRestoreArray(XZ, &xzarray);CHKERRQ(ierr);
#if 0
        /* Interpolate Xzoom -> Xfine, note that this may be on subcomms */
        ierr = DMRefine(dmz, MPI_COMM_NULL, &dmf);CHKERRQ(ierr);
        ierr = DMCreateInterpolation(dmz, dmf, &interpz, PETSC_NULL);CHKERRQ(ierr);
        ierr = DMInterpolate(dmz, interpz, dmf);CHKERRQ(ierr);
        /* Smooth Xfine using two-step smoother, normal smoother plus Kaczmarz---moves back and forth from dmzoom to dmfine */
        /* Compute residual Rfine */
        /* Restrict Rfine to Rzoom_restricted */
#endif
        /* Scatter Rzoom_restricted -> Rcoarse_restricted */
        ierr = VecGetArray(XC, &xcarray);CHKERRQ(ierr);
        ierr = VecGetArray(XZ, &xzarray);CHKERRQ(ierr);
        ierr = PetscSFReduceBegin(sfzr, MPIU_SCALAR, xzarray, xcarray, MPI_SUM);CHKERRQ(ierr);
        ierr = PetscSFReduceEnd(sfzr, MPIU_SCALAR, xzarray, xcarray, MPI_SUM);CHKERRQ(ierr);
        ierr = VecRestoreArray(XC, &xcarray);CHKERRQ(ierr);
        ierr = VecRestoreArray(XZ, &xzarray);CHKERRQ(ierr);
        ierr = DMRestoreGlobalVector(dmz, &XZ);CHKERRQ(ierr);
        /* Compute global residual Rcoarse */
        /* TauCoarse = Rcoarse - Rcoarse_restricted */

        ierr = PetscSFDestroy(&sfz);CHKERRQ(ierr);
        ierr = PetscSFDestroy(&sfzr);CHKERRQ(ierr);
        ierr = DMDestroy(&dmz);CHKERRQ(ierr);
      }
    }
  }
  ierr = DMRestoreGlobalVector(dmc, &XC);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMPatchView_Ascii"
PetscErrorCode DMPatchView_Ascii(DM dm, PetscViewer viewer)
{
  DM_Patch         *mesh = (DM_Patch *) dm->data;
  PetscViewerFormat format;
  const char       *name;
  PetscErrorCode    ierr;

  PetscFunctionBegin;
  ierr = PetscViewerGetFormat(viewer, &format);CHKERRQ(ierr);
  /* if (format == PETSC_VIEWER_ASCII_INFO_DETAIL) */
  ierr = PetscObjectGetName((PetscObject) dm, &name);CHKERRQ(ierr);
  ierr = PetscViewerASCIIPrintf(viewer, "Patch DM %s\n", name);CHKERRQ(ierr);
  ierr = PetscViewerASCIIPushTab(viewer);CHKERRQ(ierr);
  ierr = PetscViewerASCIIPrintf(viewer, "Coarse DM\n");CHKERRQ(ierr);
  ierr = DMView(mesh->dmCoarse, viewer);CHKERRQ(ierr);
  ierr = PetscViewerASCIIPopTab(viewer);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMView_Patch"
PetscErrorCode DMView_Patch(DM dm, PetscViewer viewer)
{
  PetscBool      iascii, isbinary;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
  ierr = PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);CHKERRQ(ierr);
  ierr = PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERBINARY, &isbinary);CHKERRQ(ierr);
  if (iascii) {
    ierr = DMPatchView_Ascii(dm, viewer);CHKERRQ(ierr);
#if 0
  } else if (isbinary) {
    ierr = DMPatchView_Binary(dm, viewer);CHKERRQ(ierr);
#endif
  } else SETERRQ1(((PetscObject)viewer)->comm,PETSC_ERR_SUP,"Viewer type %s not supported by this mesh object", ((PetscObject)viewer)->type_name);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMDestroy_Patch"
PetscErrorCode DMDestroy_Patch(DM dm)
{
  DM_Patch      *mesh = (DM_Patch *) dm->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  if (--mesh->refct > 0) {PetscFunctionReturn(0);}
  ierr = DMDestroy(&mesh->dmCoarse);CHKERRQ(ierr);
  /* This was originally freed in DMDestroy(), but that prevents reference counting of backend objects */
  ierr = PetscFree(mesh);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMSetUp_Patch"
PetscErrorCode DMSetUp_Patch(DM dm)
{
  DM_Patch      *mesh = (DM_Patch *) dm->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  ierr = DMSetUp(mesh->dmCoarse);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMCreateGlobalVector_Patch"
PetscErrorCode DMCreateGlobalVector_Patch(DM dm, Vec *g)
{
  DM_Patch      *mesh = (DM_Patch *) dm->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  ierr = DMCreateGlobalVector(mesh->dmCoarse, g);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMCreateLocalVector_Patch"
PetscErrorCode DMCreateLocalVector_Patch(DM dm, Vec *l)
{
  DM_Patch      *mesh = (DM_Patch *) dm->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  ierr = DMCreateLocalVector(mesh->dmCoarse, l);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMCreateSubDM_Patch"
PetscErrorCode DMCreateSubDM_Patch(DM dm, PetscInt numFields, PetscInt fields[], IS *is, DM *subdm)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  SETERRQ(((PetscObject) dm)->comm, PETSC_ERR_SUP, "Tell me to code this");
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMPatchGetCoarse"
PetscErrorCode DMPatchGetCoarse(DM dm, DM *dmCoarse)
{
  DM_Patch *mesh = (DM_Patch *) dm->data;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  *dmCoarse = mesh->dmCoarse;
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMPatchGetPatchSize"
PetscErrorCode DMPatchGetPatchSize(DM dm, MatStencil *patchSize)
{
  DM_Patch *mesh = (DM_Patch *) dm->data;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidPointer(patchSize, 2);
  *patchSize = mesh->patchSize;
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "DMPatchSetPatchSize"
PetscErrorCode DMPatchSetPatchSize(DM dm, MatStencil patchSize)
{
  DM_Patch *mesh = (DM_Patch *) dm->data;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  mesh->patchSize = patchSize;
  PetscFunctionReturn(0);
}