use std::fmt::Display;

use crate::ring::*;

mod owned;
mod submatrix;
mod transpose;

pub use owned::*;
pub use submatrix::*;
#[allow(unused_imports)]
pub use transpose::*;

/// Contains the trait [`transform::TransformTarget`], for "consumers" of elementary
/// matrix operations.
pub mod transform;

#[stability::unstable(feature = "enable")]
pub fn format_matrix<'a, M, R>(row_count: usize, col_count: usize, matrix: M, ring: R) -> impl 'a + Display
where
    R: 'a + RingStore,
    El<R>: 'a,
    M: 'a + Fn(usize, usize) -> &'a El<R>,
{
    struct DisplayWrapper<'a, R: 'a + RingStore, M: Fn(usize, usize) -> &'a El<R>> {
        matrix: M,
        ring: R,
        row_count: usize,
        col_count: usize,
    }

    impl<'a, R: 'a + RingStore, M: Fn(usize, usize) -> &'a El<R>> Display for DisplayWrapper<'a, R, M> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let strings = (0..self.row_count)
                .flat_map(|i| (0..self.col_count).map(move |j| (i, j)))
                .map(|(i, j)| format!("{}", self.ring.format((self.matrix)(i, j))))
                .collect::<Vec<_>>();
            let max_len = strings.iter().map(|s| s.chars().count()).chain([2]).max().unwrap();
            let mut strings = strings.into_iter();
            for i in 0..self.row_count {
                write!(f, "|")?;
                if self.col_count > 0 {
                    write!(f, "{:>width$}", strings.next().unwrap(), width = max_len)?;
                }
                for _ in 1..self.col_count {
                    write!(f, ",{:>width$}", strings.next().unwrap(), width = max_len)?;
                }
                if i + 1 != self.row_count {
                    writeln!(f, "|")?;
                } else {
                    write!(f, "|")?;
                }
            }
            return Ok(());
        }
    }

    DisplayWrapper {
        matrix,
        ring,
        col_count,
        row_count,
    }
}

/// Defines a very general way of comparing anything that can be interpreted as a
/// matrix in [`matrix_compare::MatrixCompare`]. Used solely in [`crate::assert_matrix_eq`],
/// with tests being the primary use case.
pub mod matrix_compare {
    use super::*;

    /// Used by [`crate::assert_matrix_eq`] to compare objects that are matrices in a
    /// very general sense.
    pub trait MatrixCompare<T> {
        /// Returns the number of rows of the matrix.
        fn row_count(&self) -> usize;

        /// Returns the number of columns of the matrix.
        fn col_count(&self) -> usize;

        /// Returns a reference to the element at the given position.
        fn at(&self, i: usize, j: usize) -> &T;
    }

    impl<T, const ROWS: usize, const COLS: usize> MatrixCompare<T> for [[T; COLS]; ROWS] {
        fn col_count(&self) -> usize { COLS }
        fn row_count(&self) -> usize { ROWS }
        fn at(&self, i: usize, j: usize) -> &T { &self[i][j] }
    }

    impl<T, const ROWS: usize, const COLS: usize> MatrixCompare<T> for [DerefArray<T, COLS>; ROWS] {
        fn col_count(&self) -> usize { COLS }
        fn row_count(&self) -> usize { ROWS }
        fn at(&self, i: usize, j: usize) -> &T { &self[i][j] }
    }

    impl<'a, V: AsPointerToSlice<T>, T> MatrixCompare<T> for Submatrix<'a, V, T> {
        fn col_count(&self) -> usize { Submatrix::col_count(self) }
        fn row_count(&self) -> usize { Submatrix::row_count(self) }
        fn at(&self, i: usize, j: usize) -> &T { Submatrix::at(self, i, j) }
    }

    impl<'a, V: AsPointerToSlice<T>, T> MatrixCompare<T> for SubmatrixMut<'a, V, T> {
        fn col_count(&self) -> usize { SubmatrixMut::col_count(self) }
        fn row_count(&self) -> usize { SubmatrixMut::row_count(self) }
        fn at(&self, i: usize, j: usize) -> &T { self.as_const().into_at(i, j) }
    }

