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)