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

# --------------------------------------------------------------------

from petsc4py import PETSc

# --------------------------------------------------------------------

class MyODE:
    """
    du/dt + u**2 = 0;
    u0[i] = 10 + i
    """
    def __init__(self):
        self.call_log = {}
        
    def _log(self, method, ts, t, u):
        self.call_log.setdefault(method, 0)
        self.call_log[method] += 1
        n  = ts.getStepNumber()
        t0 = ts.getTime()
        u0 = ts.getSolution()
        du = (u-u0).norm()
        dt = ts.getTimeStep()
        meth = '%-20s' % ('MyODE.%s()'% method)
        info = 'n->%2d | t0->%4.2f | t->%4.2f | du->%4.2f | dt->%5.3f | du/dt->%4.2f' % (n,t0,t,du,dt,du/dt)
        print '%s %s' % (meth, info)

    def function(self,ts,t,u,F):
        self._log('function', ts, t, u)
        dt = ts.getTimeStep()
        u0 = ts.getSolution()
        f = (u - u0)/dt + u * u
        f.copy(F)

    def jacobian(self,ts,t,u,J,P):
        self._log('jacobian', ts, t, u)
        u0 = ts.getSolution()
        dt = ts.getTimeStep()
        P.zeroEntries()
        diag = 1/dt + 2 * u
        P.diagonalSet(diag)
        P.assemble()
        if J != P: J.assemble()
        return False # same_nz

    def setup(self, ts, t, u):
        self._log('setup', ts, t, u)

    def presolve(self, ts, t, u):
        self._log('presolve', ts, t, u)
        
    def postsolve(self, ts, t, u):
        self._log('postsolve', ts, t, u)
        
    def prestep(self, ts, t, u):
        self._log('prestep', ts, t, u)
        
    def poststep(self, ts, t, u):
        self._log('poststep', ts, t, u)
        
    def start(self, ts, t, b, u):
        self._log('start', ts, t, u)

    def step(self, ts, t, b, u):
        self._log('step', ts, t, u)
        snes = ts.getSNES()
        snes.solve(b, u)
        
    def verify(self, ts, t, b, u):
        self._log('verify', ts, t, u)
        #return (True, 1.1*ts.time_step)
        return (True, 0.9*ts.time_step)
        #return (True, float('nan'))
        n  = ts.getStepNumber()
        t0 = ts.getTime()
        u0 = ts.getSolution()
        du = (u-u0).norm()
        dt = ts.getTimeStep()
        return
        if (du) > 0.5:
            print du, 0.5*ts.time_step
            return (False, 0.5*ts.time_step)
        if (du) < 0.3:
            print du, 2.0*ts.time_step
            return (True, 2.0*ts.time_step)
        
ts = PETSc.TS().create(PETSc.COMM_SELF)
ts.setType('user')
        
ode = MyODE()
J = PETSc.Mat().create(ts.comm)
J.setSizes(3);
J.setFromOptions()
J.setPreallocation(1)
for i in range(3):
    J[i,i] = 1
J.assemble()
J.zeroEntries()
u, f = J.getVecs()


T0, dT, nT = 0.00, 0.1, 5
T = T0 + nT*dT
ts.setTime(T0)
ts.setTimeStep(dT)
ts.setMaxTime(T)
ts.setAppCtx(ode)
ts.setFunction(ode.function, f)
ts.setJacobian(ode.jacobian, J, J)

ts.setUserFunctions(
    setup     = ode.setup,
    presolve  = ode.presolve,
    postsolve = ode.postsolve,
    prestep   = ode.prestep, 
    poststep  = ode.poststep,
    start     = ode.start,
    step      = ode.step,
    verify    = ode.verify,
    )

def monitor(ts, n, t, x):
    t0 = ts.getTime()
    dt = ts.getTimeStep()
    print 'timestep %d dt %f time %f --- t0:%f t-t0:%f  dt ok:%s' \
          % (n, dt, t, t0, t-t0, (t-t0)==dt)
#ts.setMonitor(monitor)
    

#ts.snes.ksp.pc.setType('none')
#ts.setUseFDColoring(True)
#ts.setSolution(u)
#ts.setUp()
#ts.snes.setUseMF(True)

ts.setFromOptions()

for i in range(u.size):
    u[i] = 10 + (i+1)

try:
    ts.solve(u)
except:
    PETSc.Error.view()
    J.view()
    raise

ts.setTime(T0)
ts.setTimeStep(dT)
ts.setMaxTime(T)
for i in range(u.size):
    u[i] = 10 + i
try:
    ts.solve(u)
except:
    PETSc.Error.view()
    J.view()
    raise



print '\n'
print '-----------'
call_log = ode.call_log.copy()
for k in ('function', 'jacobian'):
    print k, call_log.pop(k, 0)
print '-----------'
for k,v in call_log.items():
    print k, v
print '-----------'