    impl<'a, V: AsPointerToSlice<T>, T, const TRANSPOSED: bool> MatrixCompare<T>
        for TransposableSubmatrix<'a, V, T, TRANSPOSED>
    {
        fn col_count(&self) -> usize { TransposableSubmatrix::col_count(self) }
        fn row_count(&self) -> usize { TransposableSubmatrix::row_count(self) }
        fn at(&self, i: usize, j: usize) -> &T { TransposableSubmatrix::at(self, i, j) }
    }

    impl<'a, V: AsPointerToSlice<T>, T, const TRANSPOSED: bool> MatrixCompare<T>
        for TransposableSubmatrixMut<'a, V, T, TRANSPOSED>
    {
        fn col_count(&self) -> usize { TransposableSubmatrixMut::col_count(self) }
        fn row_count(&self) -> usize { TransposableSubmatrixMut::row_count(self) }
        fn at(&self, i: usize, j: usize) -> &T { self.as_const().into_at(i, j) }
    }

    impl<T> MatrixCompare<T> for OwnedMatrix<T> {
        fn col_count(&self) -> usize { OwnedMatrix::col_count(self) }
        fn row_count(&self) -> usize { OwnedMatrix::row_count(self) }
        fn at(&self, i: usize, j: usize) -> &T { OwnedMatrix::at(self, i, j) }
    }

    impl<T, M: MatrixCompare<T>> MatrixCompare<T> for &M {
        fn col_count(&self) -> usize { (**self).col_count() }
        fn row_count(&self) -> usize { (**self).row_count() }
        fn at(&self, i: usize, j: usize) -> &T { (**self).at(i, j) }
    }

    /// Checks whether two matrices are equal.
    ///
    /// The prime use case is [`crate::assert_matrix_eq!`], which should be applicable
    /// to all types that are, in some sense, interpretable as matrices.
    pub fn is_matrix_eq<R: ?Sized + RingBase, M1: MatrixCompare<R::Element>, M2: MatrixCompare<R::Element>>(
        ring: &R,
        lhs: &M1,
        rhs: &M2,
    ) -> bool {
        if lhs.row_count() != rhs.row_count() || lhs.col_count() != rhs.col_count() {
            return false;
        }
        for i in 0..lhs.row_count() {
            for j in 0..lhs.col_count() {
                if !ring.eq_el(lhs.at(i, j), rhs.at(i, j)) {
                    return false;
                }
            }
        }
        return true;
    }
}

/// Variant of `assert_eq!` for matrices elements, i.e. assert that two ring matrices are equal.
/// Frequently used in tests.
///
/// This takes any arguments that implement [`matrix_compare::MatrixCompare`], a very minimal trait
/// for matrices that is implemented for all types that are, in some sense, interpretable as
/// matrices.
///
/// # Example
/// ```rust
/// # use feanor_math::ring::*;
/// # use feanor_math::primitive_int::*;
/// # use feanor_math::matrix::*;
/// # use feanor_math::assert_matrix_eq;
/// let lhs = [[0, 0, 0], [0, 0, 0], [0, 0, 0]];
/// let rhs = OwnedMatrix::zero(3, 3, StaticRing::<i64>::RING);
/// assert_matrix_eq!(StaticRing::<i64>::RING, lhs, rhs);
///  ```
#[macro_export]
macro_rules! assert_matrix_eq {
    ($ring:expr,$lhs:expr,$rhs:expr) => {
        match (&$ring, &$lhs, &$rhs) {
            (ring_val, lhs_val, rhs_val) => {
                assert!(
                    $crate::matrix::matrix_compare::is_matrix_eq(ring_val.get_ring(), lhs_val, rhs_val),
                    "Assertion failed: Expected\n{}\nto be\n{}",
                    $crate::matrix::format_matrix(
                        <_ as $crate::matrix::matrix_compare::MatrixCompare<_>>::row_count(lhs_val),
                        <_ as $crate::matrix::matrix_compare::MatrixCompare<_>>::col_count(lhs_val),
                        |i, j| <_ as $crate::matrix::matrix_compare::MatrixCompare<_>>::at(lhs_val, i, j),
                        ring_val
                    ),
                    $crate::matrix::format_matrix(
                        <_ as $crate::matrix::matrix_compare::MatrixCompare<_>>::row_count(rhs_val),
                        <_ as $crate::matrix::matrix_compare::MatrixCompare<_>>::col_count(rhs_val),
                        |i, j| <_ as $crate::matrix::matrix_compare::MatrixCompare<_>>::at(rhs_val, i, j),
                        ring_val
                    )
                );
            }
        }
    };
}
