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')
from tools.postprocessK.kperm.flow import getKeff




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

	ref=1
	rank=0
	pn=1
	t0=time.time()
	
	pcomm=PETSc.COMM_SELF
	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



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

	R = PETSc.Vec().createSeq((n,None),comm=pcomm) #PETSc.COMM_WORLD
	R.setUp()

	k2, Nz, nnz2=getKref(k,1,2,ref)
	k, Nz, nnz=getKref(k,0,2,ref)


	pbc=float(Nz)


	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()



	ksp = PETSc.KSP()
	ksp.create(comm=pcomm)
	ksp.setFromOptions()
	P = R.copy()
	ksp.setType(PETSc.KSP.Type.CG)
	pc = PETSc.PC()

	pc.create(comm=pcomm)	
	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)
		print(keff,ref,nx,ny,nz)
		return keff


	return



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