from petsc4py.PETSc import KSP, PC, IS, Mat, Options, Viewer, Vec, LGMap
import numpy as np
from mpi4py import MPI
comm_world=MPI.COMM_WORLD

viewer = Viewer().createBinary('matrix.bin', 'r')
matrix = Mat().load(viewer)
viewer = Viewer().createBinary('IS_u.bin', 'r')
u = IS().load(viewer)
viewer = Viewer().createBinary('IS_p.bin', 'r')
p = IS().load(viewer)
viewer = Viewer().createBinary('IS_t.bin', 'r')
t = IS().load(viewer)
viewer = Viewer().createBinary('IS_up.bin', 'r')
up = IS().load(viewer)

comm = matrix.getComm()
rank = comm.rank

local = up.getLocalSize()

Kup = matrix.createSubMatrix(up, up)
(Istart, Iend) = Kup.getOwnershipRange()
u_up_local = u.embed(up, True)
u_up_local_idx = list(u_up_local.getIndices())
def incr(lst, i):
    return [x+i for x in lst]
u_up_local_idx = incr(u_up_local_idx, Istart)
u_up = IS().createGeneral(u_up_local_idx)
p_up_local = p.embed(up, True)
p_up_local_idx = list(p_up_local.getIndices())
p_up_local_idx = incr(p_up_local_idx, Istart)
p_up = IS().createGeneral(p_up_local_idx)
Ku = Kup.createSubMatrix(u_up, u_up)
Ku.view()
Kp = Kup.createSubMatrix(p_up, p_up)
Kp.view()

"""
# In order to build debug test
NEQ=4*4
K = Mat().create()
K.setSizes([NEQ,NEQ])
K.setUp()
K.setPreallocationNNZ(1)
for row in range(NEQ):
    K.setValue(row, row, row+1.)

K.assemble()
(Istart, Iend) = K.getOwnershipRange()
u = IS().createGeneral(sorted(list(range(Istart, Iend, 4)) + list(range(Istart + 1, Iend, 4))))
p = IS().createGeneral(list(range(Istart + 2, Iend, 4)))
t = IS().createGeneral(list(range(Istart + 3, Iend, 4)))
up = IS().createGeneral(sorted(list(range(Istart, Iend, 4)) + list(range(Istart + 1, Iend, 4)) + list(range(Istart + 2, Iend, 4))))
# u.view()
# p.view()
# t.view()
# up.view()

import os.path as osp
wdir='/tmp/'
v = Viewer().createBinary(osp.join(wdir,"matrix.bin"), "w")
v.view(K)
v = Viewer().createBinary(osp.join(wdir,"IS_u.bin"), "w")
v.view(u)
v = Viewer().createBinary(osp.join(wdir,"IS_p.bin"), "w")
v.view(p)
v = Viewer().createBinary(osp.join(wdir,"IS_t.bin"), "w")
v.view(t)
v = Viewer().createBinary(osp.join(wdir,"IS_up.bin"), "w")
v.view(up)
"""