use crate::divisibility::{DivisibilityRing, DivisibilityRingStore, Domain};
use crate::integer::IntegerRingStore;
use crate::iters::multi_cartesian_product;
use crate::primitive_int::StaticRing;
use crate::ring::*;
use crate::rings::multivariate::*;
use crate::rings::poly::*;
use crate::seq::*;
use crate::homomorphism::Homomorphism;
use crate::rings::poly::dense_poly::DensePolyRing;
use std::alloc::Allocator;
use std::cmp::min;
use std::ops::Range;
#[stability::unstable(feature = "enable")]
pub fn product_except_one<V, R>(ring: R, values: V, out: &mut [El<R>])
where R: RingStore,
V: VectorFn<El<R>>
{
assert_eq!(values.len(), out.len());
let n = values.len();
let log2_n = StaticRing::<i64>::RING.abs_log2_ceil(&n.try_into().unwrap()).unwrap();
assert!(n <= (1 << log2_n));
if n % 2 == 0 {
for i in 0..n {
out[i] = values.at(i ^ 1);
}
} else {
for i in 0..(n - 1) {
out[i] = values.at(i ^ 1);
}
out[n - 1] = ring.one();
}
for s in 1..log2_n {
for j in 0..(1 << (log2_n - s - 1)) {
let block_index = j << (s + 1);
if block_index + (1 << s) < n {
let (fst, snd) = (&mut out[block_index..min(n, block_index + (1 << (s + 1)))]).split_at_mut(1 << s);
let snd_block_prod = ring.mul_ref_fst(&snd[0], values.at(block_index + (1 << s)));
let fst_block_prod = ring.mul_ref_fst(&fst[0], values.at(block_index));
for i in 0..(1 << s) {
ring.mul_assign_ref(&mut fst[i], &snd_block_prod);
}
for i in 0..snd.len() {
ring.mul_assign_ref(&mut snd[i], &fst_block_prod);
}
}
}
}
}
#[stability::unstable(feature = "enable")]
#[derive(PartialEq, Eq, Hash, Debug, Clone, Copy)]
pub enum InterpolationError {
NotInvertible
}
#[stability::unstable(feature = "enable")]
pub fn interpolate<P, V1, V2, A: Allocator>(poly_ring: P, x: V1, y: V2, allocator: A) -> Result<El<P>, InterpolationError>
where P: RingStore,
P::Type: PolyRing,
<<P::Type as RingExtension>::BaseRing as RingStore>::Type: DivisibilityRing + Domain,
V1: VectorFn<El<<P::Type as RingExtension>::BaseRing>>,
V2: VectorFn<El<<P::Type as RingExtension>::BaseRing>>
{
assert_eq!(x.len(), y.len());
let R = poly_ring.base_ring();
let null_poly = poly_ring.prod(x.iter().map(|x| poly_ring.from_terms([(R.one(), 1), (R.negate(x), 0)])));
let mut nums = Vec::with_capacity_in(x.len(), &allocator);
let div_linear = |poly: &El<P>, a: &El<<P::Type as RingExtension>::BaseRing>| if let Some(d) = poly_ring.degree(poly) {
poly_ring.from_terms((0..d).rev().scan(R.zero(), |current, i| {
R.add_assign_ref(current, poly_ring.coefficient_at(poly, i + 1));
let result = R.clone_el(current);
R.mul_assign_ref(current, a);
return Some((result, i));
}))
} else { poly_ring.zero() };
nums.extend(x.iter().map(|x| div_linear(&null_poly, &x)));
let mut denoms = Vec::with_capacity_in(x.len(), &allocator);
denoms.extend((0..x.len()).map(|i| poly_ring.evaluate(&nums[i], &x.at(i), &R.identity())));
let mut factors = Vec::with_capacity_in(x.len(), &allocator);
factors.resize_with(x.len(), || R.zero());
product_except_one(R, (&denoms[..]).into_clone_ring_els(R), &mut factors);
let denominator = R.mul_ref(&factors[0], &denoms[0]);
for i in 0..x.len() {
R.mul_assign(&mut factors[i], y.at(i));
}
if let Some(inv) = R.invert(&denominator) {
Ok(poly_ring.inclusion().mul_map(<_ as RingStore>::sum(&poly_ring, nums.into_iter().zip(factors.into_iter()).map(|(num, c)| poly_ring.inclusion().mul_map(num, c))), inv))
} else {
let scaled_result = <_ as RingStore>::sum(&poly_ring, nums.into_iter().zip(factors.into_iter()).map(|(num, c)| poly_ring.inclusion().mul_map(num, c)));
poly_ring.try_from_terms(poly_ring.terms(&scaled_result).map(|(c, i)| R.checked_div(&c, &denominator).map(|c| (c, i)).ok_or(InterpolationError::NotInvertible)))
}
}
#[stability::unstable(feature = "enable")]
pub fn interpolate_multivariate<P, V1, V2, A, A2>(poly_ring: P, interpolation_points: V1, mut values: Vec<El<<P::Type as RingExtension>::BaseRing>, A2>, allocator: A) -> Result<El<P>, InterpolationError>
where P: RingStore,
P::Type: MultivariatePolyRing,
<<P::Type as RingExtension>::BaseRing as RingStore>::Type: DivisibilityRing + Domain,
V1: VectorFn<V2>,
V2: VectorFn<El<<P::Type as RingExtension>::BaseRing>>,
A: Allocator,
A2: Allocator
{
let dim_prod = |range: Range<usize>| <_ as RingStore>::prod(&StaticRing::<i64>::RING, range.map(|i| interpolation_points.at(i).len().try_into().unwrap())) as usize;
assert_eq!(interpolation_points.len(), poly_ring.indeterminate_count());
let n = poly_ring.indeterminate_count();
assert_eq!(values.len(), dim_prod(0..n));
let uni_poly_ring = DensePolyRing::new_with_convolution(poly_ring.base_ring(), "X", &allocator, STANDARD_CONVOLUTION);
for i in (0..n).rev() {
let leading_dim = dim_prod((i + 1)..n);
let outer_block_count = dim_prod(0..i);
let len = interpolation_points.at(i).len();
let outer_block_size = leading_dim * len;
for outer_block_index in 0..outer_block_count {
for inner_block_index in 0..leading_dim {
let block_start = inner_block_index + outer_block_index * outer_block_size;
let poly = interpolate(&uni_poly_ring, interpolation_points.at(i), (&values[..]).into_clone_ring_els(poly_ring.base_ring()).restrict(block_start..(block_start + outer_block_size + 1 - leading_dim)).step_by_fn(leading_dim), &allocator)?;
for j in 0..len {
values[block_start + leading_dim * j] = poly_ring.base_ring().clone_el(uni_poly_ring.coefficient_at(&poly, j));
}
}
}
}
return Ok(poly_ring.from_terms(
multi_cartesian_product((0..n).map(|i| 0..interpolation_points.at(i).len()), |idxs| poly_ring.get_ring().create_monomial(idxs.iter().map(|e| *e)), |_, x| *x)
.zip(values.into_iter())
.map(|(m, c)| (c, m))
));
}
#[cfg(test)]
use crate::rings::zn::zn_64::Zn;
#[cfg(test)]
use crate::rings::zn::ZnRingStore;
#[cfg(test)]
use std::alloc::Global;
#[cfg(test)]
use multivariate_impl::MultivariatePolyRingImpl;
use super::convolution::STANDARD_CONVOLUTION;
#[test]
fn test_product_except_one() {
let ring = StaticRing::<i64>::RING;
let data = [2, 3, 5, 7, 11, 13, 17, 19];
let mut actual = [0; 8];
let expected = [
3 * 5 * 7 * 11 * 13 * 17 * 19,
2 * 5 * 7 * 11 * 13 * 17 * 19,
2 * 3 * 7 * 11 * 13 * 17 * 19,
2 * 3 * 5 * 11 * 13 * 17 * 19,
2 * 3 * 5 * 7 * 13 * 17 * 19,
2 * 3 * 5 * 7 * 11 * 17 * 19,
2 * 3 * 5 * 7 * 11 * 13 * 19,
2 * 3 * 5 * 7 * 11 * 13 * 17
];
product_except_one(&ring, (&data[..]).clone_els_by(|x| *x), &mut actual);
assert_eq!(expected, actual);
let data = [2, 3, 5, 7, 11, 13, 17];
let mut actual = [0; 7];
let expected = [
3 * 5 * 7 * 11 * 13 * 17,
2 * 5 * 7 * 11 * 13 * 17,
2 * 3 * 7 * 11 * 13 * 17,
2 * 3 * 5 * 11 * 13 * 17,
2 * 3 * 5 * 7 * 13 * 17,
2 * 3 * 5 * 7 * 11 * 17,
2 * 3 * 5 * 7 * 11 * 13
];
product_except_one(&ring, (&data[..]).clone_els_by(|x| *x), &mut actual);
assert_eq!(expected, actual);
let data = [2, 3, 5, 7, 11, 13];
let mut actual = [0; 6];
let expected = [
3 * 5 * 7 * 11 * 13,
2 * 5 * 7 * 11 * 13,
2 * 3 * 7 * 11 * 13,
2 * 3 * 5 * 11 * 13,
2 * 3 * 5 * 7 * 13,
2 * 3 * 5 * 7 * 11
];
product_except_one(&ring, (&data[..]).clone_els_by(|x| *x), &mut actual);
assert_eq!(expected, actual);
}
#[test]
fn test_interpolate() {
let ring = StaticRing::<i64>::RING;
let poly_ring = DensePolyRing::new(ring, "X");
let poly = poly_ring.from_terms([(3, 0), (1, 1), (-1, 3), (2, 4), (1, 5)].into_iter());
let x = (0..6).map_fn(|x| x.try_into().unwrap());
let actual = interpolate(&poly_ring, x.clone(), x.map_fn(|x| poly_ring.evaluate(&poly, &x, &ring.identity())), Global).unwrap();
assert_el_eq!(&poly_ring, &poly, &actual);
let ring = Zn::new(29).as_field().ok().unwrap();
let poly_ring = DensePolyRing::new(ring, "X");
let x = (0..5).map_fn(|x| ring.int_hom().map(x as i32));
let y = (0..5).map_fn(|x| if x == 3 { ring.int_hom().map(6) } else { ring.zero() });
let poly = interpolate(&poly_ring, x.clone(), y.clone(), Global).unwrap();
for i in 0..5 {
assert_el_eq!(ring, y.at(i), poly_ring.evaluate(&poly, &x.at(i), &ring.identity()));
}
}
#[test]
fn test_interpolate_multivariate() {
let ring = Zn::new(29).as_field().ok().unwrap();
let poly_ring: MultivariatePolyRingImpl<_> = MultivariatePolyRingImpl::new(ring, 2);
let interpolation_points = (0..2).map_fn(|_| (0..5).map_fn(|x| ring.int_hom().map(x as i32)));
let values = (0..25).map(|x| ring.int_hom().map(x & 1)).collect::<Vec<_>>();
let poly = interpolate_multivariate(&poly_ring, &interpolation_points, values, Global).unwrap();
for x in 0..5 {
for y in 0..5 {
let expected = (x * 5 + y) & 1;
assert_el_eq!(ring, ring.int_hom().map(expected), poly_ring.evaluate(&poly, [ring.int_hom().map(x), ring.int_hom().map(y)].into_clone_ring_els(&ring), &ring.identity()));
}
}
let poly_ring: MultivariatePolyRingImpl<_> = MultivariatePolyRingImpl::new(ring, 3);
let interpolation_points = (0..3).map_fn(|i| (0..(i + 2)).map_fn(|x| ring.int_hom().map(x as i32)));
let values = (0..24).map(|x| ring.int_hom().map(x / 2)).collect::<Vec<_>>();
let poly = interpolate_multivariate(&poly_ring, &interpolation_points, values, Global).unwrap();
for x in 0..2 {
for y in 0..3 {
for z in 0..4 {
let expected = (x * 12 + y * 4 + z) / 2;
assert_el_eq!(ring, ring.int_hom().map(expected), poly_ring.evaluate(&poly, [ring.int_hom().map(x), ring.int_hom().map(y), ring.int_hom().map(z)].into_clone_ring_els(&ring), &ring.identity()));
}
}
}
}
#[test]
#[ignore]
fn large_polynomial_interpolation() {
let field = Zn::new(65537).as_field().ok().unwrap();
let poly_ring = DensePolyRing::new(field, "X");
let hom = poly_ring.base_ring().can_hom(&StaticRing::<i64>::RING).unwrap();
let actual = interpolate(&poly_ring, (0..65536).map_fn(|x| hom.map(x as i64)), (0..65536).map_fn(|x| hom.map(x as i64)), Global).unwrap();
assert_el_eq!(&poly_ring, poly_ring.indeterminate(), actual);
}