from petsc4py import PETSc
import unittest

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


class BaseTestDMStag:
    COMM = PETSc.COMM_WORLD
    STENCIL = PETSc.DMStag.StencilType.BOX
    SWIDTH = 1
    PROC_SIZES = None
    OWNERSHIP_RANGES = None

    def setUp(self):
        dim = len(self.SIZES)
        self.da = PETSc.DMStag().create(
            dim,
            dofs=self.DOFS,
            sizes=self.SIZES,
            boundary_types=self.BOUNDARY,
            stencil_type=self.STENCIL,
            stencil_width=self.SWIDTH,
            comm=self.COMM,
            proc_sizes=self.PROC_SIZES,
            ownership_ranges=self.OWNERSHIP_RANGES,
            setUp=True,
        )

        self.directda = PETSc.DMStag().create(dim)
        self.directda.setStencilType(self.STENCIL)
        self.directda.setStencilWidth(self.SWIDTH)
        self.directda.setBoundaryTypes(self.BOUNDARY)
        self.directda.setDof(self.DOFS)
        self.directda.setGlobalSizes(self.SIZES)
        if self.PROC_SIZES is not None:
            self.directda.setProcSizes(self.PROC_SIZES)
        if self.OWNERSHIP_RANGES is not None:
            self.directda.setOwnershipRanges(self.OWNERSHIP_RANGES)
        self.directda.setUp()

    def tearDown(self):
        self.da = None
        self.directda = None
        PETSc.garbage_cleanup()

    def testCoordinates(self):
        self.da.setCoordinateDMType('stag')
        self.da.setUniformCoordinates(0, 1, 0, 1, 0, 1)
        self.da.setUniformCoordinatesExplicit(0, 1, 0, 1, 0, 1)
        cda = self.da.getCoordinateDM()
        datype = cda.getType()
        self.assertEqual(datype, 'stag')
        cda.destroy()

        c = self.da.getCoordinatesLocal()
        self.da.setCoordinatesLocal(c)
        gc = self.da.getCoordinatesLocal()
        self.assertEqual(c.max()[1], gc.max()[1])
        self.assertEqual(c.min()[1], gc.min()[1])

        c = self.da.getCoordinates()
        self.da.setCoordinates(c)
        gc = self.da.getCoordinates()
        self.assertEqual(c.max()[1], gc.max()[1])
        self.assertEqual(c.min()[1], gc.min()[1])

        self.directda.setCoordinateDMType('product')
        self.directda.setUniformCoordinates(0, 1, 0, 1, 0, 1)
        self.directda.setUniformCoordinatesProduct(0, 1, 0, 1, 0, 1)
        cda = self.directda.getCoordinateDM()
        datype = cda.getType()
        self.assertEqual(datype, 'product')
        cda.destroy()

    def testGetVec(self):
        vg = self.da.getGlobalVec()
        vl = self.da.getLocalVec()

        vg.set(1.0)
        self.assertEqual(vg.max()[1], 1.0)
        self.assertEqual(vg.min()[1], 1.0)
        self.da.globalToLocal(vg, vl)
        self.assertEqual(vl.max()[1], 1.0)
        self.assertTrue(vl.min()[1] in (1.0, 0.0))

        vl.set(2.0)
        self.da.localToGlobal(vl, vg)
        self.assertEqual(vg.max()[1], 2.0)
        self.assertTrue(vg.min()[1] in (2.0, 0.0))

        self.da.restoreGlobalVec(vg)
        self.da.restoreLocalVec(vl)

    def testGetOther(self):
        _ = self.da.getLGMap()
        _ = self.directda.getLGMap()

    def testDof(self):
        dim = self.da.getDim()
        dofs = self.da.getDof()
        if dim == 1:
            dof0 = self.da.getLocationDof('left')
            dof1 = self.da.getLocationDof('element')
            self.assertEqual(dofs[0], dof0)
            self.assertEqual(dofs[1], dof1)
        if dim == 2:
            dof0 = self.da.getLocationDof('down_left')
            dof1 = self.da.getLocationDof('left')
            dof2 = self.da.getLocationDof('element')
            self.assertEqual(dofs[0], dof0)
            self.assertEqual(dofs[1], dof1)
            self.assertEqual(dofs[2], dof2)
        if dim == 3:
            dof0 = self.da.getLocationDof('back_down_right')
            dof1 = self.da.getLocationDof('down_left')
            dof2 = self.da.getLocationDof('left')
            dof3 = self.da.getLocationDof('element')
            self.assertEqual(dofs[0], dof0)
            self.assertEqual(dofs[1], dof1)
            self.assertEqual(dofs[2], dof2)
            self.assertEqual(dofs[3], dof3)

    def testMigrateVec(self):
        vec = self.da.createGlobalVec()
        dmTo = self.da.createCompatibleDMStag(self.NEWDOF)
        vecTo = dmTo.createGlobalVec()
        self.da.migrateVec(vec, dmTo, vecTo)

    def testDMDAInterface(self):
        return
        self.da.setCoordinateDMType('stag')
        self.da.setUniformCoordinates(0, 1, 0, 1, 0, 1)
        dim = self.da.getDim()
        dofs = self.da.getDof()
        vec = self.da.createGlobalVec()
        if dim == 1:
            da, davec = self.da.VecSplitToDMDA(vec, 'left', -dofs[0])
            da, davec = self.da.VecSplitToDMDA(vec, 'element', -dofs[1])
        if dim == 2:
            da, davec = self.da.VecSplitToDMDA(vec, 'down_left', -dofs[0])
            da, davec = self.da.VecSplitToDMDA(vec, 'down_left', -dofs[1])
            da, davec = self.da.VecSplitToDMDA(vec, 'down_left', -dofs[2])
        if dim == 3:
            da, davec = self.da.VecSplitToDMDA(vec, 'back_down_right', -dofs[0])
            da, davec = self.da.VecSplitToDMDA(vec, 'down_left', -dofs[1])
            da, davec = self.da.VecSplitToDMDA(vec, 'left', -dofs[2])
            da, davec = self.da.VecSplitToDMDA(vec, 'element', -dofs[3])


