use alga::general::Real;

use core::{DefaultAllocator, MatrixN, MatrixMN, Matrix, SquareMatrix};
use constraint::{ShapeConstraint, SameNumberOfRows};
use storage::{Storage, StorageMut};
use allocator::Allocator;
use dimension::{Dim, Dynamic, DimSub};

/// The cholesky decomposion of a symmetric-definite-positive matrix.
#[derive(Clone, Debug)]
pub struct Cholesky<N: Real, D: Dim>
    where DefaultAllocator: Allocator<N, D, D> {
    chol: MatrixN<N, D>
}

impl<N: Real, D: Dim> Copy for Cholesky<N, D>
    where DefaultAllocator: Allocator<N, D, D>,
          MatrixN<N, D>: Copy { }

impl<N: Real, D: DimSub<Dynamic>> Cholesky<N, D>
    where DefaultAllocator: Allocator<N, D, D> {
    /// Attempts to compute the sholesky decomposition of `matrix`.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The intput matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn new(mut matrix: MatrixN<N, D>) -> Option<Self> {
        assert!(matrix.is_square(), "The input matrix must be square.");

        let n = matrix.nrows();

        for j in 0 .. n {
            for k in 0 .. j {
                let factor = unsafe { -*matrix.get_unchecked(j, k) };

                let (mut col_j, col_k) = matrix.columns_range_pair_mut(j, k);
                let mut col_j = col_j.rows_range_mut(j ..);
                let col_k     = col_k.rows_range(j ..);

                col_j.axpy(factor, &col_k, N::one());
            }

            let diag = unsafe { *matrix.get_unchecked(j, j) };
            if diag > N::zero() {
                let denom = diag.sqrt();
                unsafe { *matrix.get_unchecked_mut(j, j) = denom; }

                let mut col = matrix.slice_range_mut(j + 1 .., j);
                col /= denom;
            }
            else {
                return None;
            }
        }

        Some(Cholesky { chol: matrix })
    }

    /// Retrieves the lower-triangular factor of the cholesky decomposition.
    pub fn unpack(mut self) -> MatrixN<N, D> {
        self.chol.fill_upper_triangle(N::zero(), 1);
        self.chol
    }

    /// Retrieves the lower-triangular factor of che cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// This is an allocation-less version of `self.l()`. The values of the strict upper-triangular
    /// part are garbage and should be ignored by further computations.
    pub fn unpack_dirty(self) -> MatrixN<N, D> {
        self.chol
    }

    /// Retrieves the lower-triangular factor of the cholesky decomposition.
    pub fn l(&self) -> MatrixN<N, D> {
        self.chol.lower_triangle()
    }

    /// Retrieves the lower-triangular factor of the cholesky decomposition, without zeroing-out
    /// its strict upper-triangular part.
    ///
    /// This is an allocation-less version of `self.l()`. The values of the strict upper-triangular
    /// part are garbage and should be ignored by further computations.
    pub fn l_dirty(&self) -> &MatrixN<N, D> {
        &self.chol
    }

    /// Solves the system `self * x = b` where `self` is the decomposed matrix and `x` the unknown.
    ///
    /// The result is stored on `b`.
    pub fn solve_mut<R2: Dim, C2: Dim, S2>(&self, b: &mut Matrix<N, R2, C2, S2>)
        where S2: StorageMut<N, R2, C2>,
              ShapeConstraint: SameNumberOfRows<R2, D> {
        self.chol.solve_lower_triangular_mut(b);
        self.chol.tr_solve_lower_triangular_mut(b);
    }

    /// Solves the system `self * x = b` where `self` is the decomposed matrix and `x` the unknown.
    ///
    /// The result is stored on `b`.
    pub fn solve<R2: Dim, C2: Dim, S2>(&self, b: &Matrix<N, R2, C2, S2>) -> MatrixMN<N, R2, C2>
        where S2: StorageMut<N, R2, C2>,
              DefaultAllocator: Allocator<N, R2, C2>,
              ShapeConstraint: SameNumberOfRows<R2, D> {
        let mut res = b.clone_owned();
        self.solve_mut(&mut res);
        res
    }

    /// Computes the inverse of the decomposed matrix.
    pub fn inverse(&self) -> MatrixN<N, D> {
        let shape = self.chol.data.shape();
        let mut res = MatrixN::identity_generic(shape.0, shape.1);

        self.solve_mut(&mut res);
        res
    }
}

impl<N: Real, D: DimSub<Dynamic>, S: Storage<N, D, D>> SquareMatrix<N, D, S>
    where DefaultAllocator: Allocator<N, D, D> {

    /// Attempts to compute the sholesky decomposition of this matrix.
    ///
    /// Returns `None` if the input matrix is not definite-positive. The intput matrix is assumed
    /// to be symmetric and only the lower-triangular part is read.
    pub fn cholesky(self) -> Option<Cholesky<N, D>> {
        Cholesky::new(self.into_owned())
    }
}