# --------------------------------------------------------------------

cdef extern from * nogil:

    ctypedef enum PetscDMStagStencilType"DMStagStencilType":
        DMSTAG_STENCIL_STAR
        DMSTAG_STENCIL_BOX
        DMSTAG_STENCIL_NONE

    ctypedef enum PetscDMStagStencilLocation"DMStagStencilLocation":
        DMSTAG_NULL_LOCATION
        DMSTAG_BACK_DOWN_LEFT
        DMSTAG_BACK_DOWN
        DMSTAG_BACK_DOWN_RIGHT
        DMSTAG_BACK_LEFT
        DMSTAG_BACK
        DMSTAG_BACK_RIGHT
        DMSTAG_BACK_UP_LEFT
        DMSTAG_BACK_UP
        DMSTAG_BACK_UP_RIGHT
        DMSTAG_DOWN_LEFT
        DMSTAG_DOWN
        DMSTAG_DOWN_RIGHT
        DMSTAG_LEFT
        DMSTAG_ELEMENT
        DMSTAG_RIGHT
        DMSTAG_UP_LEFT
        DMSTAG_UP
        DMSTAG_UP_RIGHT
        DMSTAG_FRONT_DOWN_LEFT
        DMSTAG_FRONT_DOWN
        DMSTAG_FRONT_DOWN_RIGHT
        DMSTAG_FRONT_LEFT
        DMSTAG_FRONT
        DMSTAG_FRONT_RIGHT
        DMSTAG_FRONT_UP_LEFT
        DMSTAG_FRONT_UP
        DMSTAG_FRONT_UP_RIGHT

    PetscErrorCode DMStagCreate1d(MPI_Comm, PetscDMBoundaryType, PetscInt, PetscInt, PetscInt, PetscDMStagStencilType, PetscInt, const PetscInt[], PetscDM*)
    PetscErrorCode DMStagCreate2d(MPI_Comm, PetscDMBoundaryType, PetscDMBoundaryType, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt, PetscDMStagStencilType, PetscInt, const PetscInt[], const PetscInt[], PetscDM*)
    PetscErrorCode DMStagCreate3d(MPI_Comm, PetscDMBoundaryType, PetscDMBoundaryType, PetscDMBoundaryType, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt, PetscInt, PetscDMStagStencilType, PetscInt, const PetscInt[], const PetscInt[], const PetscInt[], PetscDM*)

    PetscErrorCode DMStagGetCorners(PetscDM, PetscInt*, PetscInt*, PetscInt*, PetscInt*, PetscInt*, PetscInt*, PetscInt*, PetscInt*, PetscInt*)
    PetscErrorCode DMStagGetGhostCorners(PetscDM, PetscInt*, PetscInt*, PetscInt*, PetscInt*, PetscInt*, PetscInt*)
    PetscErrorCode DMStagGetLocalSizes(PetscDM, PetscInt*, PetscInt*, PetscInt*)
    PetscErrorCode DMStagGetEntriesPerElement(PetscDM, PetscInt*)

    PetscErrorCode DMStagGetDOF(PetscDM, PetscInt*, PetscInt*, PetscInt*, PetscInt*)
    PetscErrorCode DMStagGetNumRanks(PetscDM, PetscInt*, PetscInt*, PetscInt*)
    PetscErrorCode DMStagGetGlobalSizes(PetscDM, PetscInt*, PetscInt*, PetscInt*)
    PetscErrorCode DMStagGetBoundaryTypes(PetscDM, PetscDMBoundaryType*, PetscDMBoundaryType*, PetscDMBoundaryType*)
    PetscErrorCode DMStagGetStencilWidth(PetscDM, PetscInt*)
    PetscErrorCode DMStagGetStencilType(PetscDM, PetscDMStagStencilType*)
    PetscErrorCode DMStagGetOwnershipRanges(PetscDM, const PetscInt*[], const PetscInt*[], const PetscInt*[])

    PetscErrorCode DMStagSetDOF(PetscDM, PetscInt, PetscInt, PetscInt, PetscInt)
    PetscErrorCode DMStagSetNumRanks(PetscDM, PetscInt, PetscInt, PetscInt)
    PetscErrorCode DMStagSetGlobalSizes(PetscDM, PetscInt, PetscInt, PetscInt)
    PetscErrorCode DMStagSetBoundaryTypes(PetscDM, PetscDMBoundaryType, PetscDMBoundaryType, PetscDMBoundaryType)
    PetscErrorCode DMStagSetStencilWidth(PetscDM, PetscInt)
    PetscErrorCode DMStagSetStencilType(PetscDM, PetscDMStagStencilType)
    PetscErrorCode DMStagSetOwnershipRanges(PetscDM, const PetscInt[], const PetscInt[], const PetscInt[])

    PetscErrorCode DMStagGetLocationSlot(PetscDM, PetscDMStagStencilLocation, PetscInt, PetscInt*)
    PetscErrorCode DMStagGetLocationDOF(PetscDM, PetscDMStagStencilLocation, PetscInt*)
    PetscErrorCode DMStagGetProductCoordinateLocationSlot(PetscDM, PetscDMStagStencilLocation, PetscInt*)

    PetscErrorCode DMStagGetIsFirstRank(PetscDM, PetscBool*, PetscBool*, PetscBool*)
    PetscErrorCode DMStagGetIsLastRank(PetscDM, PetscBool*, PetscBool*, PetscBool*)

    PetscErrorCode DMStagSetUniformCoordinatesExplicit(PetscDM, PetscReal, PetscReal, PetscReal, PetscReal, PetscReal, PetscReal)
    PetscErrorCode DMStagSetUniformCoordinatesProduct(PetscDM, PetscReal, PetscReal, PetscReal, PetscReal, PetscReal, PetscReal)
    PetscErrorCode DMStagSetCoordinateDMType(PetscDM, PetscDMType)
    PetscErrorCode DMStagSetUniformCoordinates(PetscDM, PetscReal, PetscReal, PetscReal, PetscReal, PetscReal, PetscReal)

    PetscErrorCode DMStagCreateCompatibleDMStag(PetscDM, PetscInt, PetscInt, PetscInt, PetscInt, PetscDM*)
    PetscErrorCode DMStagVecSplitToDMDA(PetscDM, PetscVec, PetscDMStagStencilLocation, PetscInt, PetscDM*, PetscVec*)
    PetscErrorCode DMStagMigrateVec(PetscDM, PetscVec, PetscDM, PetscVec)