GHOSTED = PETSc.DM.BoundaryType.GHOSTED
PERIODIC = PETSc.DM.BoundaryType.PERIODIC
NONE = PETSc.DM.BoundaryType.NONE

SCALE = 4


class BaseTestDMStag_1D(BaseTestDMStag):
    SIZES = [
        100 * SCALE,
    ]
    BOUNDARY = [
        NONE,
    ]


class BaseTestDMStag_2D(BaseTestDMStag):
    SIZES = [9 * SCALE, 11 * SCALE]
    BOUNDARY = [NONE, NONE]


class BaseTestDMStag_3D(BaseTestDMStag):
    SIZES = [6 * SCALE, 7 * SCALE, 8 * SCALE]
    BOUNDARY = [NONE, NONE, NONE]


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


class TestDMStag_1D_W0_N11(BaseTestDMStag_1D, unittest.TestCase):
    SWIDTH = 0
    DOFS = (1, 1)
    NEWDOF = (2, 1)


class TestDMStag_1D_W0_N21(BaseTestDMStag_1D, unittest.TestCase):
    SWIDTH = 0
    DOFS = (2, 1)
    NEWDOF = (2, 2)


class TestDMStag_1D_W0_N12(BaseTestDMStag_1D, unittest.TestCase):
    SWIDTH = 0
    DOFS = (1, 2)
    NEWDOF = (2, 2)


class TestDMStag_1D_W2_N11(BaseTestDMStag_1D, unittest.TestCase):
    SWIDTH = 2
    DOFS = (1, 1)
    NEWDOF = (2, 1)


class TestDMStag_1D_W2_N21(BaseTestDMStag_1D, unittest.TestCase):
    SWIDTH = 2
    DOFS = (2, 1)
    NEWDOF = (2, 2)


class TestDMStag_1D_W2_N12(BaseTestDMStag_1D, unittest.TestCase):
    SWIDTH = 2
    DOFS = (1, 2)
    NEWDOF = (2, 2)


class TestDMStag_2D_W0_N112(BaseTestDMStag_2D, unittest.TestCase):
    DOFS = (1, 1, 2)
    SWIDTH = 0
    NEWDOF = (2, 2, 2)


class TestDMStag_2D_W2_N112(BaseTestDMStag_2D, unittest.TestCase):
    DOFS = (1, 1, 2)
    SWIDTH = 2
    NEWDOF = (2, 2, 2)


class TestDMStag_2D_PXY(BaseTestDMStag_2D, unittest.TestCase):
    SIZES = [13 * SCALE, 17 * SCALE]
    DOFS = (1, 1, 2)
    SWIDTH = 5
    BOUNDARY = (PERIODIC,) * 2
    NEWDOF = (2, 2, 2)


class TestDMStag_2D_GXY(BaseTestDMStag_2D, unittest.TestCase):
    SIZES = [13 * SCALE, 17 * SCALE]
    DOFS = (1, 1, 2)
    SWIDTH = 5
    BOUNDARY = (GHOSTED,) * 2
    NEWDOF = (2, 2, 2)


