import numpy as np
#import petsc4py
import math
import time
#from mpi4py import MPI
from tools.postprocessK.kperm.computeFlows import *

from petsc4py import PETSc
#petsc4py.init('-ksp_max_it 9999999999',comm=PETSc.COMM_SELF)
from tools.postprocessK.flow import getKeff






def PetscP(datadir,ref,k,saveres):

	#datadir='./data/'+str(job)+'/'


	#comm=MPI.COMM_WORLD
	#rank=comm.Get_rank()
	'''
	size=comm.Get_size()
	print(rank,size)
	pcomm = MPI.COMM_WORLD.Split(color=rank, key=rank)
	#print(new_comm.Get_rank())

	#pcomm=comm.Create(newgroup)
	print('entro')
	print pcomm.Get_rank()

	print pcomm.Get_size()
	pcomm=comm


	rank=pcomm.rank
	pn=pcomm.size
	#PETSc.COMM_WORLD.PetscSubcommCreate(pcomm,PetscSubcomm *psubcomm)
	print(rank,pn)
	'''
	#Optpetsc = PETSc.Options() 
	rank=0
	pn=1
	t0=time.time()
	
	#comm=MPI.Comm.Create()
	if k.shape[2]==1:
		refz=1
	else:
		refz=ref

	nz, ny, nx=k.shape[0]*ref,k.shape[1]*ref,k.shape[2]*refz
	n=nx*ny*nz

	print('algo')

	K = PETSc.Mat().create(comm=PETSc.COMM_SELF)
	print('algo2')
	K.setType('seqaij')
	print('algo3')
	K.setSizes(((n,None),(n,None))) # Aca igual que lo que usas arriba
	K.setPreallocationNNZ(nnz=(7,4)) # Idem anterior


	#K = PETSc.Mat('seqaij', m=n,n=n,nz=7,comm=PETSc.COMM_WORLD)
	#K = PETSc.Mat('aij', ((n,None),(n,None)), nnz=(7,4),comm=PETSc.COMM_WORLD)
	#K = PETSc.Mat().createAIJ(((n,None),(n,None)), nnz=(7,4),comm=PETSc.COMM_WORLD)
	#K = PETSc.Mat().createSeqAIJ(((n,None),(n,None)), nnz=(7,4),comm=PETSc.COMM_WORLD)
	#K.setPreallocationNNZ(nnz=(7,4))
	print('ksetup')
	#K.MatCreateSeqAIJ()
	#K=PETSc.Mat().MatCreate(PETSc.COMM_WORLD)

	#K = PETSc.Mat().createAIJ(((n,None),(n,None)), nnz=(7,4),comm=pcomm)

	K.setUp()
	print('entro2')
	R = PETSc.Vec().createSeq((n,None),comm=PETSc.COMM_SELF) #PETSc.COMM_WORLD
	R.setUp()
	print('entro2')
	k2, Nz, nnz2=getKref(k,1,2,ref)
	k, Nz, nnz=getKref(k,0,2,ref)


	pbc=float(Nz)
	#print('entro3')

	K,R = firstL(K,R,k,pbc)
	r=(k.shape[1]-2)*(k.shape[2]-2)*nnz2  #start row
	K,R =lastL(K,R,k2,r)

	k2=0

	K.assemble()
	R.assemble()


	print('entro3')
	ksp = PETSc.KSP()
	ksp.create(comm=PETSc.COMM_SELF)
	ksp.setFromOptions()
	print('entro4')
	P = R.copy()
	ksp.setType(PETSc.KSP.Type.CG)
	pc = PETSc.PC()

	pc.create(comm=PETSc.COMM_SELF)	
	print('entro4')
	pc.setType(PETSc.PC.Type.JACOBI)
	ksp.setPC(pc)
	ksp.setOperators(K)
	ksp.setUp()
	t1=time.time()
	ksp.solve(R, P)
	t2=time.time()
	p=P.getArray().reshape(nz,ny,nx)



	if rank==0:
		keff,Q=getKeff(p,k[1:-1,1:-1,1:-1],pbc,Nz)

		return keff


	return



#Ver: A posteriori error estimates and adaptive solvers for porous media flows (Martin Vohralik)