/*!
Helper traits to work with references or owned values of types and their subsets.
*/

#[derive(Clone, Copy)]
pub enum EitherDecomp<A, B> {
    Owned(A),
    Borrowed(B),
}

/// A very basic implementation of [`std::borrow::Cow`] without a [`Clone`] trait dependency.
pub type MyCow<'b, X> = EitherDecomp<X, &'b X>;

impl<'b, X> std::ops::Deref for MyCow<'b, X> {
    type Target = X;

    #[inline]
    fn deref(&self) -> &Self::Target {
        match self {
            EitherDecomp::Owned(x) => &x,
            EitherDecomp::Borrowed(x) => x,
        }
    }
}

/// Trait for ownable-by-consumption objects
pub trait Ownable {
    type OwnedVariant: Clone;

    /// Returns an owned instance, possibly consuming the original,
    /// avoiding cloning when possible.
    fn into_owned(self) -> Self::OwnedVariant;

    /// Returns an owned instance of a reference.
    fn clone_owned(&self) -> Self::OwnedVariant;

    /// Returns an owned instance or a reference to one.
    fn cow_owned<'b>(self) -> MyCow<'b, Self::OwnedVariant>
    where
        Self: 'b;

    /// Returns an owned instance or a reference to one.
    fn ref_cow_owned<'b>(&'b self) -> MyCow<'b, Self::OwnedVariant>
    where
        Self: 'b;
}

impl<'a, X: Ownable> Ownable for &'a X {
    type OwnedVariant = X::OwnedVariant;

    #[inline]
    fn into_owned(self) -> Self::OwnedVariant {
        X::clone_owned(self)
    }

    #[inline]
    fn clone_owned(&self) -> Self::OwnedVariant {
        X::clone_owned(self)
    }

    #[inline]
    fn cow_owned<'b>(self) -> MyCow<'b, Self::OwnedVariant>
    where
        Self: 'b,
    {
        X::ref_cow_owned(self)
    }

    fn ref_cow_owned<'b>(&self) -> MyCow<'b, Self::OwnedVariant>
    where
        Self: 'b,
    {
        X::ref_cow_owned(self)
    }
}

impl<'a, X: Ownable> Ownable for &'a mut X {
    type OwnedVariant = X::OwnedVariant;

    #[inline]
    fn into_owned(self) -> Self::OwnedVariant {
        X::clone_owned(self)
    }

    #[inline]
    fn clone_owned(&self) -> Self::OwnedVariant {
        X::clone_owned(self)
    }

    #[inline]
    fn cow_owned<'b>(self) -> MyCow<'b, Self::OwnedVariant>
    where
        Self: 'b,
    {
        X::ref_cow_owned(self)
    }

    fn ref_cow_owned<'b>(&'b self) -> MyCow<'b, Self::OwnedVariant>
    where
        Self: 'b,
    {
        X::ref_cow_owned(self)
    }
}

impl<'a, A, B> Ownable for EitherDecomp<A, B>
where
    A: Ownable,
    B: Ownable<OwnedVariant = A::OwnedVariant>,
{
    type OwnedVariant = A::OwnedVariant;

    #[inline]
    fn into_owned(self) -> Self::OwnedVariant {
        match self {
            EitherDecomp::Owned(a) => A::into_owned(a),
            EitherDecomp::Borrowed(b) => B::into_owned(b),
        }
    }

    #[inline]
    fn clone_owned(&self) -> Self::OwnedVariant {
        match self {
            EitherDecomp::Owned(a) => A::clone_owned(a),
            EitherDecomp::Borrowed(b) => B::clone_owned(b),
        }
    }

    #[inline]
    fn cow_owned<'b>(self) -> MyCow<'b, Self::OwnedVariant>
    where
        A: 'b,
        B: 'b,
    {
        match self {
            EitherDecomp::Owned(a) => A::cow_owned(a),
            EitherDecomp::Borrowed(b) => B::cow_owned(b),
        }
    }

    #[inline]
    fn ref_cow_owned<'b>(&'b self) -> MyCow<'b, Self::OwnedVariant>
    where
        Self: 'b,
    {
        match self {
            EitherDecomp::Owned(a) => A::ref_cow_owned(a),
            EitherDecomp::Borrowed(b) => B::ref_cow_owned(b),
        }
    }
}

