//#![feature(const_generics, const_evaluatable_checked)]

use crate::dimutils;


/// Implement by specifying the const DIM to give your implementation a dimension.
pub trait TensorDimension {
    const DIM : usize;
    const ISSCALAR : bool = Self::DIM == 0;
    fn is_scalar(&self) -> bool {Self::ISSCALAR}
}



///  Expresses tensor-like objects with a size. Due to current compiler constraints, takes the dimension as an argument 
///  in addition to it being an associated type, to avoid internal compiler errors in use.
pub trait TensorSize<const DIM:usize> : TensorDimension {
    fn size(&self) -> [usize;DIM];

    fn inbounds(&self,index : [usize;DIM]) -> bool {
        index.iter().zip(self.size().iter()).all(|(i,s)| i < s)   
    }

    /// Returns an iterator over all valid indices.
    fn eachindex(&self) -> dimutils::EachIndex<DIM> {
        dimutils::ndim_iterator(self.size())
    }

    /// If this is false, you have an invalid implementation.
    const VALIDDIMS : bool = DIM == Self::DIM;
}

/// The main trait of this crate. Somehat like iterator, is effectively a lazy multidimensional array-like object
/// instead of a lazy sequence.
/// You can map over functions, zip them, take cartesian products, etc etc in a way that does
/// not erase structure, and produces another object with random access indexing unlike iterators.
pub trait Broadcastable<const DIM : usize> : TensorSize<DIM> + Sized {
    type Element;

    /// This should return none iff the default
    /// inbounds method returns false.
    fn bget(&self, index:[usize;DIM]) -> Option<Self::Element>;

    /// Takes indices modulo size. Useful for periodic shifts of indices.
    fn mod_bget(&self,index:[isize;DIM]) -> Self::Element {
        self.bget(dimutils::modular_index(index,self.size())).expect("Broken contract in Broadcast impl. Modular access out of bounds.")
    }

    /// Writes the contents of the Broadcastable to a receiver, such as an actual (multidimensional) array.
    fn feedto(&self,receiver :&mut impl BroadcastReceiver<DIM,Element=Self::Element>) -> Option<()>{
        receiver.receive(self)
    }

    /// Returns an object with the provided size. If NEWDIM > Self::DIM, the additional indices are ignored.
    /// All entries of size input must be either equal to the size of input, except for entries where the original size was 1.
    fn lazy_updim<const NEWDIM : usize>(&self, size : [usize;NEWDIM] ) -> LazyUpdim<Self,DIM,NEWDIM> 
    {
        assert!(Self::VALIDDIMS, "Invalid dimensions entered in Trait implementation."); // These asserts can be eliminated with more advanced
        assert!(NEWDIM >= DIM, "Updimmed tensor cannot have fewer indices than the initial one."); // const generic bounds on nightly. ( )
        LazyUpdim {size,reference:&self}
    }

    /// The main way to create multidimensional arrays with this crate, 
    /// but can resize any broadcastable. Resized broadcastables must have the
    /// same number of elements as the old one. Preserves iteration order.
    /// Reshapes a broadcastable by linearizing and delinearizing indices.
    fn reshaped<const NEWDIM : usize>(self,size: [usize;NEWDIM]) -> ReShaped<Self,NEWDIM,DIM> {
        ReShaped {underlying : self,index_multipliers:dimutils::linearization_coeffs(size),size}
    }

    /// Provides a more flexible but less safe way to reshape a broadcastable. 
    /// Useful for slicing and dicing arrays and broadcastables, taking diagonals, etc etc.
    fn mapindex<F:Fn([usize;M],[usize;M]) -> [usize;DIM],const M : usize>(&self,indexclosure:F,sizeclosure : impl Fn([usize;DIM]) -> [usize;M] ) 
    -> MapIndex<Self,F,DIM,M>
    {
        MapIndex {reference: self,indexclosure,size:sizeclosure(self.size())}
    }


