/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <cusolverDn.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include "mmio.h"

/* avoid Windows warnings (for example: strcpy, fscanf, etc.) */
#if defined(_WIN32)
#define _CRT_SECURE_NO_WARNINGS
#endif

/* various __inline__ __device__  function to initialize a T_ELEM */
template <typename T_ELEM> __inline__ T_ELEM cuGet(int);
template <> __inline__ float                 cuGet<float>(int x) { return float(x); }

template <> __inline__ double cuGet<double>(int x) { return double(x); }

template <> __inline__ cuComplex cuGet<cuComplex>(int x) { return (make_cuComplex(float(x), 0.0f)); }

template <> __inline__ cuDoubleComplex cuGet<cuDoubleComplex>(int x) { return (make_cuDoubleComplex(double(x), 0.0)); }

template <typename T_ELEM> __inline__ T_ELEM cuGet(int, int);
template <> __inline__ float                 cuGet<float>(int x, int y) { return float(x); }

template <> __inline__ double cuGet<double>(int x, int y) { return double(x); }

template <> __inline__ cuComplex cuGet<cuComplex>(int x, int y) { return make_cuComplex(float(x), float(y)); }

template <> __inline__ cuDoubleComplex cuGet<cuDoubleComplex>(int x, int y)
{
    return (make_cuDoubleComplex(double(x), double(y)));
}

template <typename T_ELEM> __inline__ T_ELEM cuGet(float);
template <> __inline__ float                 cuGet<float>(float x) { return float(x); }

template <> __inline__ double cuGet<double>(float x) { return double(x); }

template <> __inline__ cuComplex cuGet<cuComplex>(float x) { return (make_cuComplex(float(x), 0.0f)); }

template <> __inline__ cuDoubleComplex cuGet<cuDoubleComplex>(float x)
{
    return (make_cuDoubleComplex(double(x), 0.0));
}

template <typename T_ELEM> __inline__ T_ELEM cuGet(float, float);
template <> __inline__ float                 cuGet<float>(float x, float y) { return float(x); }

template <> __inline__ double cuGet<double>(float x, float y) { return double(x); }

template <> __inline__ cuComplex cuGet<cuComplex>(float x, float y) { return (make_cuComplex(float(x), float(y))); }

template <> __inline__ cuDoubleComplex cuGet<cuDoubleComplex>(float x, float y)
{
    return (make_cuDoubleComplex(double(x), double(y)));
}

template <typename T_ELEM> __inline__ T_ELEM cuGet(double);
template <> __inline__ float                 cuGet<float>(double x) { return float(x); }

template <> __inline__ double cuGet<double>(double x) { return double(x); }

template <> __inline__ cuComplex cuGet<cuComplex>(double x) { return (make_cuComplex(float(x), 0.0f)); }

template <> __inline__ cuDoubleComplex cuGet<cuDoubleComplex>(double x)
{
    return (make_cuDoubleComplex(double(x), 0.0));
}

template <typename T_ELEM> __inline__ T_ELEM cuGet(double, double);
template <> __inline__ float                 cuGet<float>(double x, double y) { return float(x); }

template <> __inline__ double cuGet<double>(double x, double y) { return double(x); }

template <> __inline__ cuComplex cuGet<cuComplex>(double x, double y) { return (make_cuComplex(float(x), float(y))); }

template <> __inline__ cuDoubleComplex cuGet<cuDoubleComplex>(double x, double y)
{
    return (make_cuDoubleComplex(double(x), double(y)));
}

static void compress_index(const int *Ind, int nnz, int m, int *Ptr, int base)
{
    int i;

    /* initialize everything to zero */
    for (i = 0; i < m + 1; i++) {
        Ptr[i] = 0;
    }
    /* count elements in every row */
    Ptr[0] = base;
    for (i = 0; i < nnz; i++) {
        Ptr[Ind[i] + (1 - base)]++;
    }
    /* add all the values */
    for (i = 0; i < m; i++) {
        Ptr[i + 1] += Ptr[i];
    }
}