#[macro_export]
macro_rules! self_ownable {
    ($type:ty where $($qual:tt)*) => {
        impl<$($qual)*> $crate::instance::Ownable for $type {
            type OwnedVariant = Self;

            #[inline]
            fn into_owned(self) -> Self::OwnedVariant {
                self
            }

            fn clone_owned(&self) -> Self::OwnedVariant {
                self.clone()
            }

            fn cow_owned<'b>(self) -> $crate::instance::MyCow<'b, Self::OwnedVariant>
            where
                Self: 'b,
            {
                $crate::instance::MyCow::Owned(self)
            }

            fn ref_cow_owned<'b>(&'b self) -> $crate::instance::MyCow<'b, Self::OwnedVariant>
            where
                Self: 'b,
            {
                $crate::instance::MyCow::Borrowed(self)
            }
        }
    };
}

self_ownable!(Vec<T> where T : Clone);

impl<'a, T: Clone> Ownable for &'a [T] {
    type OwnedVariant = Vec<T>;

    #[inline]
    fn into_owned(self) -> Self::OwnedVariant {
        Vec::from(self)
    }

    #[inline]
    fn clone_owned(&self) -> Self::OwnedVariant {
        Vec::from(*self)
    }

    #[inline]
    fn cow_owned<'b>(self) -> MyCow<'b, Self::OwnedVariant>
    where
        Self: 'b,
    {
        MyCow::Owned(Vec::from(self))
    }

    fn ref_cow_owned<'b>(&'b self) -> MyCow<'b, Self::OwnedVariant>
    where
        Self: 'b,
    {
        MyCow::Owned(Vec::from(*self))
    }
}

/// Trait for abitrary mathematical spaces.
pub trait Space: Ownable<OwnedVariant = Self::Principal> + Sized {
    /// Principal, typically owned realisation of the space.
    type Principal: ClosedSpace;

    /// Default decomposition for the space
    type Decomp: Decomposition<Self>;
}

mod private {
    pub trait Sealed {}
}

/// Helper trait for working with own types.
pub trait Owned: Ownable<OwnedVariant = Self> + private::Sealed {}
impl<X: Ownable<OwnedVariant = X>> private::Sealed for X {}
impl<X: Ownable<OwnedVariant = X>> Owned for X {}

/// Helper trait for working with closed spaces, operations in which should
/// return members of the same space
pub trait ClosedSpace: Space<Principal = Self> + Owned + Instance<Self> {}
impl<X: Space<Principal = Self> + Owned + Instance<Self>> ClosedSpace for X {}

#[macro_export]
macro_rules! impl_basic_space {
    ($($type:ty)*) => {
        $( $crate::impl_basic_space!($type where ); )*
    };
    ($type:ty where $($where:tt)*) => {
        impl<$($where)*> $crate::instance::Space for $type {
            type Principal = Self;
            type Decomp = $crate::instance::BasicDecomposition;
        }

        impl<$($where)*> $crate::instance::Ownable for $type {
            type OwnedVariant = Self;

            #[inline]
            fn into_owned(self) -> Self::OwnedVariant {
                self
            }

            #[inline]
            fn clone_owned(&self) -> Self::OwnedVariant {
                *self
            }

            #[inline]
            fn cow_owned<'b>(self) -> MyCow<'b, Self::OwnedVariant> where Self : 'b {
                MyCow::Owned(self)
            }

            #[inline]
            fn ref_cow_owned<'b>(&self) -> MyCow<'b, Self::OwnedVariant> where Self : 'b {
                MyCow::Owned(*self)
            }
        }
    };
}

