ommx 3.0.0-beta.5

use super::{
    operation::{
        associative_expression, binary_expression, unary_expression, AssociativeOperator,
        BinaryOperator, UnaryOperator,
    },
    *,
};
use crate::{
    Coefficient, Linear, LinearParameters, Polynomial, PolynomialParameters, Quadratic,
    QuadraticParameters, VariableID,
};
use proptest::{prelude::*, strategy::Union};

const DEFAULT_MAX_DEPTH: u32 = 4;
const DEFAULT_DESIRED_SIZE: u32 = 16;
const MIN_POWI_EXPONENT: i32 = -3;
const MAX_POWI_EXPONENT: i32 = 3;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum FunctionSpace {
    PolynomialOnly,
    Full,
}

/// Parameters for random [`Function`] generation.
///
/// [`FunctionParameters::default`] and [`FunctionParameters::full`] generate
/// both compact polynomials and every composed operation kind. The composed
/// strategy uses a bounded depth and small integer-power exponents so generated
/// cases remain practical; these sampling bounds are not semantic limits on
/// [`Function`]. Use
/// [`FunctionParameters::polynomial_only`] when a property or consumer has an
/// explicit polynomial precondition.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FunctionParameters {
    polynomial: PolynomialParameters,
    space: FunctionSpace,
}

impl FunctionParameters {
    /// Generate compact leaves and every composed operation kind using the
    /// supplied polynomial leaf parameters.
    pub fn full(polynomial: PolynomialParameters) -> Self {
        Self {
            polynomial,
            space: FunctionSpace::Full,
        }
    }

    /// Generate only compact polynomial `Function` variants.
    pub fn polynomial_only(polynomial: PolynomialParameters) -> Self {
        Self {
            polynomial,
            space: FunctionSpace::PolynomialOnly,
        }
    }

    /// Parameters used to generate compact polynomial leaves.
    pub fn polynomial_parameters(&self) -> PolynomialParameters {
        self.polynomial
    }

    /// Largest variable ID generated by any compact leaf.
    pub fn max_id(&self) -> VariableID {
        self.polynomial.max_id()
    }

    fn is_polynomial_only(&self) -> bool {
        self.space == FunctionSpace::PolynomialOnly
    }
}

impl Default for FunctionParameters {
    fn default() -> Self {
        Self::full(PolynomialParameters::default())
    }
}

impl Arbitrary for FunctionParameters {
    type Parameters = ();
    type Strategy = BoxedStrategy<Self>;

    fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
        (PolynomialParameters::arbitrary(), any::<bool>())
            .prop_map(|(polynomial, polynomial_only)| {
                if polynomial_only {
                    Self::polynomial_only(polynomial)
                } else {
                    Self::full(polynomial)
                }
            })
            .boxed()
    }
}

fn arbitrary_polynomial_function(p: PolynomialParameters) -> BoxedStrategy<Function> {
    if p.num_terms() == 0 {
        return prop_oneof![
            Just(Function::Zero),
            Just(Function::Linear(Linear::zero())),
            Just(Function::Quadratic(Quadratic::zero())),
            Just(Function::Polynomial(Polynomial::zero())),
        ]
        .boxed();
    }
    if p.max_degree() == 0 {
        debug_assert_eq!(p.num_terms(), 1);
        return Coefficient::arbitrary()
            .prop_map(Function::Constant)
            .boxed();
    }
    let polynomial = Polynomial::arbitrary_with(p);
    let linear = LinearParameters::new(p.num_terms(), p.max_id())
        .ok()
        .map(Linear::arbitrary_with);
    let quad = if p.max_degree() == 1 {
        linear
            .clone()
            .map(|linear| linear.prop_map(Quadratic::from).boxed())
    } else {
        QuadraticParameters::new(p.num_terms(), p.max_id())
            .ok()
            .map(|parameters| Quadratic::arbitrary_with(parameters).boxed())
    };
    let mut candidates = vec![polynomial.prop_map(Function::Polynomial).boxed()];
    if let Some(linear) = linear {
        candidates.push(linear.prop_map(Function::Linear).boxed());
    }
    if let Some(quadratic) = quad {
        candidates.push(quadratic.prop_map(Function::Quadratic).boxed());
    }
    Union::new(candidates).boxed()
}

impl Arbitrary for Function {
    type Parameters = FunctionParameters;
    type Strategy = BoxedStrategy<Self>;

