
"""PKTron Algorithms Science — ODE Solvers."""
import numpy as np
_R = np.float32

def euler_method(f, y0, t_span, n_steps=1000):
    """
    Explicit Euler method: y'=f(t,y).
    Returns (t_arr, y_arr) float32.
    """
    t0,tf  = float(t_span[0]), float(t_span[1])
    dt     = _R((tf-t0)/n_steps)
    y      = np.asarray(y0, dtype=_R).copy()
    dim    = y.shape[0] if y.ndim>0 else 1
    t_arr  = np.zeros(n_steps+1, dtype=_R)
    y_arr  = np.zeros((n_steps+1,)+y.shape, dtype=_R)
    t      = _R(t0)
    t_arr[0] = t; y_arr[0] = y
    for i in range(n_steps):
        y     = y + dt*np.asarray(f(t,y), dtype=_R)
        t     = t + dt
        t_arr[i+1] = t; y_arr[i+1] = y
    return t_arr, y_arr

def rk4(f, y0, t_span, n_steps=1000):
    """
    Classic 4th-order Runge-Kutta.
    Returns (t_arr, y_arr) float32.
    """
    t0,tf  = float(t_span[0]), float(t_span[1])
    dt     = _R((tf-t0)/n_steps)
    y      = np.asarray(y0, dtype=_R).copy()
    t      = _R(t0)
    t_arr  = np.zeros(n_steps+1, dtype=_R)
    y_arr  = np.zeros((n_steps+1,)+y.shape, dtype=_R)
    t_arr[0]=t; y_arr[0]=y
    for i in range(n_steps):
        k1 = dt*np.asarray(f(t,       y),         dtype=_R)
        k2 = dt*np.asarray(f(t+dt/2,  y+k1/2),    dtype=_R)
        k3 = dt*np.asarray(f(t+dt/2,  y+k2/2),    dtype=_R)
        k4 = dt*np.asarray(f(t+dt,    y+k3),       dtype=_R)
        y  = y + (k1+2*k2+2*k3+k4)/_R(6)
        t  = t + dt
        t_arr[i+1]=t; y_arr[i+1]=y
    return t_arr, y_arr

def solve_ode(f, y0, t_span, method="rk4", n_steps=500):
    """Dispatch to euler or rk4 solver. Returns (t, y)."""
    if method=="euler":
        return euler_method(f, y0, t_span, n_steps)
    return rk4(f, y0, t_span, n_steps)

def lorenz_system(sigma=10.0, rho=28.0, beta=8/3):
    """Return Lorenz attractor RHS function."""
    def f(t, y):
        x,y_,z = float(y[0]),float(y[1]),float(y[2])
        return np.array([sigma*(y_-x), x*(rho-z)-y_, x*y_-beta*z], dtype=np.float32)
    return f

def harmonic_oscillator(omega=1.0, zeta=0.0):
    """Return damped harmonic oscillator f(t,[x,v]) -> [v, a]."""
    def f(t, y):
        x,v = float(y[0]), float(y[1])
        a   = -omega**2*x - 2*zeta*omega*v
        return np.array([v, a], dtype=np.float32)
    return f