class TestDMStag_3D_W0_N1123(BaseTestDMStag_3D, unittest.TestCase):
    DOFS = (1, 1, 2, 3)
    SWIDTH = 0
    NEWDOF = (2, 2, 3, 3)


class TestDMStag_3D_W2_N1123(BaseTestDMStag_3D, unittest.TestCase):
    DOFS = (1, 1, 2, 3)
    SWIDTH = 2
    NEWDOF = (2, 2, 3, 3)


class TestDMStag_3D_PXYZ(BaseTestDMStag_3D, unittest.TestCase):
    SIZES = [11 * SCALE, 13 * SCALE, 17 * SCALE]
    DOFS = (1, 1, 2, 3)
    NEWDOF = (2, 2, 3, 3)
    SWIDTH = 3
    BOUNDARY = (PERIODIC,) * 3


class TestDMStag_3D_GXYZ(BaseTestDMStag_3D, unittest.TestCase):
    SIZES = [11 * SCALE, 13 * SCALE, 17 * SCALE]
    DOFS = (1, 1, 2, 3)
    NEWDOF = (2, 2, 3, 3)
    SWIDTH = 3
    BOUNDARY = (GHOSTED,) * 3


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

DIM = (1, 2, 3)
DOF0 = (0, 1, 2)
DOF1 = (0, 1, 2)
DOF2 = (0, 1, 2)
DOF3 = (0, 1, 2)
BOUNDARY_TYPE = ('none', 'ghosted', 'periodic')
STENCIL_TYPE = ('none', 'star', 'box')
STENCIL_WIDTH = (0, 1, 2, 3)


class TestDMStagCreate(unittest.TestCase):
    pass


