// July 2017, Mark Gates, Innovative Computing Laboratory, UTK

#include <stdio.h>

#include <exception>

#include "magma.h"
#include "magma_lapack.h"

int main( int argc, char** argv )
{
    #define  A(i_, j_) (A  + (i_) + (j_)*lda)
    #define AP(i_, j_) (AP + (i_) + (j_)*lda)
    
    // --------------------    
    // constants
    const int ione = 1;
    const double c_zero    = 0;
    const double c_one     = 1;
    const double c_neg_one = -1;
    
    // --------------------
    // variables
    double *A    = NULL;
    double *AP   = NULL;
    double *Q    = NULL;
    double *R    = NULL;
    double *tau  = NULL;
    double *work = NULL;
    int    *jpvt = NULL;
    
    double tmp = 0;
    int m = 20;
    int n = 10;
    if (argc > 2) {
        m = atoi( argv[1] );
        n = atoi( argv[2] );
    }
    assert( m >= n );  // assumption of tester
    int lda   = m;
    int minmn = n;
    int info  = 0;
    int lwork, size;
    
    // larnv options
    int idist = 1;  // uniform (0,1)
    int iseed[4] = { 0, 1, 2, 3 };

    int status = 0;

    // --------------------
    magma_init();
    
    try {
        magma_dmalloc_cpu( &A,    lda*n );
        magma_dmalloc_cpu( &AP,   lda*n );
        magma_dmalloc_cpu( &Q,    lda*n );
        magma_dmalloc_cpu( &R,    n*n );
        magma_dmalloc_cpu( &tau,  minmn );
        magma_imalloc_cpu( &jpvt, n     );
        if (A    == NULL) { throw std::exception(); }
        if (tau  == NULL) { throw std::exception(); }
        if (jpvt == NULL) { throw std::exception(); }
        
        // initialize A matrix; keep A for checking results and work in copy Q
        size = lda*n;
        lapackf77_dlarnv( &idist, iseed, &size, A );
        lapackf77_dlacpy( "general", &m, &n, A, &lda, Q, &lda );
        printf( "A = " );  magma_dprint( m, n, A, lda );
        
        // clear jpvt array
        for (int i = 0; i < n; ++i) {
            jpvt[i] = 0;
        }
        
        // query for workspace size
        magma_dgeqp3( m, n, Q, lda, jpvt, tau, &tmp, -1, &info );
        lwork = int( tmp );
        magma_dmalloc_cpu( &work, lwork );
        if (work == NULL) { throw std::exception(); }
        
        // factor A*P = Q*R
        // input:  A stored in Q
        // output: implicit Q and R stored in Q and tau, P stored in jpvt
        magma_dgeqp3( m, n, Q, lda, jpvt, tau, work, lwork, &info );
        if (info != 0) {
            printf( "magma_dgeqp3 failed: %s (%d)\n", magma_strerror(info), info );
            throw std::exception();
        }
        printf( "QP3 = "   ); magma_dprint( m, n, Q, lda );
        printf( "tau = "   ); magma_dprint( 1, minmn, tau, 1 );
        printf( "jpvt = [" );
        for (int i = 0; i < n; ++i) {
            printf( " %3d", jpvt[i] );
        }
        printf( "];\n" );
        
        // copy upper triangular n-by-n R
        lapackf77_dlaset( "lower", &n, &n, &c_zero, &c_zero, R, &n );
        lapackf77_dlacpy( "upper", &n, &n, Q, &lda, R, &n );
        printf( "R = " ); magma_dprint( n, n, R, n );
        
        // generate Q in place
        // magma_dorgqr requires extra dT array,
        // while magma_dorgqr2 has same API as LAPACK
        magma_dorgqr2( m, n, n, Q, lda, tau, &info );
        if (info != 0) {
            printf( "magma_dorgqr failed: %s (%d)\n", magma_strerror(info), info );
            throw std::exception();
        }
        printf( "Q = " ); magma_dprint( m, n, Q, lda );
        
        // permute A => AP
        for (int j = 0; j < n; ++j) {
            // jpvt is 1-based, so subtract 1
            blasf77_dcopy( &m, A(0, jpvt[j] - 1), &ione, AP(0, j), &ione );
        }
        printf( "AP = " ); magma_dprint( m, n, AP, lda );
        
        // compute error || AP - QR || // (m ||A||)
        blasf77_dgemm( "no", "no", &m, &n, &n, &c_neg_one, Q, &lda, R, &n, &c_one, AP, &lda );
        double Anorm = lapackf77_dlange( "fro", &m, &n, A,  &lda, work );
        double error = lapackf77_dlange( "fro", &m, &n, AP, &lda, work ) / (m * Anorm);
        printf( "%% Error || AP - QR || / (m ||A||) = %.2e\n", error );
        
        // compute error || I - Q^H Q || / m
        lapackf77_dlaset( "general", &n, &n, &c_zero, &c_one, AP, &lda );
        blasf77_dgemm( "conj", "no", &n, &n, &m, &c_neg_one, Q, &lda, Q, &lda, &c_one, AP, &lda );
        error = lapackf77_dlange( "fro", &n, &n, AP, &lda, work ) / m;
        printf( "%% Error || I - Q^H Q || / m       = %.2e\n", error );
        
        // Matlab code to verify results
        // Error that Matlab computes depends on number of digits printed;
        // by default MAGMA prints only 4 digits, so expect errors ~ 1e-04 or less.
        //
        // [m, n] = size( A );
        // AP2 = A( :, jpvt );
        // assert( all( all( AP == AP2 )));
        // I = eye( n, n );
        // fprintf( '|| I - Q^T Q || / m         = %.2e\n', norm( I - Q'*Q ) / (m) );
        // fprintf( '||  AP - QR  || / (m ||A||) = %.2e\n', norm( AP - Q*R ) / (m*norm(A)) );
    }
    catch (std::exception& e) {
        printf( "caught exception\n" );
        status = 1;
    }
    
    magma_free( A    );
    magma_free( AP   );
    magma_free( Q    );
    magma_free( R    );
    magma_free( tau  );
    magma_free( jpvt );
    magma_free( work );
    
    magma_finalize();
    return status;
}