use std::collections::BTreeMap;

use super::*;
use crate::{random::unique_integers, Function, PolynomialParameters};
use anyhow::{anyhow, Result};
use proptest::prelude::*;

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

    fn arbitrary_with(_params: Self::Parameters) -> Self::Strategy {
        prop_oneof![
            Just(Equality::EqualToZero),
            Just(Equality::LessThanOrEqualToZero),
        ]
        .boxed()
    }
}

impl Arbitrary for Constraint<Created> {
    type Parameters = PolynomialParameters;
    type Strategy = BoxedStrategy<Self>;
    fn arbitrary_with(params: Self::Parameters) -> Self::Strategy {
        (Function::arbitrary_with(params), Equality::arbitrary())
            .prop_map(|(function, equality)| Constraint {
                equality,
                stage: CreatedData { function },
            })
            .boxed()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ConstraintIDParameters {
    size: usize,
    max_id: ConstraintID,
}

impl ConstraintIDParameters {
    pub fn new(size: usize, max_id: ConstraintID) -> Result<Self> {
        if size > max_id.0 as usize + 1 {
            return Err(anyhow!(
                "size {} is greater than `max_id {} + 1`",
                size,
                max_id.0
            ));
        }
        Ok(Self { size, max_id })
    }
}

impl Default for ConstraintIDParameters {
    fn default() -> Self {
        Self {
            size: 5,
            max_id: ConstraintID(10),
        }
    }
}

pub fn arbitrary_constraints(
    id_parameters: ConstraintIDParameters,
    parameters: PolynomialParameters,
) -> impl Strategy<Value = BTreeMap<ConstraintID, Constraint<Created>>> {
    let unique_ids_strategy = unique_integers(0, id_parameters.max_id.0, id_parameters.size);
    let constraints_strategy =
        proptest::collection::vec(Constraint::arbitrary_with(parameters), id_parameters.size);
    (unique_ids_strategy, constraints_strategy)
        .prop_map(|(ids, constraints)| {
            ids.into_iter()
                .map(ConstraintID::from)
                .zip(constraints)
                .collect()
        })
        .boxed()
}
