import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc


# Number of components of the exact solution
N=10

# Create the exact solution vector
expected_soln_vec = 100*np.random.randn(N)



def residue(expected_soln_vec,  x, f):
    ''' Residue equation implementation '''
    print("Call")
    f = (x - expected_soln_vec)**2
    print("Residue:", f)


# this user class is an application
# context for the nonlinear problem
# at hand; it contains some parameters
# and knows how to compute residuals
class SimpleEqn:
    def __init__(self, m, expected_soln_vec, impl='python'):
        self.m = m
        self.expected_soln_vec = expected_soln_vec
        
        if impl == 'python':
            order = 'c'
        elif impl == 'fortran':
            order = 'f'
        else:
            raise ValueError('invalid implementation')
        self.compute = residue
        self.order = order

    def evalFunction(self, snes, X, F):
        m = self.m
        expected_soln_vec = self.expected_soln_vec
        order = self.order
        x = X.getArray(readonly=1).reshape(m, order=order)
        f = F.getArray(readonly=0).reshape(m, order=order)

        self.compute(expected_soln_vec, x, f)


    
# convenience access to
# PETSc options database
OptDB = PETSc.Options()
m = OptDB.getInt('m', N)
impl  = OptDB.getString('impl', 'python')

# create application context
# and PETSc nonlinear solver
appc = SimpleEqn(m, expected_soln_vec, impl)
snes = PETSc.SNES().create()

# register the function in charge of
# computing the nonlinear residual
f = PETSc.Vec().createSeq(m)
snes.setFunction(appc.evalFunction, f)

# configure the nonlinear solver
# to use a matrix-free Jacobian
snes.setUseMF(True)
snes.getKSP().setType('cg')
snes.setFromOptions()


# solve the nonlinear problem
b, x = None, f.duplicate()
x.set(0) # zero initial guess
snes.solve(b, x)