import sys, slepc4py
slepc4py.init(sys.argv)
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
from numpy import exp

opts = PETSc.Options()
n = opts.getInt('n', 30)


#Set up the matrix A

A = PETSc.Mat().create()
A.setSizes([n, n])
A.setFromOptions()
A.setUp()

rstart, rend = A.getOwnershipRange()

# first row
if rstart == 0:
    A[0, :2] = [2, -1]
    rstart += 1
# last row
if rend == n:
    A[n-1, -2:] = [-1, 2]
    rend -= 1
# other rows
for i in range(rstart, rend):
    A[i, i-1:i+2] = [-1, 2, -1]

A.assemble()


Print = PETSc.Sys.Print

#Set up the right vector b

b = PETSc.Vec().create()
b.setSizes(n)
b.setFromOptions()
b.setUp()
for i in range(n):
    b[i] = 1

b.assemble()

#Set up the output vector x

x = PETSc.Vec().create()
x.setSizes(n)
x.setFromOptions()
x.setUp()





E = SLEPc.MFN(); E.create()

F =  SLEPc.FN().create()
#fy = SLEPc.FN().setType('exp')


E.setOperator(A)
E.getFN()
F.setType(SLEPc.FN.Type.EXP)
E.setType(SLEPc.MFN.Type.KRYLOV)
#for alpha in range(1,3):
t=3
F.setScale(t, 1.0 )
E.setFromOptions()
E.solve(b, x)
b=x
E.view()

E.destroy()
#print(b)
"""
myviewer = PETSc.Viewer().createASCII(
    "test.txt", format=PETSc.Viewer.Format.ASCII_COMMON,
    comm= PETSc.COMM_WORLD)
b.view(myviewer)
"""
viewer = PETSc.Viewer().createASCII('test1.dat', 'w')
viewer(x)