impl_basic_space!(u8 u16 u32 u64 u128 usize
                  i8 i16 i32 i64 i128 isize
                  f32 f64);

/// Marker type for decompositions to be used with [`Instance`].
pub trait Decomposition<X: Space>: Sized {
    /// Possibly owned form of the decomposition
    type Decomposition<'b>: Instance<X, Self>
    where
        X: 'b;
    /// Unlikely owned form of the decomposition.
    /// Type for a lightweight intermediate conversion that does not own the original variable.
    /// Usually this is just a reference, but may also be a lightweight structure that
    /// contains references; see the implementation for [`crate::direct_product::Pair`].
    type Reference<'b>: Instance<X, Self> + Copy
    where
        X: 'b;

    /// Lift the lightweight reference type into a full decomposition type.
    fn lift<'b>(r: Self::Reference<'b>) -> Self::Decomposition<'b>;
}

/// Most common [`Decomposition`] (into `Either<X, &'b X>`) that allows working with owned
/// values and all sorts of references.
#[derive(Copy, Clone, Debug)]
pub struct BasicDecomposition;

impl<X: Space> Decomposition<X> for BasicDecomposition {
    type Decomposition<'b>
        = MyCow<'b, X>
    where
        X: 'b;
    type Reference<'b>
        = &'b X
    where
        X: 'b;

    #[inline]
    fn lift<'b>(r: Self::Reference<'b>) -> Self::Decomposition<'b> {
        MyCow::Borrowed(r)
    }
}

/// Helper trait for functions to work with either owned values or references to either the
/// “principal type” `X` or types some present a subset of `X`. In the latter sense, this
/// generalises [`std::borrow::ToOwned`], [`std::borrow::Borrow`], and [`std::borrow::Cow`].
///
/// This is used, for example, by [`crate::mapping::Mapping::apply`].
pub trait Instance<X, D = <X as Space>::Decomp>: Sized
where
    X: Space,
    D: Decomposition<X>,
{
    /// Decomposes self according to `decomposer`, and evaluate `f` on the result.
    /// Consumes self.
    #[inline]
    fn eval_decompose<'b, R>(self, f: impl FnOnce(D::Decomposition<'b>) -> R) -> R
    where
        X: 'b,
        Self: 'b,
    {
        f(self.decompose())
    }

    /// Does a light decomposition of self `decomposer`, and evaluates `f` on the result.
    /// Does not consume self.
    fn eval_ref<'b, R>(&'b self, f: impl FnOnce(D::Reference<'b>) -> R) -> R
    where
        X: 'b,
        Self: 'b;

    /// Returns an owned instance of `X`, cloning or converting non-true instances when necessary.
    fn own(self) -> X::Principal;

    fn decompose<'b>(self) -> D::Decomposition<'b>
    where
        Self: 'b;

    /// Returns an owned instance or reference to `X`, converting non-true instances when necessary.
    ///
    /// Default implementation uses [`Self::own`]. Consumes the input.
    fn cow<'b>(self) -> MyCow<'b, X::Principal>
    where
        Self: 'b;

    #[inline]
    /// Evaluates `f` on a reference to self.
    ///
    /// Default implementation uses [`Self::cow`]. Consumes the input.
    fn eval<'b, R>(self, f: impl FnOnce(&X::Principal) -> R) -> R
    where
        X: 'b,
        Self: 'b,
    {
        f(&*self.cow())
    }
}