    /// Elements of returned broadcastable are the output of the provided closure, 
    /// applied to the input broadcastable.
    fn bmap<T,F :Fn(Self::Element) -> T>(&self,foo : F) -> BMap<T,Self,F,DIM>{
        assert!(Self::VALIDDIMS, "Invalid dimensions entered in Trait implementation.");
        BMap {reference:self,closure : foo}
    }

    /// For inputs whose elements are references, clones them.
    fn bcloned(self) -> BCloned<Self,DIM> {
        BCloned {underlying:self}
    }
    
    /// Unique operation not found in Iterators. If the dimensions of the two broadcastable, zips them. 
    /// If not, will lazy_updim all inputs to the same size. As such, can also provide cartesian or tensor products depending on input.
    fn broadcast2<'a,'b,T : Broadcastable<N>,const N : usize >(&'a self,foo: &'b T) 
        -> Broadcast2<LazyUpdim<'a,Self,DIM,{dimutils::cmax(DIM,N)}>,LazyUpdim<'b,T,N,{dimutils::cmax(DIM,N)}>,{dimutils::cmax(DIM,N)}> 
    {
        let commondims : [usize;dimutils::cmax(DIM,N)] = dimutils::commondims(self.size(),foo.size()).expect("F error handling");
        Broadcast2 {first : self.lazy_updim(commondims),second : foo.lazy_updim(commondims)}
    }

    /// Takes a broadcastable and returns an iterator over its elements.
    fn bc_iter(&self) -> BroadcastIterator<Self,DIM> {
        BroadcastIterator {reference:self,state : self.eachindex()}
    }
}

/// Trait that emulates write access to reshaped arrays.
pub trait BroadcastReceiver<const DIM : usize> : TensorSize<DIM> + Sized {
    type Element;
    fn bget_mut<'a>(&'a mut self, index: [usize;DIM]) -> Option<&'a mut Self::Element>;

    fn receive(&mut self, broadcast: &impl Broadcastable<DIM,Element=Self::Element>) -> Option<()> {
        if self.size() != broadcast.size() {return None};
        for index in self.eachindex() {
            *(self.bget_mut(index)?) = broadcast.bget(index)?;
        }
        Some(())
    }
    // TODO: properly review the implementation of iter_mut, which is unsafe. This should need GATs to be safe.
    /// To be safe, the BroadcastReceiver must NOT alias bget_mut to the same underlying field for different inputs.
    /// You may implement a safe iter_mut method using this on any receiver where that invariant is preserved.
    /// If in doubt, use eachindex + bget_mut for a safe solution.
    unsafe fn bc_iter_mut<'a>(&'a mut self) -> BroadcastIterMut<'a,Self,DIM> {
        BroadcastIterMut {reference : self.into() ,state: self.eachindex(),_marker: core::marker::PhantomData}
    }
}


/// Return type of bc_iter.
pub struct BroadcastIterator<'a,T : Broadcastable<DIM>,const DIM : usize> {
    reference : &'a T,
    state : dimutils::EachIndex<DIM>,
}

impl<'a,T : Broadcastable<DIM>,const DIM : usize> Iterator for BroadcastIterator<'a,T,DIM> {
    type Item = T::Element;
    fn next(&mut self) -> Option<T::Element> {
        let index = self.state.next()?;
        self.reference.bget(index)
    }
    fn size_hint(&self) -> (usize,Option<usize>){
        self.state.size_hint()
    }
}

impl<'a,T : Broadcastable<DIM>,const DIM : usize> ExactSizeIterator for BroadcastIterator<'a,T,DIM> {}
impl<'a,T : Broadcastable<DIM>,const DIM : usize> core::iter::FusedIterator for BroadcastIterator<'a,T,DIM> {}

/// Return type of bc_iter_mut.
pub struct BroadcastIterMut<'a,T : BroadcastReceiver<DIM>,const DIM : usize> {
    reference : core::ptr::NonNull<T>,
    state :dimutils::EachIndex<DIM>,
    _marker: core::marker::PhantomData<&'a mut T>,
}

