"""
Example usage for ODE solvers. Can be run as a script.
"""
import numpy as np
from solve_first_order_ode import solve_first_order_ode
from solve_second_order_ode import solve_second_order_ode

if __name__ == '__main__':
    # --- First-order ODE example: dy/dt = -y*t with y(0)=1 ---
    def first_order_example(t, y):
        return -y * t

    print("Solving dy/dt = -y*t, y(0)=1 from t=0 to t=5")
    solution1 = solve_first_order_ode(first_order_example, (0, 5), [1], t_eval_count=50)
    if solution1['success']:
        print(f"First-order ODE solved. Message: {solution1['message']}")
        if solution1['plot_path']:
            print(f"Plot saved to {solution1['plot_path']}")
    else:
        print(f"First-order ODE failed. Message: {solution1['message']}")

    # --- First-order system example: Lotka-Volterra ---
    a, b, c, d = 1.5, 0.8, 0.5, 0.9
    def lotka_volterra(t, y):
        prey, predator = y[0], y[1]
        d_prey_dt = a * prey - b * prey * predator
        d_predator_dt = c * prey * predator - d * predator
        return [d_prey_dt, d_predator_dt]

    print("\nSolving Lotka-Volterra system from t=0 to t=20 with y0=[10, 5]")
    solution_lv = solve_first_order_ode(lotka_volterra, (0, 20), [10, 5], t_eval_count=200)
    if solution_lv['success']:
        print(f"Lotka-Volterra solved. Plot saved to {solution_lv['plot_path']}")
    else:
        print(f"Lotka-Volterra failed. Message: {solution_lv['message']}")

    # --- Second-order ODE example: d²y/dt² = -sin(y) (simple pendulum) ---
    g_L = 9.81 / 1.0 # Example: g/L = 9.81
    def pendulum_ode(t, y_angle, dy_dt_angular_velocity):
        return -g_L * np.sin(y_angle)

    print("\nSolving d²y/dt² = -g/L*sin(y), y(0)=pi/4, dy/dt(0)=0 from t=0 to t=10")
    solution2 = solve_second_order_ode(pendulum_ode, (0, 10), y0=np.pi/4, dy0_dt=0, t_eval_count=100)
    if solution2['success']:
        print(f"Second-order ODE solved. Message: {solution2['message']}")
        if solution2['plot_path']:
            print(f"Plot saved to {solution2['plot_path']}")
    else:
        print(f"Second-order ODE failed. Message: {solution2['message']}")