use super::Matrix;
use super::Pivot;
use super::SingularMatrixError;
use super::SizeError;
use num_traits::One;
use num_traits::Zero;
use std::ops::Add;
use std::ops::Div;
use std::ops::DivAssign;
use std::ops::Index;
use std::ops::IndexMut;
use std::ops::Mul;
use std::ops::MulAssign;
use std::ops::Neg;
use std::ops::SubAssign;

pub trait MatrixImpl<'a, TYPE>: Matrix<TYPE>
where
    TYPE: 'a,
    Self: Index<(usize, usize), Output = TYPE>
        + IndexMut<(usize, usize), Output = TYPE>
        + IntoIterator<Item = (usize, usize, TYPE), IntoIter = Self::IteratorType>
        + 'a,
    &'a Self: IntoIterator<Item = (usize, usize, &'a TYPE), IntoIter = Self::RefIteratorType<'a>>,
    &'a mut Self:
        IntoIterator<Item = (usize, usize, &'a mut TYPE), IntoIter = Self::MutRefIteratorType<'a>>,
    Self::IteratorType: Iterator<Item = (usize, usize, TYPE)>,
    Self::RefIteratorType<'a>: Iterator<Item = (usize, usize, &'a TYPE)>,
    Self::MutRefIteratorType<'a>: Iterator<Item = (usize, usize, &'a mut TYPE)>,
{
    fn swap(&mut self, idx1: (usize, usize), idx2: (usize, usize)) {
        unsafe {
            // Can't take two mutable loans from one matrix, so instead just cast
            // them to their raw pointers to do the swap
            let pa: *mut TYPE = &mut self[idx1];
            let pb: *mut TYPE = &mut self[idx2];
            std::ptr::swap(pa, pb);
        }
    }

    fn transpose_in_place(&mut self) -> Result<(), SizeError>
    where
        Self: Sized,
    {
        for i in 0..self.rows() {
            for j in 0..self.columns() {
                self.swap((i, j), (j, i));
            }
        }
        Ok(())
    }

    fn swap_rows(&mut self, row1: usize, row2: usize) {
        for i in 0..self.columns() {
            self.swap((row1, i), (row2, i))
        }
    }

    fn swap_rows_pivot<P: Pivot>(&mut self, pivot: &mut P, row1: usize, row2: usize) {
        self.swap_rows(row1, row2);
        pivot.swap(row1, row2);
    }

    fn swap_rows_mul_pivot<P: Pivot>(
        &mut self,
        other: &mut Self,
        pivot: &mut P,
        row1: usize,
        row2: usize,
    ) {
        self.swap_rows_pivot(pivot, row1, row2);
        other.swap_rows(row1, row2);
    }

    fn swap_rows_mul(&mut self, other: &mut Self, row1: usize, row2: usize) {
        self.swap_rows(row1, row2);
        other.swap_rows(row1, row2);
    }

    fn swap_columns(&mut self, col1: usize, col2: usize) {
        for i in 0..self.rows() {
            self.swap((i, col1), (i, col2))
        }
    }

    type IteratorType;
    type RefIteratorType<'b>;
    type MutRefIteratorType<'b>;
}

pub trait LUDecompositionImpl<TYPE, MATRIX, PIVOT: Pivot>
where
    MATRIX: Matrix<TYPE>,
    for<'b> TYPE: One + MulAssign<&'b TYPE> + Neg<Output = TYPE>,
{
    fn get_matrix(&self) -> &MATRIX;

    fn get_pivot(&self) -> &PIVOT;

    fn det(&self) -> TYPE {
        let mut result = TYPE::one();
        for i in 0..self.get_matrix().rows() {
            result *= &self.get_matrix()[(i, i)];
        }
        if self.get_pivot().permutations() % 2 != 0 {
            result = -result;
        }
        result
    }
}

pub trait LUDecomposition<TYPE, MATRIX: Matrix<TYPE>, PIVOT: Pivot>:
    LUDecompositionImpl<TYPE, MATRIX, PIVOT>
where
    for<'b> TYPE: One + MulAssign<&'b TYPE> + Neg<Output = TYPE>,
{
    fn l(&self) -> MATRIX;

    fn u(&self) -> MATRIX;

    fn invert(&self) -> MATRIX;

    fn pivot(&self, m: MATRIX) -> MATRIX;

    fn det(&self) -> TYPE {
        LUDecompositionImpl::det(self)
    }
}