impl<'a,T : BroadcastReceiver<DIM>,const DIM : usize> Iterator for BroadcastIterMut<'a,T,DIM> {
    type Item = &'a mut T::Element;
    fn next(&mut self) -> Option< &'a mut T::Element> {
        let index = self.state.next()?;
        let castreference =  unsafe {self.reference.as_mut::<'a>()};
        castreference.bget_mut(index)
    }
    fn size_hint(&self) -> (usize,Option<usize>){
        self.state.size_hint()
    }
}


/// Return type of reshaped.
pub struct ReShaped<T:Broadcastable<M>,const N : usize,const M : usize> {
    underlying : T,
    index_multipliers : [usize;N],
    size : [usize;N]
}

impl<T:Broadcastable<M>,const N : usize,const M : usize> TensorDimension for ReShaped<T,N,M> {
    const DIM : usize = T::DIM;
}
impl<T:Broadcastable<M>,const N : usize,const M : usize> TensorSize<N> for ReShaped<T,N,M> {
    fn size(&self) -> [usize;N] {self.size}
}
impl<T:Broadcastable<M>,const N : usize,const M : usize> Broadcastable<N> for ReShaped<T,N,M> {
    type Element = T::Element;
    fn bget(&self,index:[usize;N])-> Option<T::Element> {
        if !self.inbounds(index) {return None}
        let linearindex = dimutils::linearize_index(index,self.index_multipliers);
        let innerindex = dimutils::delinearize_index(linearindex,self.underlying.size());
        self.underlying.bget(innerindex)
    }
}
impl<T:Broadcastable<M> + BroadcastReceiver<M>,const N : usize,const M : usize> BroadcastReceiver<N> for ReShaped<T,N,M> {
    type Element = <T as BroadcastReceiver<M> >::Element;
    fn bget_mut(&mut self, index:[usize;N]) -> Option<&mut Self::Element> {
        if !self.inbounds(index) {return None}
        let linearindex = dimutils::linearize_index(index,self.index_multipliers);
        let innerindex = dimutils::delinearize_index(linearindex,self.underlying.size());
        self.underlying.bget_mut(innerindex)
    }
}

impl<'a,T:Broadcastable<M>,const N : usize,const M : usize> IntoIterator for &'a ReShaped<T,N,M> {
    type Item = T::Element;
    type IntoIter = BroadcastIterator<'a,ReShaped<T,N,M>,N>;
    fn into_iter(self) -> Self::IntoIter  {
        self.bc_iter()
    }
}





/// Return type of bcloned.
pub struct BCloned<T:Broadcastable<N>,const N : usize> {
    underlying : T,
}

impl<'a,T:Broadcastable<N>,const N : usize> TensorDimension for BCloned<T,N> {
    const DIM : usize = T::DIM;
}
impl<'a,T:Broadcastable<N>,const N : usize> TensorSize<N> for BCloned<T,N> {
    fn size(&self) -> [usize;N] {self.underlying.size()}
}
impl<'b,T:Broadcastable<N,Element=&'b E>,E : 'b + Clone,const N : usize> Broadcastable<N> for BCloned<T,N> 
{
    type Element = E;
    fn bget(&self,index:[usize;N]) -> Option<Self::Element>  {
         self.underlying.bget(index).cloned()
    }
}
impl<'a,'b,T:Broadcastable<N,Element=&'b E>,E : Clone + 'a + 'b,const N : usize> IntoIterator for &'a BCloned<T,N> {
    type Item = E;
    type IntoIter = BroadcastIterator<'a,BCloned<T,N>,N>;
    fn into_iter(self) -> Self::IntoIter  {
        self.bc_iter()
    }
}


/// Return type of bmap.
pub struct MapIndex<'a,T:Broadcastable<N>, F : Fn([usize;M],[usize;M]) -> [usize;N] ,const N : usize,const M : usize> {
    reference : &'a T,
    indexclosure : F,
    size : [usize;M],
}

