import sys
import petsc4py

petsc4py.init(sys.argv)

from petsc4py import PETSc


class BouncingBall:
    n = 2
    comm = PETSc.COMM_SELF

    def __init__(self):
        pass

    def initialCondition(self, u):
        u[0] = 0.0
        u[1] = 20.0
        u.assemble()

    def evalRHSFunction(self, ts, t, u, f):
        f[0] = u[1]
        f[1] = -9.8
        f.assemble()

    def evalRHSJacobian(self, ts, t, u, A, B):
        J = A
        J[0, 0] = 0.0
        J[1, 0] = 0.0
        J[0, 1] = 1.0
        J[1, 1] = 0.0
        J.assemble()
        if A != B:
            B.assemble()


class Monitor:
    def __init__(self):
        pass

    def __call__(self, ts, k, t, x):
        PETSc.Sys.Print(f'Position at time {t}: {x[0]}')


ode = BouncingBall()

J = PETSc.Mat().create()
J.setSizes([ode.n, ode.n])
J.setType('aij')
J.setUp()
J.assemble()

u = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
f = u.duplicate()

ts = PETSc.TS().create(comm=ode.comm)
ts.setProblemType(ts.ProblemType.NONLINEAR)

ts.setType(ts.Type.BEULER)
ts.setRHSFunction(ode.evalRHSFunction, f)
ts.setRHSJacobian(ode.evalRHSJacobian, J)

# ts.setSaveTrajectory()
ts.setTime(0.0)
ts.setTimeStep(0.01)
ts.setMaxTime(15.0)
# ts.setMaxSteps(1000)
ts.setExactFinalTime(PETSc.TS.ExactFinalTime.MATCHSTEP)
# ts.setMonitor(Monitor())

direction = [-1]
terminate = [False]


def event(ts, t, X, fvalue):
    fvalue[0] = X[0]


def postevent(ts, events, t, X, forward):
    X[0] = 0.0
    X[1] = -0.9 * X[1]
    X.assemble()


ts.setEventHandler(direction, terminate, event, postevent)
ts.setEventTolerances(1e-6, vtol=[1e-9])

ts.setFromOptions()

ode.initialCondition(u)
ts.solve(u)