# --------------------------------------------------------------------

cdef inline PetscDMStagStencilType asStagStencil(object stencil) \
    except <PetscDMStagStencilType>(-1):
    if isinstance(stencil, str):
        if   stencil == "star":  return DMSTAG_STENCIL_STAR
        elif stencil == "box":   return DMSTAG_STENCIL_BOX
        elif stencil == "none":  return DMSTAG_STENCIL_NONE
        else: raise ValueError("unknown stencil type: %s" % stencil)
    return stencil

cdef inline object toStagStencil(PetscDMStagStencilType stype):
    if   stype == DMSTAG_STENCIL_STAR:  return "star"
    elif stype == DMSTAG_STENCIL_BOX:   return "box"
    elif stype == DMSTAG_STENCIL_NONE:  return "none"

cdef inline PetscDMStagStencilLocation asStagStencilLocation(object stencil_location) \
    except <PetscDMStagStencilLocation>(-1):
    if isinstance(stencil_location, str):
        if   stencil_location == "null":                return DMSTAG_NULL_LOCATION
        elif stencil_location == "back_down_left":      return DMSTAG_BACK_DOWN_LEFT
        elif stencil_location == "back_down":           return DMSTAG_BACK_DOWN
        elif stencil_location == "back_down_right":     return DMSTAG_BACK_DOWN_RIGHT
        elif stencil_location == "back_left":           return DMSTAG_BACK_LEFT
        elif stencil_location == "back":                return DMSTAG_BACK
        elif stencil_location == "back_right":          return DMSTAG_BACK_RIGHT
        elif stencil_location == "back_up_left":        return DMSTAG_BACK_UP_LEFT
        elif stencil_location == "back_up":             return DMSTAG_BACK_UP
        elif stencil_location == "back_up_right":       return DMSTAG_BACK_UP_RIGHT
        elif stencil_location == "down_left":           return DMSTAG_DOWN_LEFT
        elif stencil_location == "down":                return DMSTAG_DOWN
        elif stencil_location == "down_right":          return DMSTAG_DOWN_RIGHT
        elif stencil_location == "left":                return DMSTAG_LEFT
        elif stencil_location == "element":             return DMSTAG_ELEMENT
        elif stencil_location == "right":               return DMSTAG_RIGHT
        elif stencil_location == "up_left":             return DMSTAG_UP_LEFT
        elif stencil_location == "up":                  return DMSTAG_UP
        elif stencil_location == "up_right":            return DMSTAG_UP_RIGHT
        elif stencil_location == "front_down_left":     return DMSTAG_FRONT_DOWN_LEFT
        elif stencil_location == "front_down":          return DMSTAG_FRONT_DOWN
        elif stencil_location == "front_down_right":    return DMSTAG_FRONT_DOWN_RIGHT
        elif stencil_location == "front_left":          return DMSTAG_FRONT_LEFT
        elif stencil_location == "front":               return DMSTAG_FRONT
        elif stencil_location == "front_right":         return DMSTAG_FRONT_RIGHT
        elif stencil_location == "front_up_left":       return DMSTAG_FRONT_UP_LEFT
        elif stencil_location == "front_up":            return DMSTAG_FRONT_UP
        elif stencil_location == "front_up_right":      return DMSTAG_FRONT_UP_RIGHT
        else: raise ValueError("unknown stencil location type: %s" % stencil_location)
    return stencil_location


