# -*- coding: utf-8 -*-
"""
Created on Mon Feb  8 08:06:19 2016

@author: nemenman
"""
import numpy as np


def Euler(xPrime, t0=0.0, x0=0.0, T=1.0, dt=0.1):
    """
    Solves 1-d ODE using the Euler method.
    Euler(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1):
    """
    t = np.arange(t0, T+dt, dt)
    x = np.zeros(t.size)
    x[0] = x0
    for i in range(1, t.size):
        x[i] = x[i-1] + dt * xPrime(x[i-1])

    return (t, x)


def RK2(xPrime, t0=0.0, x0=0.0, T=1.0, dt=0.1):
    """
    Solves 1-d ODE using the Runge-Kutta 2 method.

    RK2(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1):
    """

    t = np.arange(t0, T+dt, dt)
    x = np.zeros(t.size)
    x[0] = x0
    for i in range(1, t.size):
        fx = xPrime(x[i-1])
        guess = x[i-1] + dt * fx
        fxdx = xPrime(guess)
        x[i] = x[i-1]+0.5*(fx+fxdx)*dt

    return (t, x)


def EulerArg(xPrime, t0=0.0, x0=0.0, T=1.0, dt=0.1, arg=()):
    """
    Solves 1-d ODE using the Euler method.

    EulerArg(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1),arg=():
    """
    t = np.arange(0, T+dt, dt)
    x = np.zeros(t.size)
    x[0] = x0
    for i in range(1, t.size):
        x[i] = x[i-1] + dt * xPrime(x[i-1], t[i-1], arg)

    return (t, x)


def RK2Arg(xPrime, t0=0.0, x0=0.0, T=1.0, dt=0.1, arg=()):
    """
    Solves 1-d ODE using the Runge-Kutta 2 method.

    RK2Arg(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1,arg=()):
    """

    t = np.arange(0, T+dt, dt)
    x = np.zeros(t.size)
    x[0] = x0
    for i in range(1, t.size):
        fx = xPrime(x[i-1], t[i-1], arg)
        guess = x[i-1] + dt * fx
        fxdx = xPrime(guess, t[i-1], arg)
        x[i] = x[i-1]+0.5*(fx+fxdx)*dt

    return (t, x)