counter = 0
for dim in DIM:
    for dof0 in DOF0:
        for dof1 in DOF1:
            for dof2 in DOF2:
                if dim == 1 and dof2 > 0:
                    continue
                for dof3 in DOF3:
                    if dim == 2 and dof3 > 0:
                        continue
                    if dof0 == 0 and dof1 == 0 and dof2 == 0 and dof3 == 0:
                        continue
                    dofs = [dof0, dof1, dof2, dof3][: dim + 1]
                    for boundary in BOUNDARY_TYPE:
                        if boundary == 'periodic':
                            continue  # XXX broken
                        for stencil in STENCIL_TYPE:
                            if stencil == 'none' and boundary != 'none':
                                continue
                            for width in STENCIL_WIDTH:
                                if stencil == 'none' and width > 0:
                                    continue
                                if stencil in ['star', 'box'] and width == 0:
                                    continue
                                kargs = {
                                    'dim': dim,
                                    'dofs': dofs,
                                    'boundary_type': boundary,
                                    'stencil_type': stencil,
                                    'stencil_width': width,
                                }

                                def testCreate(self, kargs=kargs):
                                    kargs = dict(kargs)
                                    cda = PETSc.DMStag().create(
                                        kargs['dim'],
                                        dofs=kargs['dofs'],
                                        sizes=[
                                            8 * SCALE,
                                        ]
                                        * kargs['dim'],
                                        boundary_types=[
                                            kargs['boundary_type'],
                                        ]
                                        * kargs['dim'],
                                        stencil_type=kargs['stencil_type'],
                                        stencil_width=kargs['stencil_width'],
                                        setUp=True,
                                    )

                                    dda = PETSc.DMStag().create(kargs['dim'])
                                    dda.setStencilType(kargs['stencil_type'])
                                    dda.setStencilWidth(kargs['stencil_width'])
                                    dda.setBoundaryTypes(
                                        [
                                            kargs['boundary_type'],
                                        ]
                                        * kargs['dim']
                                    )
                                    dda.setDof(kargs['dofs'])
                                    dda.setGlobalSizes(
                                        [
                                            8 * SCALE,
                                        ]
                                        * kargs['dim']
                                    )
                                    dda.setUp()

                                    cdim = cda.getDim()
                                    cdof = cda.getDof()
                                    cgsizes = cda.getGlobalSizes()
                                    clsizes = cda.getLocalSizes()
                                    cboundary = cda.getBoundaryTypes()
                                    cstencil_type = cda.getStencilType()
                                    cstencil_width = cda.getStencilWidth()
                                    centries_per_element = cda.getEntriesPerElement()
                                    cstarts, csizes, cnextra = cda.getCorners()
                                    cisLastRank = cda.getIsLastRank()
                                    cisFirstRank = cda.getIsFirstRank()
                                    cownershipranges = cda.getOwnershipRanges()
                                    cprocsizes = cda.getProcSizes()

                                    ddim = dda.getDim()
                                    ddof = dda.getDof()
                                    dgsizes = dda.getGlobalSizes()
                                    dlsizes = dda.getLocalSizes()
                                    dboundary = dda.getBoundaryTypes()
                                    dstencil_type = dda.getStencilType()
                                    dstencil_width = dda.getStencilWidth()
                                    dentries_per_element = dda.getEntriesPerElement()
                                    dstarts, dsizes, dnextra = dda.getCorners()
                                    disLastRank = dda.getIsLastRank()
                                    disFirstRank = dda.getIsFirstRank()
                                    downershipranges = dda.getOwnershipRanges()
                                    dprocsizes = dda.getProcSizes()

                                    self.assertEqual(cdim, kargs['dim'])
                                    self.assertEqual(cdof, tuple(kargs['dofs']))
                                    self.assertEqual(
                                        cboundary,
                                        tuple(
                                            [
                                                kargs['boundary_type'],
                                            ]
                                            * kargs['dim']
                                        ),
                                    )
                                    self.assertEqual(
                                        cstencil_type, kargs['stencil_type']
                                    )
                                    self.assertEqual(
                                        cstencil_width, kargs['stencil_width']
                                    )
                                    self.assertEqual(
                                        cgsizes,
                                        tuple(
                                            [
                                                8 * SCALE,
                                            ]
                                            * kargs['dim']
                                        ),
                                    )

                                    self.assertEqual(cdim, ddim)
                                    self.assertEqual(cdof, ddof)
                                    self.assertEqual(cgsizes, dgsizes)
                                    self.assertEqual(clsizes, dlsizes)
                                    self.assertEqual(cboundary, dboundary)
                                    self.assertEqual(cstencil_type, dstencil_type)
                                    self.assertEqual(cstencil_width, dstencil_width)
                                    self.assertEqual(
                                        centries_per_element, dentries_per_element
                                    )
                                    self.assertEqual(cstarts, dstarts)
                                    self.assertEqual(csizes, dsizes)
                                    self.assertEqual(cnextra, dnextra)
                                    self.assertEqual(cisLastRank, disLastRank)
                                    self.assertEqual(cisFirstRank, disFirstRank)
                                    self.assertEqual(cprocsizes, dprocsizes)
                                    for co, do in zip(
                                        cownershipranges, downershipranges
                                    ):
                                        for i, j in zip(co, do):
                                            self.assertEqual(i, j)

                                    self.assertEqual(cdim + 1, len(cdof))
                                    self.assertEqual(cdim, len(cgsizes))
                                    self.assertEqual(cdim, len(clsizes))
                                    self.assertEqual(cdim, len(cboundary))
                                    self.assertEqual(cdim, len(cstarts))
                                    self.assertEqual(cdim, len(csizes))
                                    self.assertEqual(cdim, len(cnextra))
                                    self.assertEqual(cdim, len(cisLastRank))
                                    self.assertEqual(cdim, len(cisLastRank))
                                    if cdim == 1:
                                        self.assertEqual(
                                            centries_per_element, cdof[0] + cdof[1]
                                        )
                                    if cdim == 2:
                                        self.assertEqual(
                                            centries_per_element,
                                            cdof[0] + 2 * cdof[1] + cdof[2],
                                        )
                                    if cdim == 3:
                                        self.assertEqual(
                                            centries_per_element,
                                            cdof[0]
                                            + 3 * cdof[1]
                                            + 3 * cdof[2]
                                            + cdof[3],
                                        )
                                    for i in range(cdim):
                                        self.assertEqual(csizes[i], clsizes[i])
                                        if cisLastRank[i]:
                                            self.assertEqual(cnextra[i], 1)
                                        if cnextra[i] == 1:
                                            self.assertTrue(cisLastRank[i])
                                        if cisFirstRank[i]:
                                            self.assertEqual(cstarts[i], 0)
                                    self.assertEqual(
                                        len(cprocsizes), len(cownershipranges)
                                    )
                                    self.assertEqual(len(cprocsizes), cdim)
                                    for i, m in enumerate(cprocsizes):
                                        self.assertEqual(m, len(cownershipranges[i]))
                                    dda.destroy()
                                    cda.destroy()

                                setattr(
                                    TestDMStagCreate,
                                    'testCreate%05d' % counter,
                                    testCreate,
                                )

                                del testCreate
                                counter += 1

del counter, dim, dofs, dof0, dof1, dof2, dof3, boundary, stencil, width

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

if __name__ == '__main__':
    unittest.main()

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