use mop_common::rayon::iter::plumbing::bridge;
use mop_common::rayon::iter::plumbing::Producer;
use mop_common::rayon::iter::plumbing::ProducerCallback;
use mop_common::rayon::iter::plumbing::{Consumer, UnindexedConsumer};
use mop_common::rayon::iter::IndexedParallelIterator;
use mop_common::rayon::iter::ParallelIterator;

use crate::matrix::dr_matrix::{DrMatrixRowIter, DrMatrixRowIterMut};
use crate::utils::{ParallelIteratorWrapper, ParallelProducerWrapper};

macro_rules! impl_par_iter {
    ($dr_matrix_row_iter:ident, $return_type:ty) => {
        impl<'a, T> ParallelIterator for ParallelIteratorWrapper<$dr_matrix_row_iter<'a, T>>
        where
            T: Send + Sync,
        {
            type Item = $return_type;
            fn drive_unindexed<C>(self, consumer: C) -> C::Result
            where
                C: UnindexedConsumer<Self::Item>,
            {
                bridge(self, consumer)
            }

            fn opt_len(&self) -> Option<usize> {
                Some(self.0.len())
            }
        }

        impl<'a, T> IndexedParallelIterator for ParallelIteratorWrapper<$dr_matrix_row_iter<'a, T>>
        where
            T: Send + Sync,
        {
            fn with_producer<Cb>(self, callback: Cb) -> Cb::Output
            where
                Cb: ProducerCallback<Self::Item>,
            {
                callback.callback(ParallelProducerWrapper(self.0))
            }

            fn len(&self) -> usize {
                ExactSizeIterator::len(&self.0)
            }

            fn drive<C>(self, consumer: C) -> C::Result
            where
                C: Consumer<Self::Item>,
            {
                bridge(self, consumer)
            }
        }

        impl<'a, T> IntoIterator for ParallelProducerWrapper<$dr_matrix_row_iter<'a, T>> {
            type IntoIter = $dr_matrix_row_iter<'a, T>;
            type Item = <Self::IntoIter as Iterator>::Item;

            fn into_iter(self) -> Self::IntoIter {
                self.0
            }
        }

        impl<'a, T> Producer for ParallelProducerWrapper<$dr_matrix_row_iter<'a, T>> {
            type IntoIter = $dr_matrix_row_iter<'a, T>;
            type Item = <Self::IntoIter as Iterator>::Item;

            fn into_iter(self) -> Self::IntoIter {
                self.0
            }

            fn split_at(self, i: usize) -> (Self, Self) {
                let (a, b) = self.0.split_at(i);
                (ParallelProducerWrapper(a), ParallelProducerWrapper(b))
            }
        }
    };
}

impl_par_iter!(DrMatrixRowIter, &'a [T]);
impl_par_iter!(DrMatrixRowIterMut, &'a mut [T]);