struct cooFormat
{
    int i;
    int j;
    int p; // permutation
};

int cmp_cooFormat_csr(struct cooFormat *s, struct cooFormat *t)
{
    if (s->i < t->i) {
        return -1;
    }
    else if (s->i > t->i) {
        return 1;
    }
    else {
        return s->j - t->j;
    }
}

int cmp_cooFormat_csc(struct cooFormat *s, struct cooFormat *t)
{
    if (s->j < t->j) {
        return -1;
    }
    else if (s->j > t->j) {
        return 1;
    }
    else {
        return s->i - t->i;
    }
}

typedef int (*FUNPTR)(const void *, const void *);
typedef int (*FUNPTR2)(struct cooFormat *s, struct cooFormat *t);

static FUNPTR2 fptr_array[2] = {
    cmp_cooFormat_csr,
    cmp_cooFormat_csc,
};

static int verify_pattern(int m, int nnz, int *csrRowPtr, int *csrColInd)
{
    int i, col, start, end, base_index;
    int error_found = 0;

    if (nnz != (csrRowPtr[m] - csrRowPtr[0])) {
        fprintf(stderr,
                "Error (nnz check failed): (csrRowPtr[%d]=%d - csrRowPtr[%d]=%d) "
                "!= (nnz=%d)\n",
                0,
                csrRowPtr[0],
                m,
                csrRowPtr[m],
                nnz);
        error_found = 1;
    }

    base_index = csrRowPtr[0];
    if ((0 != base_index) && (1 != base_index)) {
        fprintf(stderr, "Error (base index check failed): base index = %d\n", base_index);
        error_found = 1;
    }

    for (i = 0; (!error_found) && (i < m); i++) {
        start = csrRowPtr[i] - base_index;
        end   = csrRowPtr[i + 1] - base_index;
        if (start > end) {
            fprintf(stderr,
                    "Error (corrupted row): csrRowPtr[%d] (=%d) > csrRowPtr[%d] (=%d)\n",
                    i,
                    start + base_index,
                    i + 1,
                    end + base_index);
            error_found = 1;
        }
        for (col = start; col < end; col++) {
            if (csrColInd[col] < base_index) {
                fprintf(stderr, "Error (column vs. base index check failed): csrColInd[%d] < %d\n", col, base_index);
                error_found = 1;
            }
            if ((col < (end - 1)) && (csrColInd[col] >= csrColInd[col + 1])) {
                fprintf(stderr,
                        "Error (sorting of the column indecis check failed): "
                        "(csrColInd[%d]=%d) >= (csrColInd[%d]=%d)\n",
                        col,
                        csrColInd[col],
                        col + 1,
                        csrColInd[col + 1]);
                error_found = 1;
            }
        }
    }
    return error_found;
}