cdef inline PetscInt asStagDims(dims,
                                PetscInt *_M,
                                PetscInt *_N,
                                PetscInt *_P) except? -1:
    cdef PetscInt dim = PETSC_DECIDE
    cdef object M=None, N=None, P=None
    dims = tuple(dims)
    dim = <PetscInt>len(dims)
    if   dim == 0: pass
    elif dim == 1: M, = dims
    elif dim == 2: M, N = dims
    elif dim == 3: M, N, P = dims
    if dim >= 1: _M[0] = asInt(M)
    if dim >= 2: _N[0] = asInt(N)
    if dim >= 3: _P[0] = asInt(P)
    return dim

cdef inline tuple toStagDims(PetscInt dim,
                             PetscInt M,
                             PetscInt N,
                             PetscInt P):
    if   dim == 0: return ()
    elif dim == 1: return (toInt(M),)
    elif dim == 2: return (toInt(M), toInt(N))
    elif dim == 3: return (toInt(M), toInt(N), toInt(P))

cdef inline PetscInt asDofs(dofs,
                            PetscInt *_dof0,
                            PetscInt *_dof1,
                            PetscInt *_dof2,
                            PetscInt *_dof3) except? -1:
    cdef PetscInt ndofs = PETSC_DECIDE
    cdef object dof0=None, dof1=None, dof2=None, dof3=None
    dofs = tuple(dofs)
    ndofs = <PetscInt>len(dofs)
    if ndofs == 2:   dof0, dof1 = dofs
    elif ndofs == 3: dof0, dof1, dof2 = dofs
    elif ndofs == 4: dof0, dof1, dof2, dof3 = dofs
    if ndofs >= 2: _dof0[0] = asInt(dof0)
    if ndofs >= 2: _dof1[0] = asInt(dof1)
    if ndofs >= 3: _dof2[0] = asInt(dof2)
    if ndofs >= 4: _dof3[0] = asInt(dof3)
    return ndofs

cdef inline tuple toDofs(PetscInt ndofs,
                         PetscInt dof0,
                         PetscInt dof1,
                         PetscInt dof2,
                         PetscInt dof3):
    if ndofs == 2: return (toInt(dof0), toInt(dof1))
    elif ndofs == 3: return (toInt(dof0), toInt(dof1), toInt(dof2))
    elif ndofs == 4: return (toInt(dof0), toInt(dof1), toInt(dof2), toInt(dof3))

cdef inline tuple asStagOwnershipRanges(object ownership_ranges,
                                        PetscInt dim,
                                        PetscInt *m, PetscInt *n, PetscInt *p,
                                        PetscInt **_x,
                                        PetscInt **_y,
                                        PetscInt **_z):
    cdef object ranges = list(ownership_ranges)
    cdef PetscInt rdim = <PetscInt>len(ranges)
    cdef PetscInt nlx=0, nly=0, nlz=0
    if dim == PETSC_DECIDE: dim = rdim
    elif dim != rdim: raise ValueError(
        "number of dimensions %d and number ownership ranges %d" %
        (toInt(dim), toInt(rdim)))
    if dim >= 1:
        ranges[0] = iarray_i(ranges[0], &nlx, _x)
        if m[0] == PETSC_DECIDE: m[0] = nlx
        elif m[0] != nlx: raise ValueError(
            "ownership range size %d and number or processors %d" %
            (toInt(nlx), toInt(m[0])))
    if dim >= 2:
        ranges[1] = iarray_i(ranges[1], &nly, _y)
        if n[0] == PETSC_DECIDE: n[0] = nly
        elif n[0] != nly: raise ValueError(
            "ownership range size %d and number or processors %d" %
            (toInt(nly), toInt(n[0])))
    if dim >= 3:
        ranges[2] = iarray_i(ranges[2], &nlz, _z)
        if p[0] == PETSC_DECIDE: p[0] = nlz
        elif p[0] != nlz: raise ValueError(
            "ownership range size %d and number or processors %d" %
            (toInt(nlz), toInt(p[0])))
    return tuple(ranges)


cdef inline tuple toStagOwnershipRanges(PetscInt dim,
                                        PetscInt m, PetscInt n, PetscInt p,
                                        const PetscInt *lx,
                                        const PetscInt *ly,
                                        const PetscInt *lz):
    # Returns tuple of arrays containing ownership ranges as Python arrays
    ranges = [array_i(m, lx)]
    if dim > 1:
        ranges.append(array_i(n, ly))
    if dim > 2:
        ranges.append(array_i(p, lz))
    return tuple(ranges)

cdef inline object toStagBoundary(PetscDMBoundaryType btype):
    if   btype == DM_BOUNDARY_NONE:       return "none"
    elif btype == DM_BOUNDARY_PERIODIC:   return "periodic"
    elif btype == DM_BOUNDARY_GHOSTED:    return "ghosted"

cdef inline tuple toStagBoundaryTypes(PetscInt dim, PetscDMBoundaryType btx, PetscDMBoundaryType bty, PetscDMBoundaryType btz):
    if dim == 1: return (toStagBoundary(btx),)
    if dim == 2: return (toStagBoundary(btx), toStagBoundary(bty))
    if dim == 3: return (toStagBoundary(btx), toStagBoundary(bty), toStagBoundary(btz))

# --------------------------------------------------------------------
