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)