"""
ODE Solvers
===========

Numerical methods for solving ordinary differential equations.

Author: VIX Spike VolOfVol Research Team
Date: March 2026
"""

import numpy as np
from typing import Callable, Tuple


def odeRK4(fcn, a, b, y0, N, *var):
    """
    ODE solver with RK4 algorithm
    Solve dY(t)/dt = f(t, Y(t)) in 'N' steps using 4th-order Runge-Kutta witdt initial condition Y[a] = y0.
    where Y(t,u), given each u, is the desired coefficient function, e.g. A(t,u), B(t,u)
    
    Parameters:
    fcn: RHS function in coefficient functions
    a,b: time interval; usually set a=0(meaning tau=0), b=T; 以到期日条件作为初始条件
    N: # of discretisation steps

    Return:
    t, Y(t,u): time path; coef function path along t, with given u
    """
    dt = (b - a) / N
    t = a + np.arange(N + 1) * dt
    nDim = var[0] # dimention of u (in ChF)
    y = np.zeros([int(nDim), t.size], dtype=complex)
    y[:,0] = y0

    for k in range(N):
        k1 = fcn(t[k], y[:,k])
        k2 = fcn(t[k] + dt / 2, y[:,k] + dt * k1 / 2)
        k3 = fcn(t[k] + dt / 2, y[:,k] + dt * k2 / 2)
        k4 = fcn(t[k] + dt, y[:,k] + dt * k3)
        y[:,k+1] = y[:,k] + dt * (k1 + 2 * (k2 + k3) + k4) / 6

    return t, y
