# baan.py - banen met en zonder luchtweerstand, met de methode van Heun (hoofdstuk 4)
import math
import numpy as np

g = 9.81


def baan(v0, hoek, h0=0.0, vt=None, wind=0.0, dt=0.001):
    """Worp vanaf hoogte h0 (m) met snelheid v0 (m/s) onder een hoek (graden) met de horizontaal.
    vt = eindsnelheid (m/s) voor kwadratische luchtweerstand, None = zonder lucht; wind in m/s (+ = mee).
    Geeft een array met rijen [t, x, y, vx, vy], tot het voorwerp de grond (y = 0) raakt."""
    def f(s):                                       # s = [x, y, vx, vy]  ->  ds/dt
        x, y, vx, vy = s
        if vt is None:
            return np.array([vx, vy, 0.0, -g])
        rx, ry = vx - wind, vy                      # snelheid tegenover de lucht
        k = g / vt**2 * math.hypot(rx, ry)          # |F_w| / m = g (v_rel / vt)^2
        return np.array([vx, vy, -k * rx, -g - k * ry])

    a = math.radians(hoek)
    s = np.array([0.0, h0, v0 * math.cos(a), v0 * math.sin(a)])
    t, rijen = 0.0, [[0.0, *s]]
    while True:
        k1 = f(s)
        k2 = f(s + dt * k1)
        s2 = s + dt / 2 * (k1 + k2)                 # Heun: gemiddelde van twee hellingen
        if s2[1] < 0:                               # de grond ligt binnen deze stap: lineair interpoleren
            q = s[1] / (s[1] - s2[1])
            s2 = s + q * (s2 - s)
            rijen.append([t + q * dt, *s2])
            return np.array(rijen)
        s, t = s2, t + dt
        rijen.append([t, *s])


if __name__ == "__main__":
    print("pakket uit de drone: 40 m hoog, 8,0 m/s horizontaal")
    for naam, vt in (("zonder lucht", None), ("met lucht (vt = 18,6 m/s)", 18.6)):
        t, x, y, vx, vy = baan(8.0, 0.0, 40.0, vt)[-1]
        hoek = math.degrees(math.atan2(-vy, vx))
        print(f"  {naam:26s} {t:5.2f} s  {x:5.1f} m verder  {math.hypot(vx, vy):5.1f} m/s onder {hoek:4.1f}°")

    print("sproeier: 19,0 m/s onder 25°, mondje 0,30 m hoog")
    for naam, vt in (("zonder lucht", None), ("druppel 1 mm", 4.03), ("druppel 2 mm", 6.49),
                     ("druppel 3 mm", 8.06), ("druppel 4 mm", 8.83)):
        b = baan(19.0, 25.0, 0.30, vt)
        print(f"  {naam:13s} bereik {b[-1, 1]:5.2f} m  top {b[:, 2].max():4.2f} m  vliegtijd {b[-1, 0]:4.2f} s")

    b1, b2 = baan(19.0, 25.0, 0.30, 8.06, dt=0.001), baan(19.0, 25.0, 0.30, 8.06, dt=0.002)
    print(f"controle 3 mm: dt = 1 ms {b1[-1, 1]:.4f} m, dt = 2 ms {b2[-1, 1]:.4f} m")