zoe 0.0.30

use crate::{
    alignment::phmm::{
        CorePhmm, EmissionParams, LayerParams, PhmmNumber, PhmmState, TransitionParams, indexing::GetLayerMut,
    },
    data::arbitrary::{ArbitrarySpecs, ArraySpecs, VecSpecs},
};
use arbitrary::{Arbitrary, Result, Unstructured};

impl<'a, T, const S: usize> Arbitrary<'a> for EmissionParams<T, S>
where
    T: Arbitrary<'a>,
{
    #[inline]
    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
        Ok(EmissionParams::from_array(<[T; S]>::arbitrary(u)?))
    }
}

impl<'a, T> Arbitrary<'a> for TransitionParams<T>
where
    T: Arbitrary<'a>,
{
    #[inline]
    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
        Ok(TransitionParams(<[[T; 3]; 3]>::arbitrary(u)?))
    }
}

impl<'a, T, const S: usize> Arbitrary<'a> for LayerParams<T, S>
where
    T: Arbitrary<'a>,
{
    #[inline]
    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
        Ok(LayerParams {
            transition:      TransitionParams::arbitrary(u)?,
            emission_match:  EmissionParams::arbitrary(u)?,
            emission_insert: EmissionParams::arbitrary(u)?,
        })
    }
}

impl<'a, T, const S: usize> Arbitrary<'a> for CorePhmm<T, S>
where
    T: Arbitrary<'a>,
{
    #[inline]
    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
        let mut layers = vec![LayerParams::<T, S>::arbitrary(u)?, LayerParams::<T, S>::arbitrary(u)?];
        layers.extend(Vec::<LayerParams<T, S>>::arbitrary(u)?);
        Ok(CorePhmm::new_unchecked(layers))
    }
}

/// Specifications for generating an arbitrary [`CorePhmm`].
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Default)]
pub struct CorePhmmSpecs<K, const S: usize> {
    /// The specifications for generating the parameters.
    pub param_specs: K,

    /// Whether to disallow invalid transitions/emissions for the first and last
    /// layers by setting the parameters to infinity (probability zero).
    pub disallow_invalid: bool,
}

impl<'a, K, const S: usize> ArbitrarySpecs<'a> for CorePhmmSpecs<K, S>
where
    K: ArbitrarySpecs<'a, Output: PhmmNumber> + Copy,
{
    type Output = CorePhmm<K::Output, S>;

    #[inline]
    fn make_arbitrary(&self, u: &mut Unstructured<'a>) -> Result<Self::Output> {
        use PhmmState::*;

        let specs = VecSpecs {
            element_specs: LayerParamsSpecs {
                param_specs: self.param_specs,
            },
            min_len:       2,
            len:           None,
            max_len:       usize::MAX,
        };

        let mut core = CorePhmm::new_unchecked(specs.make_arbitrary(u)?);

        if self.disallow_invalid {
            let first_layer = core.begin_layer_mut();
            first_layer.transition[(Delete, Delete)] = K::Output::INFINITY;
            first_layer.transition[(Delete, Match)] = K::Output::INFINITY;
            first_layer.transition[(Delete, Insert)] = K::Output::INFINITY;

            let last_layer = core.last_match_mut();
            last_layer.transition[(Delete, Delete)] = K::Output::INFINITY;
            last_layer.transition[(Insert, Delete)] = K::Output::INFINITY;
            last_layer.transition[(Match, Delete)] = K::Output::INFINITY;

            last_layer.emission_match = EmissionParams::default();
        }

        Ok(core)
    }
}

/// Specifications for generating an arbitrary [`EmissionParams`] struct.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Default)]
pub struct EmissionParamsSpecs<K, const S: usize> {
    /// The specifications for generating the parameters.
    pub param_specs: K,
}

impl<'a, K, const S: usize> ArbitrarySpecs<'a> for EmissionParamsSpecs<K, S>
where
    K: ArbitrarySpecs<'a, Output: Default> + Copy,
{
    type Output = EmissionParams<K::Output, S>;

    #[inline]
    fn make_arbitrary(&self, u: &mut Unstructured<'a>) -> Result<Self::Output> {
        let specs = ArraySpecs {
            element_specs: self.param_specs,
        };

        Ok(EmissionParams::from_array(specs.make_arbitrary(u)?))
    }
}

/// Specifications for generating an arbitrary [`TransitionParamsSpecs`] struct.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Default)]
pub struct TransitionParamsSpecs<K> {
    /// The specifications for generating the parameters.
    param_specs: K,
}

impl<'a, K> ArbitrarySpecs<'a> for TransitionParamsSpecs<K>
where
    K: ArbitrarySpecs<'a, Output: Default> + Copy,
{
    type Output = TransitionParams<K::Output>;

    #[inline]
    fn make_arbitrary(&self, u: &mut Unstructured<'a>) -> Result<Self::Output> {
        let specs = ArraySpecs {
            element_specs: ArraySpecs {
                element_specs: self.param_specs,
            },
        };

        Ok(TransitionParams(specs.make_arbitrary(u)?))
    }
}

/// Specifications for generating an arbitrary [`LayerParams`] struct.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Default)]
pub struct LayerParamsSpecs<K, const S: usize> {
    /// The specifications for generating the parameters.
    param_specs: K,
}

impl<'a, K, const S: usize> ArbitrarySpecs<'a> for LayerParamsSpecs<K, S>
where
    K: ArbitrarySpecs<'a, Output: Default> + Copy,
{
    type Output = LayerParams<K::Output, S>;

    #[inline]
    fn make_arbitrary(&self, u: &mut Unstructured<'a>) -> Result<Self::Output> {
        let transition_specs = TransitionParamsSpecs {
            param_specs: self.param_specs,
        };

        let emission_specs = EmissionParamsSpecs {
            param_specs: self.param_specs,
        };

        Ok(LayerParams {
            transition:      transition_specs.make_arbitrary(u)?,
            emission_match:  emission_specs.make_arbitrary(u)?,
            emission_insert: emission_specs.make_arbitrary(u)?,
        })
    }
}