impl<X: Space> Instance<X, BasicDecomposition> for X {
    #[inline]
    fn eval_ref<'b, R>(&'b self, f: impl FnOnce(&'b X) -> R) -> R
    where
        X: 'b,
        Self: 'b,
    {
        f(self)
    }

    #[inline]
    fn own(self) -> X::Principal {
        self.into_owned()
    }

    #[inline]
    fn cow<'b>(self) -> MyCow<'b, X::Principal>
    where
        Self: 'b,
    {
        self.cow_owned()
    }

    #[inline]
    fn decompose<'b>(self) -> MyCow<'b, X>
    where
        Self: 'b,
    {
        MyCow::Owned(self)
    }
}

impl<'a, X: Space> Instance<X, BasicDecomposition> for &'a X {
    #[inline]
    fn eval_ref<'b, R>(&'b self, f: impl FnOnce(&'b X) -> R) -> R
    where
        X: 'b,
        Self: 'b,
    {
        f(*self)
    }

    #[inline]
    fn own(self) -> X::Principal {
        self.into_owned()
    }

    #[inline]
    fn cow<'b>(self) -> MyCow<'b, X::Principal>
    where
        Self: 'b,
    {
        self.cow_owned()
    }

    #[inline]
    fn decompose<'b>(self) -> MyCow<'b, X>
    where
        Self: 'b,
    {
        MyCow::Borrowed(self)
    }
}

impl<'a, X: Space> Instance<X, BasicDecomposition> for &'a mut X {
    #[inline]
    fn eval_ref<'b, R>(&'b self, f: impl FnOnce(&'b X) -> R) -> R
    where
        X: 'b,
        Self: 'b,
    {
        f(*self)
    }

    #[inline]
    fn own(self) -> X::Principal {
        self.into_owned()
    }

    #[inline]
    fn cow<'b>(self) -> MyCow<'b, X::Principal>
    where
        Self: 'b,
    {
        self.cow_owned()
    }

    #[inline]
    fn decompose<'b>(self) -> MyCow<'b, X>
    where
        Self: 'b,
    {
        MyCow::Borrowed(self)
    }
}

impl<'a, X: Space> Instance<X, BasicDecomposition> for MyCow<'a, X> {
    #[inline]
    fn eval_ref<'b, R>(&'b self, f: impl FnOnce(&'b X) -> R) -> R
    where
        X: 'b,
        Self: 'b,
    {
        match self {
            MyCow::Borrowed(a) => f(a),
            MyCow::Owned(b) => f(&b),
        }
    }

    #[inline]
    fn own(self) -> X::Principal {
        self.into_owned()
    }

    #[inline]
    fn cow<'b>(self) -> MyCow<'b, X::Principal>
    where
        Self: 'b,
    {
        self.cow_owned()
    }

    #[inline]
    fn decompose<'b>(self) -> MyCow<'b, X>
    where
        Self: 'b,
    {
        self
    }
}

/// Marker type for mutable decompositions to be used with [`InstanceMut`].
pub trait DecompositionMut<X: Space>: Sized {
    type ReferenceMut<'b>: InstanceMut<X, Self>
    where
        X: 'b;
}

/// Helper trait for functions to work with mutable references.
pub trait InstanceMut<X: Space, D = <X as Space>::Decomp>: Sized
where
    D: DecompositionMut<X>,
{
    /// Returns a mutable decomposition of self.
    fn ref_instance_mut(&mut self) -> D::ReferenceMut<'_>;
}

impl<X: Space> DecompositionMut<X> for BasicDecomposition {
    type ReferenceMut<'b>
        = &'b mut X
    where
        X: 'b;
}

/// This impl may seem pointless, but allows throwaway mutable scratch variables
impl<'a, X: Space> InstanceMut<X, BasicDecomposition> for X {
    #[inline]
    fn ref_instance_mut(
        &mut self,
    ) -> <BasicDecomposition as DecompositionMut<X>>::ReferenceMut<'_> {
        self
    }
}

impl<'a, X: Space> InstanceMut<X, BasicDecomposition> for &'a mut X {
    #[inline]
    fn ref_instance_mut(
        &mut self,
    ) -> <BasicDecomposition as DecompositionMut<X>>::ReferenceMut<'_> {
        self
    }
}
