from util FuncVec2 import Core algorithm def _runge_kutta_ FuncVec2 ve

from util.FuncVec2 import *


# ############## #
# Core algorithm #
# ############## #


def _runge_kutta_(t, f: FuncVec2, vec: Vec2, h: float):
    r1 = f.apply(t,           vec)
    r2 = f.apply(t + h / 2.0, vec + r1 * h / 2.0)
    r3 = f.apply(t + h / 2.0, vec + r2 * h / 2.0)
    r4 = f.apply(t + h,       vec + r3 * h)

    r_result = (r1 + r2 * 2.0 + r3 * 2.0 + r4) * h / 6.0

    return r_result.x, r_result.y


def runge_kutta(start: Vec2, func: FuncVec2, h: float):
    t   = 0
    xs  = [0]
    ys  = [0]
    vec = start

    while True:
        t += h

        x = xs[len(xs) - 1] + vec.x
        y = ys[len(ys) - 1] + vec.y

        xs.append(x)
        ys.append(y)

        vx, vy = _runge_kutta_(t, func, vec, h)
        vec    = Vec2(vx, vy)

        if y <= 0:
            break

    return xs, ys