template <typename T_ELEM>
int loadMMSparseMatrix(char    *filename,
                       char     elem_type,
                       bool     csrFormat,
                       int     *m,
                       int     *n,
                       int     *nnz,
                       T_ELEM **aVal,
                       int    **aRowInd,
                       int    **aColInd,
                       int      extendSymMatrix)
{
    MM_typecode       matcode;
    double           *tempVal;
    int              *tempRowInd, *tempColInd;
    double           *tval;
    int              *trow, *tcol;
    int              *csrRowPtr, *cscColPtr;
    int               i, j, error, base, count;
    struct cooFormat *work;

    /* read the matrix */
    error = mm_read_mtx_crd(filename, m, n, nnz, &trow, &tcol, &tval, &matcode);
    if (error) {
        fprintf(stderr, "!!!! can not open file: '%s'\n", filename);
        return 1;
    }

    /* start error checking */
    if (mm_is_complex(matcode) && ((elem_type != 'z') && (elem_type != 'c'))) {
        fprintf(stderr, "!!!! complex matrix requires type 'z' or 'c'\n");
        return 1;
    }

    if (mm_is_dense(matcode) || mm_is_array(matcode) || mm_is_pattern(matcode) /*|| mm_is_integer(matcode)*/) {
        fprintf(stderr, "!!!! dense, array, pattern and integer matrices are not supported\n");
        return 1;
    }

    /* if necessary symmetrize the pattern (transform from triangular to full) */
    if ((extendSymMatrix) && (mm_is_symmetric(matcode) || mm_is_hermitian(matcode) || mm_is_skew(matcode))) {
        // count number of non-diagonal elements
        count = 0;
        for (i = 0; i < (*nnz); i++) {
            if (trow[i] != tcol[i]) {
                count++;
            }
        }
        // allocate space for the symmetrized matrix
        tempRowInd = (int *)malloc((*nnz + count) * sizeof(int));
        tempColInd = (int *)malloc((*nnz + count) * sizeof(int));
        if (mm_is_real(matcode) || mm_is_integer(matcode)) {
            tempVal = (double *)malloc((*nnz + count) * sizeof(double));
        }
        else {
            tempVal = (double *)malloc(2 * (*nnz + count) * sizeof(double));
        }
        // copy the elements regular and transposed locations
        for (j = 0, i = 0; i < (*nnz); i++) {
            tempRowInd[j] = trow[i];
            tempColInd[j] = tcol[i];
            if (mm_is_real(matcode) || mm_is_integer(matcode)) {
                tempVal[j] = tval[i];
            }
            else {
                tempVal[2 * j]     = tval[2 * i];
                tempVal[2 * j + 1] = tval[2 * i + 1];
            }
            j++;
            if (trow[i] != tcol[i]) {
                tempRowInd[j] = tcol[i];
                tempColInd[j] = trow[i];
                if (mm_is_real(matcode) || mm_is_integer(matcode)) {
                    if (mm_is_skew(matcode)) {
                        tempVal[j] = -tval[i];
                    }
                    else {
                        tempVal[j] = tval[i];
                    }
                }
                else {
                    if (mm_is_hermitian(matcode)) {
                        tempVal[2 * j]     = tval[2 * i];
                        tempVal[2 * j + 1] = -tval[2 * i + 1];
                    }
                    else {
                        tempVal[2 * j]     = tval[2 * i];
                        tempVal[2 * j + 1] = tval[2 * i + 1];
                    }
                }
                j++;
            }
        }
        (*nnz) += count;
        // free temporary storage
        free(trow);
        free(tcol);
        free(tval);
    }
    else {
        tempRowInd = trow;
        tempColInd = tcol;
        tempVal    = tval;
    }
    // life time of (trow, tcol, tval) is over.
    // please use COO format (tempRowInd, tempColInd, tempVal)

    // use qsort to sort COO format
    work = (struct cooFormat *)malloc(sizeof(struct cooFormat) * (*nnz));
    if (NULL == work) {
        fprintf(stderr, "!!!! allocation error, malloc failed\n");
        return 1;
    }
    for (i = 0; i < (*nnz); i++) {
        work[i].i = tempRowInd[i];
        work[i].j = tempColInd[i];
        work[i].p = i; // permutation is identity
    }

    if (csrFormat) {
        /* create row-major ordering of indices (sorted by row and within each row
         * by column) */
        qsort(work, *nnz, sizeof(struct cooFormat), (FUNPTR)fptr_array[0]);
    }
    else {
        /* create column-major ordering of indices (sorted by column and within each
         * column by row) */
        qsort(work, *nnz, sizeof(struct cooFormat), (FUNPTR)fptr_array[1]);
    }

    // (tempRowInd, tempColInd) is sorted either by row-major or by col-major
    for (i = 0; i < (*nnz); i++) {
        tempRowInd[i] = work[i].i;
        tempColInd[i] = work[i].j;
    }

    // setup base
    // check if there is any row/col 0, if so base-0
    // check if there is any row/col equal to matrix dimension m/n, if so base-1
    int base0 = 0;
    int base1 = 0;
    for (i = 0; i < (*nnz); i++) {
        const int row = tempRowInd[i];
        const int col = tempColInd[i];
        if ((0 == row) || (0 == col)) {
            base0 = 1;
        }
        if ((*m == row) || (*n == col)) {
            base1 = 1;
        }
    }
    if (base0 && base1) {
        printf("Error: input matrix is base-0 and base-1 \n");
        return 1;
    }

    base = 0;
    if (base1) {
        base = 1;
    }

    /* compress the appropriate indices */
    if (csrFormat) {
        /* CSR format (assuming row-major format) */
        csrRowPtr = (int *)malloc(((*m) + 1) * sizeof(csrRowPtr[0]));
        if (!csrRowPtr)
            return 1;
        compress_index(tempRowInd, *nnz, *m, csrRowPtr, base);

        *aRowInd = csrRowPtr;
        *aColInd = (int *)malloc((*nnz) * sizeof(int));
    }
    else {
        /* CSC format (assuming column-major format) */
        cscColPtr = (int *)malloc(((*n) + 1) * sizeof(cscColPtr[0]));
        if (!cscColPtr)
            return 1;
        compress_index(tempColInd, *nnz, *n, cscColPtr, base);

        *aColInd = cscColPtr;
        *aRowInd = (int *)malloc((*nnz) * sizeof(int));
    }

    /* transfrom the matrix values of type double into one of the cusparse library
     * types */
    *aVal = (T_ELEM *)malloc((*nnz) * sizeof(T_ELEM));

    for (i = 0; i < (*nnz); i++) {
        if (csrFormat) {
            (*aColInd)[i] = tempColInd[i];
        }
        else {
            (*aRowInd)[i] = tempRowInd[i];
        }
        if (mm_is_real(matcode) || mm_is_integer(matcode)) {
            (*aVal)[i] = cuGet<T_ELEM>(tempVal[work[i].p]);
        }
        else {
            (*aVal)[i] = cuGet<T_ELEM>(tempVal[2 * work[i].p], tempVal[2 * work[i].p + 1]);
        }
    }

    /* check for corruption */
    int error_found;
    if (csrFormat) {
        error_found = verify_pattern(*m, *nnz, *aRowInd, *aColInd);
    }
    else {
        error_found = verify_pattern(*n, *nnz, *aColInd, *aRowInd);
    }
    if (error_found) {
        fprintf(stderr, "!!!! verify_pattern failed\n");
        return 1;
    }

    /* cleanup and exit */
    free(work);
    free(tempVal);
    free(tempColInd);
    free(tempRowInd);

    return 0;
}

