/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 */
package matrix4cj

/**
<P>
   For an m-by-n matrix A with m >= n, the QR decomposition is an m-by-n
   orthogonal matrix Q and an n-by-n upper triangular matrix R so that
   A = Q*R.
<P>
   The QR decompostion always exists, even if the matrix does not have
   full rank, so the constructor will never fail.  The primary use of the
   QR decomposition is in the least squares solution of nonsquare systems
   of simultaneous linear equations.  This will fail if isFullRank()
   returns false.
 */
public class QRDecomposition {
    /** Array for internal storage of decomposition*/
    private var QR: Matrix

    /**column dimension*/
    private var m: Int64 = 0

    /**row dimension*/
    private var n: Int64 = 0

    /** Array for internal storage of diagonal of R*/
    private var Rdiag: Array<Float64>

    /** 
     * QR Decomposition, computed by Householder reflections.
     * Structure to access R and the Householder vectors and compute Q.
     * @param A    Rectangular matrix
     */
    public init(A: Matrix) {
        QR = A.clone()
        m = A.rowNum
        n = A.colNum
        Rdiag = Array<Float64>(n) {_ => 0.0}
        for (k in 0..n) {
            var nrm = 0.0
            for (i in k..m) {
                nrm = Maths.hypot(nrm, QR[i, k])
            }
            if (nrm != 0.0) {
                if (QR[k, k] < 0.0) {
                    nrm = -nrm
                }
                for (i in k..m) {
                    QR[i, k] /= nrm
                }
                QR[k, k] += 1.0
                for (j in k + 1..n) {
                    var s = 0.0
                    for (i in k..m) {
                        s += QR[i, k] * QR[i, j]
                    }
                    s = -s / QR[k, k]
                    for (i in k..m) {
                        QR[i, j] += s * QR[i, k]
                    }
                }
            }
            Rdiag[k] = -nrm
        }
    }

    /** 
     * Is the matrix full rank?
     * @return     true if R, and hence A, has full rank.
     */
    public func isFullRank(): Bool {
        for (j in 0..n) {
            if (Rdiag[j] == 0.0) {
                return false
            }
        }
        return true
    }

    /** 
     * Return the Householder vectors
     * @return  Lower trapezoidal matrix whose columns define the reflections
     */
    public func getH(): Matrix {
        let X = Matrix(m, n)
        for (i in 0..m) {
            for (j in 0..n) {
                if (i >= j) {
                    X[i, j] = QR[i, j]
                } else {
                    X[i, j] = 0.0
                }
            }
        }
        return X
    }

    /** 
     * Return the upper triangular factor
     * @return     R
     */
    public func getR(): Matrix {
        let X = Matrix(n, n)
        for (i in 0..n) {
            for (j in 0..n) {
                if (i < j) {
                    X[i, j] = QR[i, j]
                } else if (i == j) {
                    X[i, j] = Rdiag[i]
                } else {
                    X[i, j] = 0.0
                }
            }
        }
        return X
    }

    /** 
     * Generate and return the (economy-sized) orthogonal factor
     * @return     Q
     */
    public func getQ(): Matrix {
        let X: Matrix = Matrix(m, n)
        for (k in (n - 1)..=0 : -1) {
            for (i in 0..m) {
                X[i, k] = 0.0
            }
            X[k, k] = 1.0
            for (j in k..n) {
                if (QR[k, k] != 0.0) {
                    var s: Float64 = 0.0
                    for (i in k..m) {
                        s += QR[i, k] * X[i, j]
                    }
                    s = -s / QR[k, k]
                    for (i in k..m) {
                        X[i, j] += s * QR[i, k]
                    }
                }
            }
        }
        return X
    }

    /** 
     * Least squares solution of A*X = B
     * @param B    A Matrix with as many rows as A and any number of columns.
     * @return     X that minimizes the two norm of Q*R*X-B.
     * @exception  IllegalArgumentException  Matrix row dimensions must agree.
     * @exception  RuntimeException  Matrix is rank deficient.
     */
    public func solve(B: Matrix): Matrix {
        if (B.rowNum != m) {
            throw IllegalArgumentException("Matrix row dimensions must agree.")
        }
        if (!this.isFullRank()) {
            throw Matrix4cjException("Matrix is rank deficient.")
        }
        let nx = B.colNum
        let X = B.clone()
        for (k in 0..n) {
            for (j in 0..nx) {
                var s: Float64 = 0.0
                for (i in k..m) {
                    s += QR[i, k] * X[i, j]
                }
                s = -s / QR[k, k]
                for (i in k..m) {
                    X[i, j] += s * QR[i, k]
                }
            }
        }
        for (k in (n - 1)..=0 : -1) {
            for (j in 0..nx) {
                X[k, j] /= Rdiag[k]
            }
            for (i in 0..k) {
                for (j in 0..nx) {
                    X[i, j] -= X[k, j] * QR[i, k]
                }
            }
        }
        return X[0..n, 0..nx] 
    }
}