"""System of ODEs - test function"""
import numpy as np
import matplotlib.pyplot as plt

#Superclass with every particular method as a subclass
class ODESolver:
    """
    Superclass for numerical methods solving scalar and vector ODEs
      du/dt = f(u, t)
    Attributes:
    t: array of time values
    u: array of solution values (at time points t)
    k: step number of the most recently computed solution
    f: callable object implementing f(u, t)
    """
    def __init__(self, f):
        if not callable(f):
            raise TypeError('f is %s, not a function' % type(f))
        # For ODE systems, f will often return a list, but
        # arithmetic operations with f in numerical methods
        # require that f is an array. Let self.f be a function
        # that first calls f(u,t) and then ensures that the
        # result is an array of floats.
        self.f = lambda u, t: np.asarray(f(u, t), float)

    def advance(self):
        """Advance solution one time step."""
        raise NotImplementedError

    def set_initial_condition(self, U0):
        if isinstance(U0, (float,int)):  # scalar ODE
            self.neq = 1
            U0 = float(U0)
        else:                            # system of ODEs
            U0 = np.asarray(U0)          # (assume U0 is sequence)
            self.neq = U0.size
        self.U0 = U0

        # Check that f returns correct length:
        try:
            f0 = self.f(self.U0, 0)
        except IndexError:
            raise IndexError('Index of u out of bounds in f(u,t) func. Legal indices are %s' % (str(range(self.neq))))
        if f0.size != self.neq:
            raise ValueError('f(u,t) returns %d components, while u has %d components' % (f0.size, self.neq))

    def solve(self, time_points, terminate=None):
        """
        Compute solution u for t values in the list/array
        time_points, as long as terminate(u,t,step_no) is False.
        terminate(u,t,step_no) is a user-given function
        returning True or False. By default, a terminate
        function which always returns False is used.
        """
        if terminate is None:
            terminate = lambda u, t, step_no: False

        if isinstance(time_points, (float,int)):
            raise TypeError('solve: time_points is not a sequence')
        self.t = np.asarray(time_points)
        if self.t.size <= 1:
            raise ValueError('ODESolver.solve requires time_points array with at least 2 time points')

        n = self.t.size
        if self.neq == 1:  # scalar ODEs
            self.u = np.zeros(n)
        else:              # systems of ODEs
            self.u = np.zeros((n,self.neq))

        # Assume that self.t[0] corresponds to self.U0
        self.u[0] = self.U0

        # Time loop
        for k in range(n-1):
            self.k = k
            self.u[k+1] = self.advance()
            if terminate(self.u, self.t, self.k+1):
                break  # terminate loop over k
        return self.u[:k+2], self.t[:k+2]


class ForwardEuler(ODESolver):
    def advance(self):
        u, f, k, t = self.u, self.f, self.k, self.t
        dt = t[k+1] - t[k]
        u_new = u[k] + dt*f(u[k], t[k])
        return u_new

class RungeKutta4(ODESolver):
    def advance(self):
        u, f, k, t = self.u, self.f, self.k, self.t
        dt = t[k+1] - t[k]
        dt2 = dt/2.0
        K1 = dt*f(u[k], t[k])
        K2 = dt*f(u[k] + 0.5*K1, t[k] + dt2)
        K3 = dt*f(u[k] + 0.5*K2, t[k] + dt2)
        K4 = dt*f(u[k] + K3, t[k] + dt)
        u_new = u[k] + (1/6.0)*(K1 + 2*K2 + 2*K3 + K4)
        return u_new
    
class Heun(ODESolver):
    def advance(self):
        u, f, k, t = self.u, self.f, self.k, self.t
        dt = t[k+1] - t[k]
        fval=f(u[k], t[k])
        u_star=u[k] + dt*fval
        u_new=u[k] + 0.5*dt*fval + 0.5*dt*f(u_star, t[k+1])
        return u_new
    
class Leapfrog(ODESolver):
    def advance(self):
        u, f, k, t = self.u, self.f, self.k, self.t
        if k==0: #need to use a different method, e.g., ForwardEuler
            dt = t[k+1] - t[k]
            u_new = u[k] + dt*f(u[k], t[k])            
        else:    
            dt_double = t[k+1] - t[k-1]
            u_new=u[k-1]+dt_double*f(u[k], t[k])        
        return u_new


def test_Leapfrog():
    """Test for system of ODEs with known solutions"""
    def f(u,t):
        return [3+(3+4*t-u[1])**3,4+(2+3*t-u[0])**4]
    
    def v_exact(t):
        return 2+3*t
    
    def w_exact(t):
        return 3+4*t
    
    lf=Leapfrog(f)
    lf.set_initial_condition([2,3])
    tpoints=np.linspace(0,10,201)
    u, times=lf.solve(tpoints)
    vok=v_exact(times)
    wok=w_exact(times)
    assert np.allclose(vok, u[:,0]), "Numerical solution for v not close enough to exact value"
    assert np.allclose(wok, u[:,1]), "Numerical solution for w not close enough to exact value"

    

if __name__ == '__main__':
    test_Leapfrog()  
    import glob, os
    #Study convergence of several methods
    def f(u,t):
        return u
    t0=0
    tn=8
    u0=1
    nlist=[2**k for k in range(1,14)]
    i=0
    for n in nlist:
        i+=1
        tpoints=np.linspace(t0,tn,n)
        plt.figure()
        for method in [ForwardEuler, RungeKutta4, Heun, Leapfrog]:
            met=method(f)
            met.set_initial_condition(u0)
            um, tm=met.solve(tpoints)
            plt.plot(tm, um, label=met.__class__.__name__)
        plt.legend(loc=2)
        plt.axis([0,8,0,3000])
        plt.xlabel("u")
        plt.ylabel("t")
        plt.title("u'=u, n=%g"%n)
        plt.savefig("ejE27c%02d-n=%g.png"%(i,n))
        plt.close()
    os.system('convert -delay 100 ejE27c*.png ejE27c.gif')
    
    
    ##Part d
    def f(u,t):
        return -u
    tn=5
    for n in nlist:
        i-=1 #to get inverse order in .gif
        tpoints=np.linspace(t0,tn,n)
        plt.figure()
        for method in [ForwardEuler, RungeKutta4, Heun, Leapfrog]:
            met=method(f)
            met.set_initial_condition(u0)
            um, tm=met.solve(tpoints)
            plt.plot(tm, um, label=met.__class__.__name__)
        plt.legend(loc=1)
        plt.axis([0,5,0,1])
        plt.xlabel("u")
        plt.ylabel("t")
        plt.title("u'=-u, n=%g"%n)
        plt.savefig("ejE27d%02d-n=%g.png"%(i,n))
        plt.close()
    os.system('convert -delay 100 ejE27d*.png ejE27d.gif')
    

            
            
