####################################################################
#
# A class to solve first-order differential equations of the form
#
#   dx_i(t)/dt = f_i(x(t),t)
#
# where x and f can be vectors with components labelled with i.
#
# In general, a solution to such a problem is provided by
# discretisation in time, i.e. by labelling timesteps in t such
# that t(j+1) = t(j)+delta t, and simultaneously, by replacing
# x_i(t) with x_i(j), with j labelling the time steps.  Then
# x_i(j+1) = x_i(j)+delta x_i, and various methods differ in how
# they determine the delta x_i.
#
# Different methods will be supplemented by the students through 
# the homework exercises, in particular a simple Euler method, 
# Runge-Kutta algorithms of order 2 and 4, and the Euler-Cromer 
# method for oscillatory problems.  In the latter, a second-order 
# differential equation is translated into two first-order 
# differential equations.
#
# In the initialisation of this class, through 
#
#   __init__(self,kernel,stepper),
#
# the function f needs to be provided ("kernel"), and the 
# algorithm for the solution must be picked - the choices above 
# are selected through the names "Euler", "RungeKutta2", 
# "RungeKutta4", etc.
# Note that in the implementation here, the kernel must provide 
# a method 
#  
#   dx_dt(x,t), 
#
# which depends on an array x and a number t.
#
# In order to solve the differential equation(s), the method
#  
#  solve(x0,t0,t1,delta_t)
#
# is invoked.  Here the arguments denote the boundary conditions
# x0 = x(t0), keeping in mind that the x (and therefore the x0)
# are vectors of length Nx, which are presented as numpy arrays.  
# The x_i(t) are evolved from the starting time t0 to the end time 
# t1 in steps of the size delta_t.  The result of this evolution 
# is stored in a list 'history'.
#
# In the evolution above, the boundary conditions etc. are used
# to initialise the problem, in 
#
#  initialise(self,x0,t0,t1,delta_t)
#
# Each new step is added to the history through the method
#
#  record(self,t,x)
#
# The students will, in their homeworks, supplement the class 
# DEq_Solver by providing expressions for the individual steps
# yielding the delta x_i.
#
# These expressions will be coded in the methods EulerStep,
# RK2Step, RK4Step, etc.

class DEq_Solver:
    def __init__(self,kernel,stepper):
        self.kernel  = kernel
        if stepper == 'Euler':
            self.stepper = self.EulerStep
        elif stepper == 'RK2':
            self.stepper = self.RK2Step
        elif stepper == 'RK4':
            self.stepper = self.RK4Step
        elif stepper == 'EulerCromer':
            self.stepper = self.EulerCromerStep
        else:
            print 'ERROR: Unknown stepper',stepper
            exit(1)

    def solve(self,x0,t0,t1,delta_t,positive_index=None,circular_index=None):
        self.initialise(x0,t0,t1,delta_t)
        while self.t <= self.t1:
            self.record(self.t,self.x)
            if positive_index is not None and self.x[positive_index] < 0:
                break
            self.stepper()
        return self.history

    def initialise(self,x0,t0,t1,delta_t):
        self.x       = x0.copy()
        self.t       = t0
        self.t1      = t1
        self.delta_t = delta_t
        self.history = []

    def record(self,t,x):
        self.history.append(  ( t, x.copy() )  )

    # Different stepper choices, can be selected through the stepper
    # argument in __init__

    def EulerStep(self):
        dx_dt = self.kernel.dx_dt(self.x, self.t)
        self.x += dx_dt * self.delta_t
        self.t += self.delta_t

    def RK2Step(self):
        dx_dt = self.kernel.dx_dt(self.x, self.t)
        x1 = self.x + dx_dt * self.delta_t / 2.
        t1 = self.t + self.delta_t / 2.
        dx_dt_1 = self.kernel.dx_dt(x1, t1)
        self.x += dx_dt_1 * self.delta_t
        self.t += self.delta_t

    def RK4Step(self):
        x1 = self.x
        t1 = self.t
        f_1 = self.kernel.dx_dt(x1, t1)

        x2 = self.x + f_1 * self.delta_t / 2.
        t2 = self.t + self.delta_t / 2.
        f_2 = self.kernel.dx_dt(x2, t2)

        x3 = self.x + f_2 * self.delta_t / 2.
        t3 = self.t + self.delta_t / 2.
        f_3 = self.kernel.dx_dt(x3, t3)

        x4 = self.x + f_3 * self.delta_t
        t4 = self.t + self.delta_t
        f_4 = self.kernel.dx_dt(x4, t4)

        delta_x = ( f_1 + 2 * f_2 + 2 * f_3 + f_4 ) / 6.
        self.x += delta_x * self.delta_t
        self.t += self.delta_t
        

    def EulerCromerStep(self):
        dx_dt = self.kernel.dx_dt_EC(self.x, self.t, self.delta_t)
        self.x += dx_dt * self.delta_t
        self.t += self.delta_t