1use ocas_atom::Symbol;
4use ocas_calc::ode::{
5 ODE, ODESolution, ODEType, classify_ode as rs_classify, dsolve as rs_dsolve,
6 dsolve_ivp as rs_dsolve_ivp,
7};
8use pyo3::exceptions::PyValueError;
9use pyo3::prelude::*;
10
11use crate::expression::Expression;
12
13#[pyfunction]
22#[pyo3(name = "classify_ode")]
23pub fn py_classify_ode(equation: &Expression, func: &str, var: &str) -> PyResult<Vec<String>> {
24 let (ctx, ode) = build_ode(equation, func, var)?;
25 Ok(rs_classify(ctx, ode)
26 .iter()
27 .map(ode_type_name)
28 .map(str::to_owned)
29 .collect())
30}
31
32#[pyfunction]
43#[pyo3(name = "dsolve", signature = (equation, func, var, hint=None))]
44pub fn py_dsolve(
45 equation: &Expression,
46 func: &str,
47 var: &str,
48 hint: Option<&str>,
49) -> PyResult<String> {
50 let hint_type = hint.map(parse_ode_type).transpose()?;
51 let (ctx, ode) = build_ode(equation, func, var)?;
52 let sol = rs_dsolve(ctx, ode, hint_type);
53 Ok(format_solution(&sol))
54}
55
56#[pyfunction]
64#[pyo3(name = "dsolve_ivp", signature = (equation, func, var, y0, y1=None))]
65pub fn py_dsolve_ivp(
66 equation: &Expression,
67 func: &str,
68 var: &str,
69 y0: &str,
70 y1: Option<&str>,
71) -> PyResult<String> {
72 let (ctx, ode) = build_ode(equation, func, var)?;
73 let y0_atom = parse_in(ctx, y0)?;
74 let y1_atom = y1.map(|s| parse_in(ctx, s)).transpose()?;
75 let sol = rs_dsolve_ivp(ctx, ode, y0_atom, y1_atom);
76 Ok(format_solution(&sol))
77}
78
79fn build_ode(
86 equation: &Expression,
87 func: &str,
88 var: &str,
89) -> PyResult<(&'static ocas_atom::AtomArena<'static>, ODE<'static>)> {
90 let ctx = equation.ctx_ref();
91 let x = ctx.var(var);
92 let func_atom = ctx.fun(func, &[x]);
93 let ode = ODE {
94 equation: equation.atom(),
95 func: func_atom,
96 var: Symbol::new(var),
97 };
98 Ok((ctx, ode))
99}
100
101fn parse_in<'a>(ctx: &'a ocas_atom::AtomArena<'a>, input: &str) -> PyResult<ocas_atom::Atom<'a>> {
102 ocas_parse::parse(ctx, input).map_err(|e| PyValueError::new_err(format!("parse error: {e}")))
103}
104
105fn ode_type_name(t: &ODEType) -> &'static str {
106 match t {
107 ODEType::Separable => "Separable",
108 ODEType::LinearFirst => "LinearFirst",
109 ODEType::Bernoulli => "Bernoulli",
110 ODEType::Exact => "Exact",
111 ODEType::Homogeneous => "Homogeneous",
112 ODEType::LinearConstantCoeff => "LinearConstantCoeff",
113 ODEType::CauchyEuler => "CauchyEuler",
114 ODEType::ReductionOfOrder => "ReductionOfOrder",
115 ODEType::PowerSeries => "PowerSeries",
116 }
117}
118
119fn parse_ode_type(name: &str) -> PyResult<ODEType> {
120 match name {
121 "Separable" => Ok(ODEType::Separable),
122 "LinearFirst" => Ok(ODEType::LinearFirst),
123 "Bernoulli" => Ok(ODEType::Bernoulli),
124 "Exact" => Ok(ODEType::Exact),
125 "Homogeneous" => Ok(ODEType::Homogeneous),
126 "LinearConstantCoeff" => Ok(ODEType::LinearConstantCoeff),
127 "CauchyEuler" => Ok(ODEType::CauchyEuler),
128 "ReductionOfOrder" => Ok(ODEType::ReductionOfOrder),
129 "PowerSeries" => Ok(ODEType::PowerSeries),
130 other => Err(PyValueError::new_err(format!(
131 "unknown ODE type hint: {other}"
132 ))),
133 }
134}
135
136fn format_solution(sol: &ODESolution<'_>) -> String {
137 match sol {
138 ODESolution::Explicit(e) => format!("y = {e}"),
139 ODESolution::Implicit(e) => format!("{e} = C"),
140 ODESolution::Parametric(a, b) => format!("x = {a}, y = {b}"),
141 ODESolution::Series(e, n) => format!("series({n} terms): y = {e}"),
142 ODESolution::System(comps) => {
143 let joined: Vec<String> = comps
144 .iter()
145 .enumerate()
146 .map(|(i, c)| format!("y{} = {c}", i + 1))
147 .collect();
148 joined.join(", ")
149 }
150 ODESolution::Unsolved(_) => "unsolved".to_string(),
151 }
152}