use ::burn::tensor::{Bool, Tensor, backend::Backend};
use super::super::roots;
pub type Coefficients<B, const D: usize> = (Tensor<B, D>, Tensor<B, D>);
pub fn evaluate<B: Backend, const D: usize>(
x_coefficients: Coefficients<B, D>,
y_coefficients: Coefficients<B, D>,
x: Tensor<B, D>,
y: Tensor<B, D>,
) -> Tensor<B, D> {
let (x1, x0) = x_coefficients;
let (y1, y0) = y_coefficients;
let roots = roots::linear::solve(y1.clone(), y0 - y);
let x_offset = x - x0;
let sign = y1.greater_elem(0.0).float() * 2.0 - 1.0;
root_contribution(roots.t, roots.valid, x_offset, x1, sign)
}
pub fn point_coefficients<B: Backend, const D: usize>(
p0: Tensor<B, D>,
p1: Tensor<B, D>,
) -> Coefficients<B, D> {
(p1 - p0.clone(), p0)
}
fn root_contribution<B: Backend, const D: usize>(
t: Tensor<B, D>,
valid: Tensor<B, D, Bool>,
x_offset: Tensor<B, D>,
x_coefficient: Tensor<B, D>,
sign: Tensor<B, D>,
) -> Tensor<B, D> {
let curve_x = x_coefficient * t.clone();
let valid = valid
.bool_and(t.clone().greater_equal_elem(0.0))
.bool_and(t.lower_equal_elem(1.0))
.bool_and(curve_x.greater(x_offset));
valid.float() * sign
}