import petsc4py, sys
petsc4py.init(sys.argv)
from petsc4py import PETSc
import mpi4py.rc
mpi4py.rc.finalize=False
#from mpi4py import MPI

import numpy as np


pComm = PETSc.COMM_WORLD
pRank = pComm.getRank()
pSize = pComm.Get_size()
#PETSc.Sys.Print('petsc rank={}, petsc size={}'.format(pRank, pSize))
print('petsc rank={}, petsc size={}'.format(pRank, pSize))

# break into communicators

NumProcsPerSubComm = 2 
color = pRank % NumProcsPerSubComm 
NumSubComms = pSize/NumProcsPerSubComm 
#subcomm = PETSc.Comm(MPI.COMM_WORLD.Split(color))
#subcomm = MPI.COMM_WORLD.Split(color)
subcomm = pComm.tompi4py().Split(color)
subRank = subcomm.Get_rank()
subSize = subcomm.Get_size()

PETSc.Sys.Print('number of subcomms = {}'.format(NumSubComms))
for i in range(pSize):
    pComm.barrier()
    #PETSc.Sys.Print('sub rank {}, psub size{}'.format(pRank, pSize))
    if (pComm.rank == i):
        print('sub rank {}/{}, color:{}'.format(subRank, subSize, color))


n =100
a = np.random.randn(n, n)
asym = a + a.T

# if subcomm != MPI.COMM_NULL:
if color == 1:
    print('creating A in subcomm {}= {}, {}'.format(color, subSize, subRank))
    A = PETSc.Mat().createDense([n,n], comm=subcomm)
    A.setUp()
    row1, row2 = A.getOwnershipRange()
    PETSc.Sys.Print('rows={},{}'.format(row1, row2))
    for ip in range(row1, row2):
        for k in range(n):
            A.setValues(ip, k, asym[ip, k])

    A.assemble()        

    ksp = PETSc.KSP()
    ksp.create(subcomm)

    ksp.setType('cg')
    ksp.getPC().setType('none')

    x, b = A.getVecs()
    x.set(0)
    b.set(1)

    ksp.setOperators(A)
    ksp.setFromOptions()
    ksp.view()
    ksp.solve(b, x)
    PETSc.Sys.Print('||x|| = {}'.format(x.norm()))