/* specific instantiation */
template int loadMMSparseMatrix<float>(char   *filename,
                                       char    elem_type,
                                       bool    csrFormat,
                                       int    *m,
                                       int    *n,
                                       int    *nnz,
                                       float **aVal,
                                       int   **aRowInd,
                                       int   **aColInd,
                                       int     extendSymMatrix);

template int loadMMSparseMatrix<double>(char    *filename,
                                        char     elem_type,
                                        bool     csrFormat,
                                        int     *m,
                                        int     *n,
                                        int     *nnz,
                                        double **aVal,
                                        int    **aRowInd,
                                        int    **aColInd,
                                        int      extendSymMatrix);

template int loadMMSparseMatrix<cuComplex>(char       *filename,
                                           char        elem_type,
                                           bool        csrFormat,
                                           int        *m,
                                           int        *n,
                                           int        *nnz,
                                           cuComplex **aVal,
                                           int       **aRowInd,
                                           int       **aColInd,
                                           int         extendSymMatrix);

template int loadMMSparseMatrix<cuDoubleComplex>(char             *filename,
                                                 char              elem_type,
                                                 bool              csrFormat,
                                                 int              *m,
                                                 int              *n,
                                                 int              *nnz,
                                                 cuDoubleComplex **aVal,
                                                 int             **aRowInd,
                                                 int             **aColInd,
                                                 int               extendSymMatrix);