impl<'a,T:Broadcastable<N>, F : Fn([usize;M],[usize;M]) -> [usize;N] ,const N : usize,const M : usize> TensorDimension for MapIndex<'a,T,F,N,M> {
    const DIM : usize = M;
}
impl<'a,T:Broadcastable<N>, F : Fn([usize;M],[usize;M]) -> [usize;N] ,const N : usize,const M : usize> TensorSize<M> for MapIndex<'a,T,F,N,M> {
    fn size(&self) -> [usize;M] {self.size}
}
impl<'a,T:Broadcastable<N>, F : Fn([usize;M],[usize;M]) -> [usize;N] ,const N : usize,const M : usize> Broadcastable<M> for MapIndex<'a,T,F,N,M> {
    type Element = T::Element;
    fn bget(&self,index : [usize;M]) -> Option<T::Element> {
        if !self.inbounds(index) {return None};
        let size = self.size;
        let indexclosure = &self.indexclosure;
        let newindex : [usize;N] = indexclosure(index,size);
        self.reference.bget(newindex)
    }
}

impl<'b,'a,T:Broadcastable<N>, F : Fn([usize;M],[usize;M]) -> [usize;N] ,const N : usize,const M : usize> IntoIterator for &'b MapIndex<'a,T,F,N,M> {
    type Item = T::Element;
    type IntoIter = BroadcastIterator<'b,MapIndex<'a,T,F,N,M>,M>;
    fn into_iter(self) -> Self::IntoIter  {
        self.bc_iter()
    }
} 

impl<'a,T:Broadcastable<N>, F : Fn([usize;M],[usize;M]) -> [usize;N] ,const N : usize,const M : usize>  MapIndex<'a,T,F,N,M> {
    pub fn iter(&self) -> BroadcastIterator<MapIndex<'a,T,F,N,M>,M> {
        self.into_iter()
    }
}



/// Return type of broadcast2.
pub struct Broadcast2<A : Broadcastable<N>,B : Broadcastable<N>, const N : usize> {
    first: A,
    second: B
}

impl<A : Broadcastable<N>,B : Broadcastable<N>, const N : usize> TensorDimension for Broadcast2<A,B,N> {
    const DIM : usize = N;
}
impl<A : Broadcastable<N>,B : Broadcastable<N>, const N : usize> TensorSize<N> for Broadcast2<A,B,N> {
    fn size(&self) -> [usize;N] {
        self.first.size()
    }
}
impl<A : Broadcastable<N>,B : Broadcastable<N>, const N : usize> Broadcastable<N> for Broadcast2<A,B,N> {
    type Element = (A::Element,B::Element);
    fn bget(&self, index:[usize;N]) -> Option<Self::Element> {
        Some((self.first.bget(index)?,self.second.bget(index)?))
    }
}

impl<'b,A : Broadcastable<N>,B : Broadcastable<N>, const N : usize> IntoIterator for &'b Broadcast2<A,B,N> {
    type Item = (A::Element,B::Element);
    type IntoIter = BroadcastIterator<'b, Broadcast2<A,B,N> ,N>;
    fn into_iter(self) -> Self::IntoIter  {
        self.bc_iter()
    }
}
impl<A : Broadcastable<N>,B : Broadcastable<N>, const N : usize> Broadcast2<A,B,N> {
    pub fn iter(&self) -> BroadcastIterator<Broadcast2<A,B,N> ,N> {
        self.bc_iter()
    }
}


/// Jointly broadcasts an arbitrary number of Broadcastables of the same type together.
pub fn broadcast_k<'a,A : Broadcastable<N>, const K : usize, const N : usize>(bcables : [&'a A;K]) -> BroadcastK<LazyUpdim<'a,A,N,N>,N,K> {
    let mut size : [usize;N] = [1;N];
    for &broadcastable in bcables.iter() {
        size = dimutils::commondims_samedims(size,broadcastable.size()).expect("F error handling")
    }
    BroadcastK {factors: bcables.map(|x| x.lazy_updim(size)) }
}


