from __future__ import division
import numpy as np


def __compute_k__(fun, ti, vi, h, *args):
    """Computes the k values for the rk44 method.

    :fun: ODE function, must be in the form of f(t, v), return v
    :ti: time for previous iteration
    :vi: previous time step values (in iterable form)
    :h: time step
    :returns:
        :k: array of k values (two-dimensional numpy array)

    """
    # convert to numpy array
    vi = np.array(vi)

    # initialize k
    k = []

    # compute values for k from fun = f(ti, vi)
    k.append(h*np.array(fun(ti, vi, *args)))
    k.append(h*np.array(fun(ti+h/2, vi+k[-1]/2, *args)))
    k.append(h*np.array(fun(ti+h/2, vi+k[-1]/2, *args)))
    k.append(h*np.array(fun(ti+h, vi+k[-1], *args)))

    return np.array(k).T


def step(fun, ti, vi, h, *args):
    """Performs a single step of rk44 integration for second order ode.

    :fun: ODE function, must be in the form of f(t, v)
    :ti: time at start of time step
    :vi: values at start of time step (in iterable form)
    :h: time step
    :returns:
        :t: time at end of time step
        :v: values at end of time step (as list)

    """
    k = __compute_k__(fun, ti, vi, h, *args)

    # weights of k's for rk44
    wts = [1, 2, 2, 1]

    # compute velocity, use numpy's dot product to compute weighted sum
    v = vi + k.dot(wts)/6

    # increment time step
    t = ti + h

    return t, list(v)


def rk44(fun, ti, vi, tf, h, *args):
    """Performs Runge-Kutta integration for second order ODE.

    :fun: ODE function, must be in the form of f(t, r, v)
    :ti: initial time
    :vi: initial conditions (in iterable form)
    :tf: final evaluation time
    :h: time step
    :returns:
        :t: final time value
        :v: final values at tf (as list)

    """

    # initialization of solution variables
    t = ti
    v = list(vi)

    # step until time reaches final time
    while t < tf:

        # allow for h to change to ensure time stops at tf (if necessary)
        hstep = min(h, tf-t)

        t, v = step(fun, t, v, hstep, *args)

    return t, v