    fn arbitrary_with(parameters: Self::Parameters) -> Self::Strategy {
        let leaf = arbitrary_polynomial_function(parameters.polynomial_parameters());
        if parameters.is_polynomial_only() {
            return leaf;
        }

        leaf.prop_recursive(DEFAULT_MAX_DEPTH, DEFAULT_DESIRED_SIZE, 2, |inner| {
            prop_oneof![
                inner
                    .clone()
                    .prop_map(|operand| unary_expression(UnaryOperator::Neg, operand)),
                inner
                    .clone()
                    .prop_map(|operand| unary_expression(UnaryOperator::Abs, operand)),
                inner
                    .clone()
                    .prop_map(|operand| unary_expression(UnaryOperator::Signum, operand)),
                (inner.clone(), MIN_POWI_EXPONENT..=MAX_POWI_EXPONENT).prop_map(
                    |(operand, exponent)| {
                        unary_expression(UnaryOperator::Powi(exponent), operand)
                    },
                ),
                (inner.clone(), inner.clone()).prop_map(|(lhs, rhs)| {
                    associative_expression(AssociativeOperator::Add, lhs, rhs)
                }),
                (inner.clone(), inner.clone()).prop_map(|(lhs, rhs)| {
                    associative_expression(AssociativeOperator::Mul, lhs, rhs)
                }),
                (inner.clone(), inner.clone()).prop_map(|(lhs, rhs)| {
                    associative_expression(AssociativeOperator::Min, lhs, rhs)
                }),
                (inner.clone(), inner.clone()).prop_map(|(lhs, rhs)| {
                    associative_expression(AssociativeOperator::Max, lhs, rhs)
                }),
                (inner.clone(), inner)
                    .prop_map(|(lhs, rhs)| { binary_expression(BinaryOperator::Div, lhs, rhs) }),
            ]
        })
        .boxed()
    }
}

#[cfg(test)]
mod tests {
    use super::super::operation::{instructions, Instruction};
    use super::*;
    use crate::Evaluate;

    proptest! {
        #[test]
        fn polynomial_only_space_generates_only_polynomials(
            (parameters, function) in PolynomialParameters::arbitrary().prop_flat_map(|parameters| {
                Function::arbitrary_with(FunctionParameters::polynomial_only(parameters))
                    .prop_map(move |function| (parameters, function))
            }),
        ) {
            prop_assert!(function.is_polynomial());
            prop_assert_eq!(function.num_terms(), Some(parameters.num_terms()));
            prop_assert!(function.degree().is_some_and(|degree| degree <= parameters.max_degree()));
        }

        #[test]
        fn full_space_generates_valid_functions(function in any::<Function>()) {
            if let Function::Expression(expression) = function {
                let instructions = instructions(&expression);
                let contains_operation = instructions
                    .iter()
                    .any(|instruction| !matches!(instruction, Instruction::Push(_)));
                prop_assert!(instructions.len() >= 2);
                prop_assert!(contains_operation);
            }
        }

        #[test]
        fn expression_operations_collect_operand_ids(lhs in any::<Function>(), rhs in any::<Function>()) {
            let mut expected = lhs.required_ids();
            expected.extend(rhs.required_ids());

            let functions = [
                associative_expression(AssociativeOperator::Add, lhs.clone(), rhs.clone()),
                associative_expression(AssociativeOperator::Mul, lhs.clone(), rhs.clone()),
                associative_expression(AssociativeOperator::Min, lhs.clone(), rhs.clone()),
                associative_expression(AssociativeOperator::Max, lhs.clone(), rhs.clone()),
                binary_expression(BinaryOperator::Div, lhs.clone(), rhs),
            ];
            for function in functions {
                prop_assert_eq!(function.required_ids(), expected.clone());
            }

            prop_assert_eq!(
                unary_expression(UnaryOperator::Powi(2), lhs.clone()).required_ids(),
                lhs.required_ids(),
            );
        }
    }

    #[test]
    fn default_space_reaches_composed_expressions() {
        let strategy = Function::arbitrary();
        let mut runner = proptest::test_runner::TestRunner::deterministic();
        let reaches_expression = (0..64).any(|_| {
            matches!(
                strategy
                    .new_tree(&mut runner)
                    .expect("full Function strategy must generate")
                    .current(),
                Function::Expression(_)
            )
        });
        assert!(reaches_expression);
    }
}