#!/usr/bin/python
"""
author: francesco.caimmi@polimi.it

This code plays with DMPlex overlaps

"""


from __future__ import print_function
import argparse
import numpy as np
import sys
import petsc4py


def create_dm(fname=None):
    """
    Creates the dmplex object
    """
    if fname is None:
        #create a 2D box mesh
        dim = 2
        dm = PETSc.DMPlex().createBoxMesh(dim)
    else:
        #read from file
        dm = PETSc.DMPlex().createExodusFromFile(fname)
        dim = dm.getDimension()

    return dm

def create_section(dm):
    """
    Creates the default section

    """
    dim = dm.getDimension()
    #define a scalar field on vertices
    numFields = 1
    numComp = np.zeros(numFields, dtype=PETSc.IntType)
    numComp[0] = 1
    numDof = np.zeros(numFields*(dim+1), dtype=PETSc.IntType)
    numDof[0] = 1 #dofs on vertex only
    vector_section = dm.createSection(numComp, numDof)
    vector_section.setName('vector')
    dm.setDefaultSection(vector_section)
    return vector_section

def main(fname='cantilever.e', overlap=0):

    dm = create_dm(fname=fname)
    dm.view()
#    dm.setAdjacencyUseClosure(False)
#    dm.setAdjacencyUseCone(True)
    if PETSc.COMM_WORLD.size == 1:
        PETSc.Sys.Print('This example should be run on multiple processes')
        sys.exit(1)
    else:
        dm.distribute(overlap = overlap)

    section = create_section(dm)
    v = dm.createLocalVec()
    v.setName('rank')
    v.set(PETSc.COMM_WORLD.rank)
    viewer = PETSc.Viewer().createVTK('out.vtu')
    v.view(viewer)
    del viewer
    PETSc.Sys.Print('DM after distribution with overlap', overlap)
    dm.view()
#    PETSc.Sys.Print('Section after distribution with overlap', overlap)
#    section.view()
    PETSc.Sys.Print('Local vector length distribution with overlap', overlap)
    PETSc.Sys.syncPrint(v.getSize())
    PETSc.Sys.syncFlush()



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=
                                     'Playing with dmplex distrubute overlay')
    parser.add_argument('-o', '--overlap', help="Overlap value", type=int,
                        default=0)
    parser.add_argument('-f', '--file', help="Overlap value",
                        default='cantilever.e')
    options, petsc_args = parser.parse_known_args(sys.argv[1:])
    petsc4py.init(petsc_args)
    from petsc4py import PETSc
    options = vars(options)

    main(fname=options['file'], overlap=options['overlap'])