pub trait NumericalMatrixImpl<'a, TYPE, PIVOT: Pivot, LU: LUDecomposition<TYPE, Self, PIVOT>>:
    MatrixImpl<'a, TYPE>
where
    for<'b> TYPE: Zero
        + One
        + PartialEq
        + Neg<Output = TYPE>
        + Div<TYPE, Output = TYPE>
        + MulAssign<&'b TYPE>,
    for<'c> &'c TYPE: Mul<&'c TYPE, Output = TYPE>
        + Add<&'c TYPE, Output = TYPE>
        + Div<&'c TYPE, Output = TYPE>
        + Neg<Output = TYPE>,
    Self: Sized,
    TYPE: 'a,
    Self: Index<(usize, usize), Output = TYPE>
        + IndexMut<(usize, usize), Output = TYPE>
        + IntoIterator<Item = (usize, usize, TYPE), IntoIter = Self::IteratorType>
        + 'a,
    &'a Self: IntoIterator<Item = (usize, usize, &'a TYPE), IntoIter = Self::RefIteratorType<'a>>,
    &'a mut Self:
        IntoIterator<Item = (usize, usize, &'a mut TYPE), IntoIter = Self::MutRefIteratorType<'a>>,
    Self::IteratorType: Iterator<Item = (usize, usize, TYPE)>,
    Self::RefIteratorType<'a>: Iterator<Item = (usize, usize, &'a TYPE)>,
    Self::MutRefIteratorType<'a>: Iterator<Item = (usize, usize, &'a mut TYPE)>,
{
    fn identity(size: usize) -> Self;

    fn add_row(&mut self, source_index: usize, dest_index: usize, factor: &TYPE)
    where
        TYPE: Mul<Output = TYPE> + Add<Output = TYPE>,
        for<'f> &'f TYPE:
            Mul<Output = TYPE> + Add<Output = TYPE> + Div<Output = TYPE> + Neg<Output = TYPE>,
    {
        for i in 0..self.columns() {
            let tmp = &(self[(source_index, i)]) * factor;
            self[(dest_index, i)] = &(self[(dest_index, i)]) + &tmp;
        }
    }

    fn add_row_mul(
        &mut self,
        other: &mut Self,
        source_index: usize,
        dest_index: usize,
        factor: &TYPE,
    ) where
        TYPE: Add<Output = TYPE> + Mul<Output = TYPE>,
        for<'f> &'f TYPE:
            Mul<Output = TYPE> + Add<Output = TYPE> + Div<Output = TYPE> + Neg<Output = TYPE>,
    {
        self.add_row(source_index, dest_index, factor);
        other.add_row(source_index, dest_index, factor);
    }

    fn gauss_jordan_low<P: Pivot>(&mut self, pivot: &mut P)
    where
        TYPE: Zero
            + PartialEq
            + Add<Output = TYPE>
            + Mul<Output = TYPE>
            + Neg<Output = TYPE>
            + Div<Output = TYPE>,
        for<'f> &'f TYPE:
            Mul<Output = TYPE> + Add<Output = TYPE> + Div<Output = TYPE> + Neg<Output = TYPE>,
    {
        for i in 0..self.rows() {
            if self[(i, i)] == TYPE::zero() {
                for j in (i + 1)..self.columns() {
                    if self[(j, i)] != TYPE::zero() {
                        self.swap_rows_pivot(pivot, i, j);
                        break;
                    }
                }
            }
            for j in (i + 1)..self.rows() {
                if self[(i, i)] != TYPE::zero() {
                    let neg = -(&self[(j, i)]);
                    let factor = &neg / &(self[(i, i)]);
                    self.add_row(i, j, &factor)
                }
            }
        }
    }

    fn gauss_jordan_low_mul<P: Pivot>(&mut self, other: &mut Self, pivot: &mut P)
    where
        TYPE: Zero
            + PartialEq
            + Add<Output = TYPE>
            + Mul<Output = TYPE>
            + Neg<Output = TYPE>
            + Div<Output = TYPE>,
        for<'f> &'f TYPE:
            Mul<Output = TYPE> + Add<Output = TYPE> + Div<Output = TYPE> + Neg<Output = TYPE>,
    {
        for i in 0..self.rows() {
            if self[(i, i)] == TYPE::zero() {
                for j in (i + 1)..self.columns() {
                    if self[(j, i)] != TYPE::zero() {
                        self.swap_rows_mul_pivot(other, pivot, i, j);
                        break;
                    }
                }
            }
            for j in (i + 1)..self.rows() {
                if self[(i, i)] != TYPE::zero() {
                    let neg = -(&self[(j, i)]);
                    let factor = &neg / &(self[(i, i)]);
                    self.add_row_mul(other, i, j, &factor)
                }
            }
        }
    }

    fn gauss_jordan_high<P: Pivot>(&mut self, pivot: &mut P)
    where
        TYPE: Zero
            + PartialEq
            + Add<Output = TYPE>
            + Mul<Output = TYPE>
            + Neg<Output = TYPE>
            + Div<Output = TYPE>,
        for<'f> &'f TYPE:
            Mul<Output = TYPE> + Add<Output = TYPE> + Div<Output = TYPE> + Neg<Output = TYPE>,
    {
        for i in (usize::zero()..self.rows()).rev() {
            if self[(i, i)] == TYPE::zero() {
                for j in (usize::zero()..i).rev() {
                    if self[(j, i)] != TYPE::zero() {
                        self.swap_rows_pivot(pivot, i, j);
                        break;
                    }
                }
            }
            for j in (usize::zero()..i).rev() {
                if self[(i, i)] != TYPE::zero() {
                    let neg = -(&self[(j, i)]);
                    let factor = &neg / &(self[(i, i)]);
                    self.add_row(i, j, &factor)
                }
            }
        }
    }

    fn gauss_jordan_high_mul<P: Pivot>(&mut self, other: &mut Self, pivot: &mut P)
    where
        TYPE: Zero
            + PartialEq
            + Add<Output = TYPE>
            + Mul<Output = TYPE>
            + Neg<Output = TYPE>
            + Div<Output = TYPE>,
        for<'f> &'f TYPE:
            Mul<Output = TYPE> + Add<Output = TYPE> + Div<Output = TYPE> + Neg<Output = TYPE>,
    {
        for i in (usize::zero()..self.rows()).rev() {
            if self[(i, i)] == TYPE::zero() {
                for j in i..usize::zero() {
                    if self[(j, i)] != TYPE::zero() {
                        self.swap_rows_mul_pivot(other, pivot, i, j);
                        break;
                    }
                }
            }
            for j in (usize::zero()..i).rev() {
                if self[(i, i)] != TYPE::zero() {
                    let neg = -(&self[(j, i)]);
                    let factor = &neg / &(self[(i, i)]);
                    self.add_row_mul(other, i, j, &factor)
                }
            }
        }
    }

    fn triu_impl(&'a mut self, diag_replacement: Option<TYPE>)
    where
        TYPE: Zero + Clone,
    {
        match diag_replacement {
            Some(replacement) => {
                for (i, j, value) in self {
                    match &i.cmp(&j) {
                        std::cmp::Ordering::Equal => *value = replacement.clone(),
                        std::cmp::Ordering::Greater => *value = TYPE::zero(),
                        _ => {}
                    }
                }
            }
            None => {
                for (i, j, value) in self {
                    if i > j {
                        *value = TYPE::zero()
                    }
                }
            }
        }
    }

    fn tril_impl(&'a mut self, diag_replacement: Option<TYPE>)
    where
        TYPE: Zero + Clone,
    {
        match diag_replacement {
            Some(replacement) => {
                for (i, j, value) in self {
                    match i.cmp(&j) {
                        std::cmp::Ordering::Equal => *value = replacement.clone(),
                        std::cmp::Ordering::Less => *value = TYPE::zero(),
                        _ => {}
                    }
                }
            }
            None => {
                for (i, j, value) in self {
                    if i < j {
                        *value = TYPE::zero()
                    }
                }
            }
        }
    }

    fn invert(&mut self, pivot: &mut PIVOT, result: &mut Self)
    where
        TYPE: Clone + Zero + One + PartialEq + Neg<Output = TYPE> + Div<Output = TYPE>,
        for<'f> &'f TYPE:
            Mul<Output = TYPE> + Add<Output = TYPE> + Div<Output = TYPE> + Neg<Output = TYPE>,
    {
        self.gauss_jordan_low_mul(result, pivot);
        self.gauss_jordan_high_mul(result, pivot);
        for i in 0..result.rows() {
            let f = self[(i, i)].clone();
            for j in 0..result.columns() {
                result[(i, j)] = &(result[(i, j)]) / &f;
            }
        }
    }

    fn det(mut self, mut pivot: PIVOT) -> TYPE {
        self.gauss_jordan_low(&mut pivot);
        (0..self.rows())
            .map(|i| &self[(i, i)])
            .fold(TYPE::one(), |a, b| &a * b)
    }

    fn lu_pivot(&mut self, i: usize, pivot: &mut PIVOT)
    where
        TYPE: PartialOrd + Zero + Clone,
        for<'f> &'f TYPE: Neg<Output = TYPE>,
    {
        let mut max = abs(&self[(i, i)]);
        let mut max_index = i;
        for j in (i + 1)..self.rows() {
            if abs(&self[(j, i)]) > max {
                max = abs(&self[(i, j)]);
                max_index = j;
            }
        }
        if max_index != i {
            self.swap_rows_pivot(pivot, i, max_index)
        }
    }

    fn lu(&mut self) -> Result<(), SingularMatrixError>
    where
        TYPE: Zero + Clone,
        for<'b> &'b TYPE: Mul<&'b TYPE, Output = TYPE>,
        TYPE: SubAssign<TYPE> + DivAssign<TYPE>,
    {
        for i in 0..self.rows() {
            self.lu_row(i)?
        }
        Ok(())
    }

    fn lup(&mut self, p: &mut PIVOT) -> Result<(), SingularMatrixError>
    where
        TYPE: Zero + Clone + PartialOrd,
        for<'b> &'b TYPE: Mul<&'b TYPE, Output = TYPE> + Neg<Output = TYPE>,
        TYPE: SubAssign<TYPE> + DivAssign<TYPE>,
    {
        for i in 0..self.rows() {
            self.lu_pivot(i, p);
            self.lu_row(i)?
        }
        Ok(())
    }

    fn lu_invert(&self, result: &mut Self, p: &mut PIVOT)
    where
        TYPE: Zero + One + Clone + PartialOrd,
        for<'b> &'b TYPE: Mul<&'b TYPE, Output = TYPE>,
        for<'c> TYPE: DivAssign<&'c TYPE>,
        TYPE: SubAssign<TYPE>,
    {
        for i in 0..self.rows() {
            for j in 0..self.rows() {
                if p[j] == i {
                    result[(j, i)] = TYPE::one();
                } else {
                    result[(j, i)] = TYPE::zero();
                }
                for k in 0..j {
                    let delta = &self[(j, k)] * &result[(k, i)];
                    result[(j, i)] -= delta;
                }
            }
            for j in (0..self.rows()).rev() {
                for k in (j + 1)..self.rows() {
                    let delta = &self[(j, k)] * &result[(k, i)];
                    result[(j, i)] -= delta;
                }
                result[(j, i)] /= &self[(j, j)];
            }
        }
    }

    fn lu_row(&mut self, i: usize) -> Result<(), SingularMatrixError>
    where
        TYPE: Zero + Clone,
        for<'b> &'b TYPE: Mul<&'b TYPE, Output = TYPE>,
        TYPE: SubAssign<TYPE> + DivAssign<TYPE>,
    {
        if self[(i, i)].is_zero() {
            return Err(SingularMatrixError {});
        }
        for j in i..self.columns() {
            for k in 0..i {
                let delta = &self[(i, k)] * &self[(k, j)];
                self[(i, j)] -= delta;
            }
        }
        for j in (i + 1)..self.columns() {
            for k in 0..i {
                let delta = &self[(j, k)] * &self[(k, i)];
                self[(j, i)] -= delta;
            }
            let div = self[(i, i)].clone();
            self[(j, i)] /= div;
        }
        Ok(())
    }

    fn squared_norm2(&'a self) -> TYPE {
        self.into_iter().fold(TYPE::zero(), move |acc, (_, _, v)| { acc + v * v })
    }
}

fn abs<TYPE>(value: &TYPE) -> TYPE
where
    TYPE: PartialOrd + Zero + Clone,
    for<'a> &'a TYPE: Neg<Output = TYPE>,
{
    if value >= &TYPE::zero() {
        value.clone()
    } else {
        -value
    }
}