/// Return type of broadcastK.
pub struct BroadcastK<A : Broadcastable<N>, const N : usize,const K : usize> {
    factors : [A;K]
}
impl<A : Broadcastable<N>, const N : usize,const K : usize> TensorDimension for BroadcastK<A,N,K> {
    const DIM : usize = N;
}
impl<A : Broadcastable<N>, const N : usize,const K : usize> TensorSize<N> for BroadcastK<A,N,K> {
    fn size(&self) -> [usize;N] {
        self.factors[1].size()
    }
}
impl<A : Broadcastable<N>, const N : usize,const K : usize> Broadcastable<N> for BroadcastK<A,N,K> {
    type Element = [A::Element;K];
    fn bget(&self, index:[usize;N]) -> Option<Self::Element> {
        if self.inbounds(index) {
        Some(self.factors.each_ref().map(|x|x.bget(index).expect("Factors of broadcast_k have invalid sizes!")))} else {None}
    }
}
impl<'b ,A : Broadcastable<N>, const N : usize,const K : usize> IntoIterator for &'b BroadcastK<A,N,K> {
    type Item = [A::Element;K];
    type IntoIter = BroadcastIterator<'b, BroadcastK<A,N,K> ,N>;
    fn into_iter(self) -> Self::IntoIter  {
        self.bc_iter()
    }
}
impl<'b ,A : Broadcastable<N>, const N : usize,const K : usize> BroadcastK<A,N,K> {
    pub fn iter(&self) -> BroadcastIterator<BroadcastK<A,N,K> ,N> {
        self.bc_iter()
    }
}


/// Return type of lazy_updim.
pub struct LazyUpdim<'a,T : Broadcastable<OLDDIM>,const OLDDIM : usize, const DIM : usize> {
    size : [usize;DIM],
    reference : &'a T
}

impl<'a,T : Broadcastable<OLDDIM>,const OLDDIM : usize,const DIM : usize> TensorDimension for LazyUpdim<'a,T,OLDDIM,DIM> {
    const DIM : usize = DIM;
}
impl<'a,T : Broadcastable<OLDDIM>,const OLDDIM : usize,const DIM : usize> TensorSize<DIM> for LazyUpdim<'a,T,OLDDIM,DIM> {
    fn size(&self) -> [usize;DIM] {self.size}
}
impl<'a,T : Broadcastable<OLDDIM>,const OLDDIM : usize,const DIM : usize>  Broadcastable<DIM>  for LazyUpdim<'a,T,OLDDIM,DIM> {
    type Element = T::Element;
    fn bget(&self,index:[usize;DIM]) -> Option<Self::Element> {
        assert!(DIM >= OLDDIM);
        if !self.inbounds(index) {return None}
        let size = self.size();
        let newindex : [usize;OLDDIM] = array_init::array_init(|i| if size[i] > 1 {index[i]} else {0});
        self.reference.bget(newindex)
    }
}

impl<'b,'a,T : Broadcastable<OLDDIM>,const OLDDIM : usize,const DIM : usize> IntoIterator for &'b LazyUpdim<'a,T,OLDDIM,DIM> {
    type Item = T::Element;
    type IntoIter = BroadcastIterator<'b,LazyUpdim<'a,T,OLDDIM,DIM>,DIM>;
    fn into_iter(self) -> Self::IntoIter  {
        self.bc_iter()
    }
}

impl<'a,T : Broadcastable<OLDDIM>,const OLDDIM : usize,const DIM : usize> LazyUpdim<'a,T,OLDDIM,DIM> {
    pub fn iter(&self) ->  BroadcastIterator<LazyUpdim<'a,T,OLDDIM,DIM>,DIM>  {
        self.bc_iter()
    }
}


/// Return type of bmap.
pub struct BMap<'a,R, T : Broadcastable<DIM>, F :  Fn(T::Element) -> R  , const DIM: usize> {
    reference : &'a T,
    closure : F
}

impl<'a,R, T : Broadcastable<DIM>, F :  Fn(T::Element) -> R  , const DIM: usize> TensorDimension for BMap<'a,R,T,F,DIM> {
    const DIM : usize = DIM;
}
impl<'a,R, T : Broadcastable<DIM>, F :  Fn(T::Element) -> R  , const DIM: usize> TensorSize<DIM> for BMap<'a,R,T,F,DIM> {
    fn size(&self) -> [usize;DIM] {self.reference.size()}
}
impl<'a,R, T : Broadcastable<DIM>, F :  Fn(T::Element) -> R  , const DIM: usize> Broadcastable<DIM> for BMap<'a,R,T,F,DIM> {
    type Element = R;
    fn bget(&self,index:[usize;DIM]) -> Option<Self::Element> {
        self.reference.bget(index).map(&self.closure)
    }
}

