from typing import Callable

import numpy as np
from numpy.typing import NDArray


def rk4(ode: Callable, t0: float, v0: NDArray, weite: float, schritte: int, fein: bool = False) -> tuple[list, list]:
    """
    Standard Runge-Kutta Verfahren 4. Ordnung
    :param ode: ODE-System als Funktion
    :param t0: Startwert der unabhängigen Variable
    :param v0: Startwerte
    :param weite: Integrationsweite
    :param schritte: Schrittzahl
    :param fein: Fein-Rechnung?
    :return: Variable und Funktionswerte an jedem Stützpunkt
    """
    h = weite/schritte

    t_list = [t0]
    werte = [v0]

    for _ in range(schritte):
        t = t_list[-1]
        v = werte[-1]

        if not fein:
            v_next = rk4_step(ode, t, v, h)

        else:
            v_grob = rk4_step(ode, t, v, h)
            v_half = rk4_step(ode, t, v, 0.5 * h)
            v_fein = rk4_step(ode, t + 0.5 * h, v_half, 0.5 * h)
            v_next = v_fein + (v_fein - v_grob) / 15.0

        t_list.append(t + h)
        werte.append(v_next)

    return t_list, werte

def rk4_step(ode: Callable, t: float, v: NDArray, h: float) -> NDArray:
    """
    Ein Schritt des Runge-Kutta Verfahrens 4. Ordnung
    :param ode: ODE-System als Funktion
    :param t: unabhängige Variable
    :param v: abhängige Variablen
    :param h: Schrittweite
    :return: abhängige Variablen nach einem Schritt
    """
    k1 = ode(t, v)
    k2 = ode(t + 0.5 * h, v + 0.5 * h * k1)
    k3 = ode(t + 0.5 * h, v + 0.5 * h * k2)
    k4 = ode(t + h, v + h * k3)
    return v + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)

def rk4_end(ode: Callable, t0: float, v0: NDArray, weite: float, schritte: int, fein: bool = False):
    """
    Standard Runge-Kutta Verfahren 4. Ordnung, nur Ausgabe der letzten Variablenwerte
    :param ode: ODE-System als Funktion
    :param t0: Startwert der unabhängigen Variable
    :param v0: Startwerte
    :param weite: Integrationsweite
    :param schritte: Schrittzahl
    :param fein: Fein-Rechnung?
    :return: Variable und Funktionswerte am letzten Stützpunkt
    """
    h = weite / schritte
    t = float(t0)
    v = np.array(v0, dtype=float, copy=True)

    for _ in range(schritte):
        if not fein:
            v_next = rk4_step(ode, t, v, h)
        else:
            v_grob = rk4_step(ode, t, v, h)
            v_half = rk4_step(ode, t, v, 0.5 * h)
            v_fein = rk4_step(ode, t + 0.5 * h, v_half, 0.5 * h)
            v_next = v_fein + (v_fein - v_grob) / 15.0

        t += h
        v = v_next

    return t, v

# RK4 mit Simpson bzw. Trapez
def rk4_integral(ode: Callable, t0: float, v0: NDArray, weite: float, schritte: int, integrand_at: Callable, fein: bool = False, simpson: bool = True):
    """
    Runge-Kutta Verfahren 4. Ordnung mit Simpson bzw. Trapez
    :param ode: ODE-System als Funktion
    :param t0: Startwert der unabhängigen Variable
    :param v0: Startwerte
    :param weite: Integrationsweite
    :param integrand_at: Funktion
    :param schritte: Schrittzahl
    :param fein: Fein-Rechnung?
    :param simpson: Simpson? Wenn nein, dann Trapez
    :return: Variable und Funktionswerte am letzten Stützpunkt
    """
    h = weite / schritte
    habs = abs(h)

    t = float(t0)
    v = np.array(v0, dtype=float, copy=True)

    if simpson and (schritte % 2 == 0):
        f0 = float(integrand_at(t, v))
        odd_sum = 0.0
        even_sum = 0.0
        fN = None

        for i in range(1, schritte + 1):
            if not fein:
                v_next = rk4_step(ode, t, v, h)
            else:
                v_grob = rk4_step(ode, t, v, h)
                v_half = rk4_step(ode, t, v, 0.5 * h)
                v_fein = rk4_step(ode, t + 0.5 * h, v_half, 0.5 * h)
                v_next = v_fein + (v_fein - v_grob) / 15.0

            t += h
            v = v_next

            fi = float(integrand_at(t, v))
            if i == schritte:
                fN = fi
            elif i % 2 == 1:
                odd_sum += fi
            else:
                even_sum += fi

        S = f0 + fN + 4.0 * odd_sum + 2.0 * even_sum
        s = (habs / 3.0) * S
        return t, v, s

    f_prev = float(integrand_at(t, v))
    acc = 0.0

    for _ in range(schritte):
        if not fein:
            v_next = rk4_step(ode, t, v, h)
        else:
            v_grob = rk4_step(ode, t, v, h)
            v_half = rk4_step(ode, t, v, 0.5 * h)
            v_fein = rk4_step(ode, t + 0.5 * h, v_half, 0.5 * h)
            v_next = v_fein + (v_fein - v_grob) / 15.0

        t += h
        v = v_next

        f_cur = float(integrand_at(t, v))
        acc += 0.5 * (f_prev + f_cur)
        f_prev = f_cur

    s = habs * acc
    return t, v, s