u"""Eksempler p? svingeforl?p. ODESolver er hentet og forenklet fra Langtangens Primer on Scientific Programming With Python"""

import numpy as np

class ODESolver:
    def __init__(self, f):
        self.f = f 

    def advance(self):
        raise NotImplementedError

    def set_initial_condition(self, U0):
        self.U0 = np.asarray(U0)
        self.number_of_equations = self.U0.size

    def solve(self, time_points):
        self.t = np.asarray(time_points)
        n = self.t.size
        self.u = np.zeros((n, self.number_of_equations))
        self.u[0] = self.U0

        for i in range(n-1):
            self.i = i
            self.u[i+1] = self.advance()

        return self.u, self.t


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


class RungeKutta4(ODESolver):
    def advance(self):
        u, f, i, t = self.u, self.f, self.i, self.t
        dt = t[i+1]-t[i]
        dt2 = dt/2.0

        K1 = dt*f(u[i],             t[i]        )
        K2 = dt*f(u[i] + 0.5 * K1,  t[i] + dt2  )
        K3 = dt*f(u[i] + 0.5 * K2,  t[i] + dt2  )
        K4 = dt*f(u[i] +       K3,  t[i] + dt2  )
        unew = u[i] + (1/6.0) * (K1 + 2*K2 + + 2*K3 + K4)

        return unew


class RungeKutta2(ODESolver):
    def advance(self):
        u, f, i, t = self.u, self.f, self.i, self.t
        dt = t[i+1]-t[i]

        K1 = dt*f(u[i], t[i])
        K2 = dt*f(u[i]+0.5*K1, t[i]+0.5*dt)
        unew = u[i] + K2

        return unew


class DrivenPendulum:
    def __init__(self, m, b, k, F, w):
        self.m, self.k, self.b, self.w, self.F = m, k, b, w, F

    def __call__(self, u, t):
        m, k, b, w, F = self.m, self.k, self.b, self.w, self.F
        return np.array(
            [
                u[1], 
                1.0/m * (F*np.cos(w*t) - b*u[1]-k*u[0])
            ])

class DampedPendulum:
    def __init__(self, m, b, k):
        self.m, self.k, self.b = m, k, b

    def __call__(self, u, t):
        m, k, b = self.m, self.k, self.b
        return np.array(
            [
                u[1], 
                1.0/m * (- b*u[1]-k*u[0])
            ])

if __name__ == "__main__":
    import matplotlib.pyplot as plt
    
    # Velg system
    m = 1.0; k = 10.0; b = 0.1; F = 1.0; w = np.sqrt(k/m)*1.05
    #myPendulum = DrivenPendulum(m, b, k, F, w)
    myPendulum = DampedPendulum(m=3.0, b=0.2, k=2.0)
    
    # Velg l?ser
    mySolver = RungeKutta4(myPendulum) 
    mySolver.set_initial_condition([5,0])
    time_points = np.linspace(0, 200, 20000)
    u, t = mySolver.solve(time_points)

    plt.plot(t, u[:,0])
    plt.xlabel('t')
    plt.ylabel('x')
    plt.show()