impl<'b,'a,R, T : Broadcastable<DIM>, F :  Fn(T::Element) -> R  , const DIM: usize> IntoIterator for &'b BMap<'a,R,T,F,DIM> {
    type Item = R;
    type IntoIter = BroadcastIterator<'b,BMap<'a,R,T,F,DIM>,DIM>;
    fn into_iter(self) -> Self::IntoIter  {
        self.bc_iter()
    }
}



/// Creates a lazy Broadcastable from an arbitrary function or closure.
fn broadcast_closure<T,F : Fn([usize;N]) -> T,const N : usize>(closure : F,size : [usize;N]) -> BroadcastClosure<T,F,N> {
    BroadcastClosure {closure,size}
}

/// Return type of broadcast_closure.
pub struct BroadcastClosure<T,F : Fn([usize;N]) -> T,const N : usize> {
    size : [usize;N],
    closure : F
}
impl<T,F : Fn([usize;N]) -> T,const N : usize> TensorDimension for BroadcastClosure<T,F,N> {
    const DIM : usize = N;
}
impl<T,F : Fn([usize;N]) -> T,const N : usize> TensorSize<N> for BroadcastClosure<T,F,N> {
    fn size(&self) -> [usize;N] {
        self.size
    }
}
impl<T,F : Fn([usize;N]) -> T,const N : usize> Broadcastable<N> for BroadcastClosure<T,F,N> {
    type Element = T;
    fn bget(&self,index : [usize;N]) -> Option<T> {
        let closure = &self.closure;
        Some(closure(index))
    }
}




impl <'a,T> TensorDimension for &'a[T] {
    const DIM : usize = 1;
}
impl <'a,T> TensorSize<1> for &'a [T] {
    fn size(&self) -> [usize;1] {[self.len()]}
}
impl<'a,T> Broadcastable<1> for &'a [T] {
    type Element = &'a T;
    fn bget(&self,[index] : [usize;1]) -> Option<&'a T> {
        self.get(index)
    }
}








#[cfg(feature = "std")] 
impl<T> TensorDimension for Vec<T> {
    const DIM : usize = 1;
}
#[cfg(feature = "std")] 
impl<T> TensorSize<1> for Vec<T> {
    fn size(&self) -> [usize;1] {[self.len()]}
}
#[cfg(feature = "std")] 
impl<'a,T> TensorDimension for &'a Vec<T> {
    const DIM : usize = 1;
}
#[cfg(feature = "std")] 
impl<'a, T> TensorSize<1> for &'a Vec<T> {
    fn size(&self) -> [usize;1] {[self.len()]}
}
#[cfg(feature = "std")] 
impl<'a,T> TensorDimension for &'a mut Vec<T> {
    const DIM : usize = 1;
}
#[cfg(feature = "std")]
impl<'a, T> TensorSize<1> for &'a mut Vec<T> {
    fn size(&self) -> [usize;1] {[self.len()]}
}
#[cfg(feature = "std")]
impl<'a,T> Broadcastable<1> for &'a Vec<T> {
    type Element = &'a T;
    fn bget(&self,index : [usize;1]) -> Option<&'a T> {
        self.get(index[0])
    }
}
#[cfg(feature = "std")]
impl<T: Copy> Broadcastable<1> for Vec<T> {
    type Element = T;
    fn bget(&self,index : [usize;1]) -> Option<T> {
        self.get(index[0]).cloned()
    }
}
#[cfg(feature = "std")]
impl<'a,T> BroadcastReceiver<1> for &'a mut Vec<T> {
    type Element = T;
    fn bget_mut(&mut self,[index] : [usize;1]) -> Option<&mut T> {
        self.get_mut(index)
    }
}