/*******************************************************
 * Copyright (c) 2014, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/

#include <af/sparse.h>
#include <af/array.h>
#include <af/algorithm.h>
#include <sparse_handle.hpp>
#include <sparse.hpp>
#include <handle.hpp>
#include <backend.hpp>
#include <common/err_common.hpp>
#include <arith.hpp>
#include <lookup.hpp>
#include <platform.hpp>

using namespace detail;
using namespace common;
using af::dim4;

const SparseArrayBase& getSparseArrayBase(const af_array in, bool device_check)
{
    const SparseArrayBase *base = static_cast<SparseArrayBase*>(reinterpret_cast<void *>(in));

    if(!base->isSparse()) {
        AF_ERROR("Input is not a SparseArray and cannot be used in Sparse functions",
                 AF_ERR_ARG);
    }

    if (device_check && base->getDevId() != detail::getActiveDeviceId()) {
        AF_ERROR("Input Array not created on current device", AF_ERR_DEVICE);
    }

    return *base;
}

////////////////////////////////////////////////////////////////////////////////
// Sparse Creation
////////////////////////////////////////////////////////////////////////////////
template<typename T>
af_array createSparseArrayFromData(const af::dim4 &dims, const af_array values,
                                   const af_array rowIdx, const af_array colIdx,
                                   const af::storage stype)
{
    SparseArray<T> sparse = common::createArrayDataSparseArray(
                            dims, getArray<T>(values),
                            getArray<int>(rowIdx), getArray<int>(colIdx),
                            stype);
    return getHandle(sparse);
}

af_err af_create_sparse_array(
                 af_array *out,
                 const dim_t nRows, const dim_t nCols,
                 const af_array values, const af_array rowIdx, const af_array colIdx,
                 const af_storage stype)
{
    try {
        // Checks:
        // rowIdx and colIdx arrays are of s32 type
        // values is of floating point type
        // if COO, rowIdx, colIdx and values should have same dims
        // if CRS, colIdx and values should have same dims, rowIdx.dims = nRows
        // if CRC, rowIdx and values should have same dims, colIdx.dims = nCols
        // stype is within acceptable range
        // type is floating type

        if(!(stype == AF_STORAGE_CSR
          || stype == AF_STORAGE_CSC
          || stype == AF_STORAGE_COO)) {
            AF_ERROR("Storage type is out of range/unsupported", AF_ERR_ARG);
        }

        const ArrayInfo& vInfo = getInfo(values);
        const ArrayInfo& rInfo = getInfo(rowIdx);
        const ArrayInfo& cInfo = getInfo(colIdx);

        TYPE_ASSERT(vInfo.isFloating());
        DIM_ASSERT(3, vInfo.isLinear());
        ARG_ASSERT(4, rInfo.getType() == s32);
        DIM_ASSERT(4, rInfo.isLinear());
        ARG_ASSERT(5, cInfo.getType() == s32);
        DIM_ASSERT(5, cInfo.isLinear());

        const size_t nNZ = vInfo.elements();
        if(stype == AF_STORAGE_COO) {
          DIM_ASSERT(4, rInfo.elements() == nNZ);
          DIM_ASSERT(5, cInfo.elements() == nNZ);
        } else if(stype == AF_STORAGE_CSR) {
          DIM_ASSERT(4, rInfo.elements() == nRows + 1);
          DIM_ASSERT(5, cInfo.elements() == nNZ);
        } else if(stype == AF_STORAGE_CSC) {
          DIM_ASSERT(4, rInfo.elements() == nNZ);
          DIM_ASSERT(5, cInfo.elements() == nCols + 1);
        }

        af_array output = 0;

        af::dim4 dims(nRows, nCols);

        switch(vInfo.getType()) {
            case f32: output = createSparseArrayFromData<float  >(dims, values, rowIdx, colIdx, stype); break;
            case f64: output = createSparseArrayFromData<double >(dims, values, rowIdx, colIdx, stype); break;
            case c32: output = createSparseArrayFromData<cfloat >(dims, values, rowIdx, colIdx, stype); break;
            case c64: output = createSparseArrayFromData<cdouble>(dims, values, rowIdx, colIdx, stype); break;
            default : TYPE_ERROR(1, vInfo.getType());
        }
        std::swap(*out, output);

    } CATCHALL;

    return AF_SUCCESS;
}

template<typename T>
af_array createSparseArrayFromPtr(
        const af::dim4 &dims, const dim_t nNZ,
        const T * const values, const int * const rowIdx, const int * const colIdx,
        const af::storage stype, const af::source source)
{
    SparseArray<T> sparse = createEmptySparseArray<T>(dims, nNZ, stype);

    if(nNZ) {
        if(source == afHost)
            sparse = common::createHostDataSparseArray(
                            dims, nNZ, values, rowIdx, colIdx, stype);
        else if (source == afDevice)
            sparse = common::createDeviceDataSparseArray(
                            dims, nNZ, values, rowIdx, colIdx, stype);
    }

    return getHandle(sparse);
}

af_err af_create_sparse_array_from_ptr(
                 af_array *out,
                 const dim_t nRows, const dim_t nCols, const dim_t nNZ,
                 const void * const values, const int * const rowIdx, const int * const colIdx,
                 const af_dtype type, const af_storage stype,
                 const af_source source)
{
    try {
        // Checks:
        // rowIdx and colIdx arrays are of s32 type
        // values is of floating point type
        // if COO, rowIdx, colIdx and values should have same dims
        // if CRS, colIdx and values should have same dims, rowIdx.dims = nRows
        // if CRC, rowIdx and values should have same dims, colIdx.dims = nCols
        // stype is within acceptable range
        // type is floating type
        if(!(stype == AF_STORAGE_CSR
          || stype == AF_STORAGE_CSC
          || stype == AF_STORAGE_COO)) {
            AF_ERROR("Storage type is out of range/unsupported", AF_ERR_ARG);
        }

        TYPE_ASSERT(type == f32 || type == f64
                 || type == c32 || type == c64);


        af_array output = 0;

        af::dim4 dims(nRows, nCols);

        switch(type) {
            case f32: output = createSparseArrayFromPtr<float  >
                               (dims, nNZ, static_cast<const float  *>(values), rowIdx, colIdx, stype, source);
                      break;
            case f64: output = createSparseArrayFromPtr<double >
                               (dims, nNZ, static_cast<const double *>(values), rowIdx, colIdx, stype, source);
                      break;
            case c32: output = createSparseArrayFromPtr<cfloat >
                               (dims, nNZ, static_cast<const cfloat *>(values), rowIdx, colIdx, stype, source);
                      break;
            case c64: output = createSparseArrayFromPtr<cdouble>
                               (dims, nNZ, static_cast<const cdouble*>(values), rowIdx, colIdx, stype, source);
                      break;
            default : TYPE_ERROR(1, type);
        }
        std::swap(*out, output);

    } CATCHALL;

    return AF_SUCCESS;
}

template<typename T>
af_array createSparseArrayFromDense(
        const af::dim4 &dims, const af_array _in,
        const af_storage stype)
{
    const Array<T> in = getArray<T>(_in);

    switch(stype) {
        case AF_STORAGE_CSR:
            return getHandle(sparseConvertDenseToStorage<T, AF_STORAGE_CSR>(in));
        case AF_STORAGE_COO:
            return getHandle(sparseConvertDenseToStorage<T, AF_STORAGE_COO>(in));
        case AF_STORAGE_CSC:
            //return getHandle(sparseConvertDenseToStorage<T, AF_STORAGE_CSC>(in));
        default: AF_ERROR("Storage type is out of range/unsupported", AF_ERR_ARG);
    }
}

af_err af_create_sparse_array_from_dense(af_array *out, const af_array in,
                                         const af_storage stype)
{
    try {
        // Checks:
        // stype is within acceptable range
        // values is of floating point type

        const ArrayInfo& info = getInfo(in);

        if(!(stype == AF_STORAGE_CSR
          || stype == AF_STORAGE_CSC
          || stype == AF_STORAGE_COO)) {
            AF_ERROR("Storage type is out of range/unsupported", AF_ERR_ARG);
        }

        // Only matrices allowed
        DIM_ASSERT(1, info.ndims() == 2);

        TYPE_ASSERT(info.isFloating());

        af::dim4 dims(info.dims()[0], info.dims()[1]);

        af_array output = 0;

        switch(info.getType()) {
            case f32: output = createSparseArrayFromDense<float  >(dims, in, stype); break;
            case f64: output = createSparseArrayFromDense<double >(dims, in, stype); break;
            case c32: output = createSparseArrayFromDense<cfloat >(dims, in, stype); break;
            case c64: output = createSparseArrayFromDense<cdouble>(dims, in, stype); break;
            default: TYPE_ERROR(1, info.getType());
        }
        std::swap(*out, output);

    } CATCHALL;

    return AF_SUCCESS;
}

template<typename T>
af_array sparseConvertStorage(const af_array in_, const af_storage destStorage)
{
    const SparseArray<T> in = getSparseArray<T>(in_);

    if(destStorage == AF_STORAGE_DENSE) {
        // Returns a regular af_array, not sparse
        switch(in.getStorage()) {
            case AF_STORAGE_CSR:
                return getHandle(detail::sparseConvertStorageToDense<T, AF_STORAGE_CSR>(in));
            case AF_STORAGE_COO:
                return getHandle(detail::sparseConvertStorageToDense<T, AF_STORAGE_COO>(in));
            default:
                AF_ERROR("Invalid storage type of input array", AF_ERR_ARG);
        }
    } else if(destStorage == AF_STORAGE_CSR) {
        // Returns a sparse af_array
        switch(in.getStorage()) {
            case AF_STORAGE_CSR:
                return retainSparseHandle<T>(in_);
            case AF_STORAGE_COO:
                return getHandle(detail::sparseConvertStorageToStorage<T, AF_STORAGE_CSR, AF_STORAGE_COO>(in));
            default:
                AF_ERROR("Invalid storage type of input array", AF_ERR_ARG);
        }
    } else if(destStorage == AF_STORAGE_COO) {
        // Returns a sparse af_array
        switch(in.getStorage()) {
            case AF_STORAGE_CSR:
                return getHandle(detail::sparseConvertStorageToStorage<T, AF_STORAGE_COO, AF_STORAGE_CSR>(in));
            case AF_STORAGE_COO:
                return retainSparseHandle<T>(in_);
            default:
                AF_ERROR("Invalid storage type of input array", AF_ERR_ARG);
        }
    }

    // Shoud never come here
    return NULL;
}

af_err af_sparse_convert_to(af_array *out, const af_array in,
                            const af_storage destStorage)
{
    try {
        // Handle dense case
        const ArrayInfo& info = getInfo(in, false, true);
        if(!info.isSparse()) {    // If input is dense
            return af_create_sparse_array_from_dense(out, in, destStorage);
        }

        af_array output = 0;

        const SparseArrayBase base = getSparseArrayBase(in);

        // Dense not allowed as input -> Should never happen with SparseArrayBase
        // CSC is currently not supported
        ARG_ASSERT(1, base.getStorage() != AF_STORAGE_DENSE
                   && base.getStorage() != AF_STORAGE_CSC);

        // Conversion to and from CSC is not supported
        ARG_ASSERT(2, destStorage != AF_STORAGE_CSC);

        if(base.getStorage() == destStorage) {
            // Return a reference
            AF_CHECK(af_retain_array(out, in));
            return AF_SUCCESS;
        }

        switch(base.getType()) {
            case f32: output = sparseConvertStorage<float  >(in, destStorage); break;
            case f64: output = sparseConvertStorage<double >(in, destStorage); break;
            case c32: output = sparseConvertStorage<cfloat >(in, destStorage); break;
            case c64: output = sparseConvertStorage<cdouble>(in, destStorage); break;
            default : AF_ERROR("Output storage type is not valid", AF_ERR_ARG);
        }
        std::swap(*out, output);
    }
    CATCHALL;
    return AF_SUCCESS;
}

af_err af_sparse_to_dense(af_array *out, const af_array in)
{
    try {
        af_array output = 0;

        const SparseArrayBase base = getSparseArrayBase(in);

        // Dense not allowed as input -> Should never happen
        // To convert from dense to type, use the create* functions
        ARG_ASSERT(1, base.getStorage() != AF_STORAGE_DENSE);

        switch(base.getType()) {
            case f32: output = sparseConvertStorage<float  >(in, AF_STORAGE_DENSE); break;
            case f64: output = sparseConvertStorage<double >(in, AF_STORAGE_DENSE); break;
            case c32: output = sparseConvertStorage<cfloat >(in, AF_STORAGE_DENSE); break;
            case c64: output = sparseConvertStorage<cdouble>(in, AF_STORAGE_DENSE); break;
            default : AF_ERROR("Output storage type is not valid", AF_ERR_ARG);
        }
        std::swap(*out, output);
    }
    CATCHALL;
    return AF_SUCCESS;
}

////////////////////////////////////////////////////////////////////////////////
// Get Functions
////////////////////////////////////////////////////////////////////////////////
template<typename T>
af_array getSparseValues(const af_array in)
{
    return getHandle(getSparseArray<T>(in).getValues());
}

af_err af_sparse_get_info(af_array *values, af_array *rows, af_array *cols, af_storage *stype,
                          const af_array in)
{
    try {
        if(values != NULL) AF_CHECK(af_sparse_get_values(values, in));
        if(rows   != NULL) AF_CHECK(af_sparse_get_row_idx(rows , in));
        if(cols   != NULL) AF_CHECK(af_sparse_get_col_idx(cols , in));
        if(stype  != NULL) AF_CHECK(af_sparse_get_storage(stype, in));
    }
    CATCHALL;

    return AF_SUCCESS;
}

af_err af_sparse_get_values(af_array *out, const af_array in)
{
    try{
        const SparseArrayBase base = getSparseArrayBase(in);

        af_array output = 0;

        switch(base.getType()) {
            case f32: output = getSparseValues<float  >(in); break;
            case f64: output = getSparseValues<double >(in); break;
            case c32: output = getSparseValues<cfloat >(in); break;
            case c64: output = getSparseValues<cdouble>(in); break;
            default : TYPE_ERROR(1, base.getType());
        }
        std::swap(*out, output);
    }
    CATCHALL;
    return AF_SUCCESS;
}

af_err af_sparse_get_row_idx(af_array *out, const af_array in)
{
    try {
        const SparseArrayBase base = getSparseArrayBase(in);
        *out = getHandle(base.getRowIdx());
    } CATCHALL;
    return AF_SUCCESS;
}

af_err af_sparse_get_col_idx(af_array *out, const af_array in)
{
    try {
        const SparseArrayBase base = getSparseArrayBase(in);
        *out = getHandle(base.getColIdx());
    } CATCHALL;
    return AF_SUCCESS;
}

af_err af_sparse_get_nnz(dim_t *out, const af_array in)
{
    try {
        const SparseArrayBase base = getSparseArrayBase(in);
        *out = base.getNNZ();
    } CATCHALL;
    return AF_SUCCESS;
}

af_err af_sparse_get_storage(af_storage *out, const af_array in)
{
    try {
        const SparseArrayBase base = getSparseArrayBase(in);
        *out = base.getStorage();
    } CATCHALL;
    return AF_SUCCESS;
}
