from petsc4py import PETSc

import scipy.sparse as sps
import scipy.sparse.linalg as spslin

import numpy as np

import time

n = 4000


matrix = sps.random(n, n, format='csr') + 10 * sps.eye(n) + \
        0.1j * sps.random(n, n, format='csr')

exact_solution = 0.2 * np.arange(n) + 0.01j * np.arange(n)

rhs = matrix * exact_solution

print("")
print(" Start PETSC so hope to be parallel")
t0 = time.time()
n_dofs = rhs.shape[0]

# Create the PETSc environment & fill the sparse matrix
pA = PETSc.Mat().createAIJ(size=matrix.shape, csr=(matrix.indptr,
						   matrix.indices,
						   matrix.data))
pA.assemblyBegin()
pA.assemblyEnd()

# create linear solver
ksp = PETSc.KSP()
ksp.create(PETSc.COMM_WORLD)

# use direct method
ksp.setType('preonly')
ksp.getPC().setType('lu')
ksp.getPC().setFactorSolverType('mumps')

x, b = pA.getVecs()
b.setValues(range(n_dofs), rhs)

# and next solve
ksp.setOperators(pA)
ksp.setFromOptions()
ksp.solve(b, x)

x_sol = x.getArray()
print(" DONE for", time.time()-t0)

print("")
#print(" The error is", np.linalg.norm(res_scipy - res_petsc))