/** @file  psMatrix.c
 *
 *  @brief Provides functions for linear algebra operations on psImages and psVectors. 
 *
 *  Functions are provided to:
 *      Transpose a psImage
 *      Compute LUD
 *      Solve LUD
 *      Matrix inversion
 *      Calculate determinant
 *      Matrix addition
 *      Matrix subtraction
 *      Matrix multiplication
 *      Calculate Eigenvectors
 *      Convert matrix to vector
 *      Convert vector to matrix
 *
 *  These functions treat psImages as if they were matrices, therefore there is no psMatrix.
 *
 *  @author Ross Harman, MHPCC
 *   
 *  @version $Revision: 1.15 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2004-10-08 02:48:01 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 */

/******************************************************************************/
/*  INCLUDE FILES                                                             */
/******************************************************************************/
#include <string.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_permutation.h>
#include <gsl/gsl_eigen.h>

#include "psMemory.h"
#include "psError.h"
#include "psImage.h"
#include "psVector.h"
#include "psMatrix.h"

/** Preprocessor macro to generate error a NULL image */
#define PS_CHECK_NULL_VECTOR(NAME, RETURN)                                                          \
if (NAME == NULL || NAME->data.V == NULL) {                                                         \
    psError(__func__,"Invalid operation: %s or its data is NULL.", #NAME);                          \
    return RETURN;                                                                                  \
}

/** Preprocessor macro to create vector based on another */
#define PS_CHECK_ALLOC_VECTOR(NAME, SIZE, PS_TYPE)                                                  \
if(NAME == NULL) {                                                                                  \
    NAME = psVectorAlloc(SIZE, PS_TYPE);                                                            \
}

/** Preprocessor macro to generate error for zero length vector */
#define PS_CHECK_SIZE_VECTOR(NAME, RETURN)                                                          \
if (NAME->n < 1) {                                                                                  \
    psError(__func__,"Invalid operation: %s has zero n value.", #NAME);                             \
    return RETURN;                                                                                  \
}

/** Preprocessor macro to generate error a NULL image */
#define PS_CHECK_NULL_IMAGE(NAME, RETURN)                                                           \
if (NAME == NULL || NAME->data.V == NULL) {                                                         \
    psError(__func__,"Invalid operation: %s or its data is NULL.", #NAME);                          \
    return RETURN;                                                                                  \
}

/** Preprocessor macro to create image based on another */
#define PS_CHECK_ALLOC_IMAGE(NAME, NCOLS, NROWS, PS_TYPE)                                           \
if(NAME == NULL) {                                                                                  \
    NAME = psImageAlloc(NCOLS, NROWS, PS_TYPE);                                                     \
}

/** Preprocessor macro to generate error for zero length rows or columns */
#define PS_CHECK_SIZE_IMAGE(NAME, RETURN)                                                           \
if (NAME->numCols < 1 || NAME->numRows < 1) {                                                       \
    psError(__func__,"Invalid operation: %s has zero rows or columns (%dx%d).", #NAME,              \
            NAME->numCols, NAME->numRows);                                                          \
    return RETURN;                                                                                  \
}

/** Preprocessor macro to generate error for image dimensionality not set to PS_DIMEN_IMAGE */
#define PS_CHECK_DIMEN_AND_TYPE(NAME, PS_DIMEN, RETURN)                                             \
if (NAME->type.dimen != PS_DIMEN) {                                                                 \
    psError(__func__,"Invalid operation: %s incorrect dimensionality %d.", #NAME, PS_DIMEN);        \
    return RETURN;                                                                                  \
} else if(NAME->type.type != PS_TYPE_F64) {                                                         \
    psError(__func__, "Invalid operation: %s not PS_TYPE_F64.", #NAME);                             \
    return RETURN;                                                                                  \
}

/** Preprocessor macro to check that input is not equal to output */
#define PS_CHECK_POINTERS(NAME1, NAME2, RETURN)                                                     \
if (NAME1 == NAME2) {                                                                               \
    psError(__func__,"Invalid operation: Pointer to %s is same as %s.", #NAME1, #NAME2);            \
    return RETURN;                                                                                  \
}

/** Preprocessor macro to check that an image is square */
#define PS_CHECK_SQUARE(NAME, RETURN)                                                               \
if (NAME->numCols != NAME->numRows) {                                                               \
    psError(__func__,"Invalid operation: %s not square array.", #NAME);                             \
    return RETURN;                                                                                  \
}

/** Preprocessor macro to initalize a GSL matrix. */
#define PS_GSL_MATRIX_INITIALIZE(LHS_NAME, RHS_NAME)                                                \
LHS_NAME.size1 = numRows;                                                                           \
LHS_NAME.size2 = numCols;                                                                           \
LHS_NAME.tda   = numCols;                                                                           \
LHS_NAME.data  = RHS_NAME;

/*****************************************************************************/
/* FUNCTION IMPLEMENTATION - PUBLIC                                          */
/*****************************************************************************/

psImage* psMatrixLUD(psImage* outImage, psVector* outPerm, psImage* inImage)
{
    int signum = 0;
    int arraySize = 0;
    int numRows = 0;
    int numCols = 0;
    gsl_matrix lu;
    gsl_permutation perm;

    // Error checks
    PS_CHECK_POINTERS(inImage, outImage, outImage);
    PS_CHECK_NULL_IMAGE(inImage, outImage);
    PS_CHECK_DIMEN_AND_TYPE(inImage, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(inImage, outImage);
    PS_CHECK_ALLOC_IMAGE(outImage, inImage->numCols, inImage->numRows, inImage->type.type);
    PS_CHECK_DIMEN_AND_TYPE(outImage, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(outImage, outImage);
    PS_CHECK_ALLOC_VECTOR(outPerm, inImage->numRows, inImage->type.type);
    PS_CHECK_NULL_VECTOR(outPerm, outImage);
    PS_CHECK_DIMEN_AND_TYPE(outPerm, PS_DIMEN_VECTOR, outImage);

    // Initialize data
    numRows = inImage->numRows;
    numCols = inImage->numCols;
    arraySize = PSELEMTYPE_SIZEOF(inImage->type.type) * numRows * numCols;

    // Initialize GSL data
    perm.size = numCols;
    outPerm->n = numCols;
    perm.data = outPerm->data.V;
    PS_GSL_MATRIX_INITIALIZE(lu, outImage->data.V[0]);

    // Non-square matrices not allowed
    PS_CHECK_SQUARE(inImage, outImage);
    PS_CHECK_SQUARE(outImage, outImage);

    // Copy psImage input data into GSL matrix data to keep input data pristine
    memcpy(lu.data, inImage->data.V[0], arraySize);

    // Calculate LU decomposition
    gsl_linalg_LU_decomp(&lu, &perm, &signum);

    return outImage;
}

psVector* psMatrixLUSolve(psVector* outVector, const psImage* inImage,
                          const psVector* inVector, const psVector* inPerm)
{
    int arraySize = 0;
    int numRows = 0;
    int numCols = 0;
    gsl_matrix lu;
    gsl_permutation perm;
    gsl_vector b;
    gsl_vector x;

    // Error checks
    PS_CHECK_POINTERS(outVector, inVector, outVector);
    PS_CHECK_POINTERS(inVector, inPerm, outVector);
    PS_CHECK_POINTERS(outVector, inPerm, outVector);
    PS_CHECK_NULL_IMAGE(inImage, outVector);
    PS_CHECK_DIMEN_AND_TYPE(inImage, PS_DIMEN_IMAGE, outVector);
    PS_CHECK_SIZE_IMAGE(inImage, outVector);
    PS_CHECK_ALLOC_VECTOR(outVector, inImage->numRows, inImage->type.type);
    PS_CHECK_NULL_VECTOR(outVector, outVector);
    PS_CHECK_DIMEN_AND_TYPE(outVector, PS_DIMEN_VECTOR, outVector);
    PS_CHECK_NULL_VECTOR(inVector, outVector);
    PS_CHECK_DIMEN_AND_TYPE(inVector, PS_DIMEN_VECTOR, outVector);
    PS_CHECK_NULL_VECTOR(inPerm, outVector);
    PS_CHECK_DIMEN_AND_TYPE(inPerm, PS_DIMEN_VECTOR, outVector);

    // Initialize data
    numRows = inImage->numRows;
    numCols = inImage->numCols;
    arraySize = PSELEMTYPE_SIZEOF(inImage->type.type) * numRows * numCols;

    // Initialize GSL data
    PS_GSL_MATRIX_INITIALIZE(lu, inImage->data.V[0]);

    outVector->n = numCols;

    perm.size = inPerm->n;
    perm.data = inPerm->data.V;

    b.size = inVector->n;
    b.stride = 1;
    b.data = inVector->data.V;

    x.size = numCols;
    x.stride = 1;
    x.data = outVector->data.V;

    // Solve for {x} in equation: {b} = [A]{x}
    gsl_linalg_LU_solve(&lu, &perm, &b, &x);

    return outVector;
}

psImage* psMatrixInvert(psImage* outImage, const psImage* inImage, float *restrict det)
{
    int signum = 0;
    int arraySize = 0;
    int numRows = 0;
    int numCols = 0;
    gsl_matrix inv;
    gsl_matrix *lu;
    gsl_permutation *perm = NULL;

    // Error checks
    if (det == NULL) {
        psError(__func__, "Invalid operation: determinant argument is NULL.");
        return outImage;
    }
    PS_CHECK_POINTERS(inImage, outImage, outImage);
    PS_CHECK_NULL_IMAGE(inImage, outImage);
    PS_CHECK_DIMEN_AND_TYPE(inImage, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(inImage, outImage);
    PS_CHECK_ALLOC_IMAGE(outImage, inImage->numCols, inImage->numRows, inImage->type.type);
    PS_CHECK_DIMEN_AND_TYPE(outImage, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(outImage, outImage);

    // Initialize data
    numRows = inImage->numRows;
    numCols = inImage->numCols;
    arraySize = PSELEMTYPE_SIZEOF(inImage->type.type) * numRows * numCols;

    // Allocate GSL structs
    perm = gsl_permutation_alloc(inImage->numRows);
    lu = gsl_matrix_alloc(numRows, numCols);

    // Initialize GSL data
    PS_GSL_MATRIX_INITIALIZE(inv, outImage->data.V[0]);

    // Non-square matrices not allowed
    PS_CHECK_SQUARE(inImage, outImage);
    PS_CHECK_SQUARE(outImage, outImage);

    // Copy psImage input data into GSL matrix data to keep input data pristine
    memcpy(lu->data, inImage->data.V[0], arraySize);

    // Invert data and calculate determinant
    gsl_linalg_LU_decomp(lu, perm, &signum);
    gsl_linalg_LU_invert(lu, perm, &inv);
    *det = (float)gsl_linalg_LU_det(lu, signum);

    // Free GSL structs
    gsl_permutation_free(perm);
    gsl_matrix_free(lu);

    return outImage;
}

float *psMatrixDeterminant(const psImage* restrict inImage)
{
    int signum = 0;
    int arraySize = 0;
    int numRows = 0;
    int numCols = 0;
    float *det = NULL;
    gsl_matrix *lu = NULL;
    gsl_permutation *perm = NULL;

    // Error checks
    PS_CHECK_NULL_IMAGE(inImage, NULL);
    PS_CHECK_DIMEN_AND_TYPE(inImage, PS_DIMEN_IMAGE, NULL);
    PS_CHECK_SIZE_IMAGE(inImage, NULL);

    // Initialize data
    numRows = inImage->numRows;
    numCols = inImage->numCols;
    arraySize = PSELEMTYPE_SIZEOF(inImage->type.type) * numRows * numCols;

    // Allocate GSL structs
    perm = gsl_permutation_alloc(numRows);
    lu = gsl_matrix_alloc(numRows, numCols);

    // Non-square matrices not allowed
    PS_CHECK_SQUARE(inImage, 0);

    // Copy psImage input data into GSL matrix data to keep input data pristine
    memcpy(lu->data, inImage->data.V[0], arraySize);

    // Calculate determinant
    det = (float *)psAlloc(sizeof(float));
    gsl_linalg_LU_decomp(lu, perm, &signum);
    *det = (float)gsl_linalg_LU_det(lu, signum);

    // Free GSL structs
    gsl_permutation_free(perm);
    gsl_matrix_free(lu);

    return det;
}

psImage* psMatrixMultiply(psImage* outImage, psImage* inImage1, psImage* inImage2)
{
    int arraySize = 0;
    int numRows = 0;
    int numCols = 0;
    gsl_matrix m1;
    gsl_matrix m2;
    gsl_matrix m3;

    // Error checks
    PS_CHECK_POINTERS(inImage1, outImage, outImage);
    PS_CHECK_POINTERS(inImage1, inImage2, outImage);
    PS_CHECK_NULL_IMAGE(inImage1, outImage);
    PS_CHECK_DIMEN_AND_TYPE(inImage1, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(inImage1, outImage);
    PS_CHECK_NULL_IMAGE(inImage2, outImage);
    PS_CHECK_DIMEN_AND_TYPE(inImage2, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(inImage2, outImage);
    PS_CHECK_ALLOC_IMAGE(outImage, inImage2->numCols, inImage2->numRows, inImage2->type.type);
    PS_CHECK_DIMEN_AND_TYPE(inImage1, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(outImage, outImage);

    // Initialize data
    numRows = inImage1->numRows;
    numCols = inImage1->numCols;
    arraySize = PSELEMTYPE_SIZEOF(outImage->type.type) * numRows * numCols;

    // Initialize GSL data
    PS_GSL_MATRIX_INITIALIZE(m1, inImage1->data.V[0]);
    PS_GSL_MATRIX_INITIALIZE(m2, inImage2->data.V[0]);
    PS_GSL_MATRIX_INITIALIZE(m3, outImage->data.V[0]);

    // Non-square matrices not allowed
    PS_CHECK_SQUARE(inImage1, outImage);
    PS_CHECK_SQUARE(inImage2, outImage);
    PS_CHECK_SQUARE(outImage, outImage);

    gsl_linalg_matmult(&m1, &m2, &m3);

    return outImage;
}

psImage* psMatrixTranspose(psImage* outImage, const psImage* inImage)
{
    int arraySize = 0;
    int numRows = 0;
    int numCols = 0;
    gsl_matrix trans;

    // Error checks
    PS_CHECK_POINTERS(inImage, outImage, outImage);
    PS_CHECK_NULL_IMAGE(inImage, outImage);
    PS_CHECK_DIMEN_AND_TYPE(inImage, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(inImage, outImage);
    PS_CHECK_ALLOC_IMAGE(outImage, inImage->numCols, inImage->numRows, inImage->type.type);
    PS_CHECK_DIMEN_AND_TYPE(outImage, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(outImage, outImage);

    // Initialize data
    numRows = inImage->numRows;
    numCols = inImage->numCols;
    arraySize = PSELEMTYPE_SIZEOF(inImage->type.type) * numRows * numCols;

    // Initialize GSL data
    PS_GSL_MATRIX_INITIALIZE(trans, outImage->data.V[0]);

    // Non-square matrices not allowed
    PS_CHECK_SQUARE(inImage, outImage);
    PS_CHECK_SQUARE(outImage, outImage);

    // Copy psImage input data into psImage output data to keep input data pristine
    memcpy(outImage->data.V[0], inImage->data.V[0], arraySize);

    // Transpose data
    gsl_matrix_transpose(&trans);

    return outImage;
}

psImage* psMatrixEigenvectors(psImage* outImage, psImage* inImage)
{
    int numRows = 0;
    int numCols = 0;
    gsl_vector *eVals = NULL;
    gsl_eigen_symmv_workspace *w = NULL;
    gsl_matrix out;
    gsl_matrix in;

    // Error checks
    PS_CHECK_POINTERS(inImage, outImage, outImage);
    PS_CHECK_NULL_IMAGE(inImage, outImage);
    PS_CHECK_DIMEN_AND_TYPE(inImage, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(inImage, outImage);
    PS_CHECK_ALLOC_IMAGE(outImage, inImage->numCols, inImage->numRows, inImage->type.type);
    PS_CHECK_DIMEN_AND_TYPE(outImage, PS_DIMEN_IMAGE, outImage);
    PS_CHECK_SIZE_IMAGE(outImage, outImage);

    // Initialize data
    numRows = inImage->numRows;
    numCols = inImage->numCols;

    // Initialize GSL data
    PS_GSL_MATRIX_INITIALIZE(in, inImage->data.V[0]);
    PS_GSL_MATRIX_INITIALIZE(out, outImage->data.V[0]);

    // Allocate GSL structs
    eVals = gsl_vector_alloc(numRows);
    w = gsl_eigen_symmv_alloc(numRows);

    // Non-square matrices not allowed
    PS_CHECK_SQUARE(inImage, outImage);
    PS_CHECK_SQUARE(outImage, outImage);

    // Calculate Eigenvalues and Eigenvectors...Eigenvalues not currently used
    gsl_eigen_symmv(&in, eVals, &out, w);

    // Free GSL structs
    gsl_eigen_symmv_free(w);
    gsl_vector_free(eVals);

    return outImage;
}

psVector* psMatrixToVector(psVector* outVector, const psImage* inImage)
{
    int size = 0;

    // Error checks
    PS_CHECK_NULL_IMAGE(inImage, outVector);
    PS_CHECK_DIMEN_AND_TYPE(inImage, PS_DIMEN_IMAGE, outVector);
    PS_CHECK_SIZE_IMAGE(inImage, outVector);

    if (inImage->numRows == 1) {
        // Create transposed row vector
        PS_CHECK_ALLOC_VECTOR(outVector, inImage->numCols, inImage->type.type);
        outVector->type.dimen = PS_DIMEN_TRANSV;
    } else if (inImage->numCols == 1) {
        // Create non-transposed column vector
        PS_CHECK_ALLOC_VECTOR(outVector, inImage->numRows, inImage->type.type);
    } else {
        psError(__func__, "Image does not have dim with 1 col or 1 row: (%d x %d).", inImage->numRows,
                inImage->numCols);
        return outVector;
    }

    PS_CHECK_NULL_VECTOR(outVector, outVector);

    // More checks
    if (outVector->type.dimen == PS_DIMEN_VECTOR) {
        PS_CHECK_DIMEN_AND_TYPE(outVector, PS_DIMEN_VECTOR, outVector);

        if (outVector->n == 0) {
            outVector->n = inImage->numRows;
        }

        if (outVector->n != inImage->numRows) {
            psError(__func__, "Image and vector sizes differ: (%d vs %d).", inImage->numRows, outVector->n);
            return outVector;
        }

        size = PSELEMTYPE_SIZEOF(inImage->type.type) * inImage->numRows;

    } else if (outVector->type.dimen == PS_DIMEN_TRANSV) {
        PS_CHECK_DIMEN_AND_TYPE(outVector, PS_DIMEN_TRANSV, outVector);

        if (outVector->n == 0) {
            outVector->n = inImage->numCols;
        }

        if (outVector->n != inImage->numCols) {
            psError(__func__, "Image and vector sizes differ: (%d vs %d).", inImage->numCols, outVector->n);
            return outVector;
        }

        size = PSELEMTYPE_SIZEOF(inImage->type.type) * inImage->numCols;
    }

    memcpy(outVector->data.V, inImage->data.V[0], size);

    return outVector;
}

psImage* psVectorToMatrix(psImage* outImage, const psVector* inVector)
{
    int size = 0;

    // Error checks
    PS_CHECK_NULL_VECTOR(inVector, outImage);

    if (inVector->type.dimen == PS_DIMEN_VECTOR) {
        PS_CHECK_DIMEN_AND_TYPE(inVector, PS_DIMEN_VECTOR, outImage);
        PS_CHECK_SIZE_VECTOR(inVector, outImage);
        PS_CHECK_ALLOC_IMAGE(outImage, 1, inVector->n, PS_TYPE_F64)
        // More checks for PS_DIMEN_VECTOR
        if (outImage->numCols > 1) {
            psError(__func__, "Image has more than 1 column: numCols = %d.", outImage->numCols);
            return outImage;
        } else if (outImage->numRows != inVector->n) {
            psError(__func__, "Image and vector sizes differ: (%d vs %d).", outImage->numRows, inVector->n);
            return outImage;
        }

        size = PSELEMTYPE_SIZEOF(outImage->type.type) * outImage->numRows;

    } else if (inVector->type.dimen == PS_DIMEN_TRANSV) {
        PS_CHECK_DIMEN_AND_TYPE(inVector, PS_DIMEN_TRANSV, outImage);
        PS_CHECK_SIZE_VECTOR(inVector, outImage);
        PS_CHECK_ALLOC_IMAGE(outImage, inVector->n, 1, PS_TYPE_F64)
        // More checks for PS_DIMEN_TRANSV
        if (outImage->numRows > 1) {
            psError(__func__, "Image has more than 1 row: numRows = %d.", outImage->numRows);
            return outImage;
        } else if (outImage->numCols != inVector->n) {
            psError(__func__, "Image and vector sizes differ: (%d vs %d).", outImage->numCols, inVector->n);
            return outImage;
        }

        size = PSELEMTYPE_SIZEOF(outImage->type.type) * outImage->numCols;
    }

    PS_CHECK_NULL_IMAGE(outImage, outImage);
    PS_CHECK_DIMEN_AND_TYPE(outImage, PS_DIMEN_IMAGE, outImage);

    memcpy(outImage->data.V[0], inVector->data.V, size);

    return outImage;
}
