#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <math.h>
#include <unistd.h>

#include "psMemory.h"
#include "psError.h"
#include "psLogMsg.h"
#include "psVector.h"
#include "psAssert.h"
#include "psConstants.h"
#include "psSparse.h"


#define BUFFER 100                      // Size to increment at each go

static void sparseFree(psSparse *sparse)
{
    if (!sparse) {
        return;
    }
    psFree(sparse->Aij);
    psFree(sparse->Bfj);
    psFree(sparse->Qii);
    psFree(sparse->Si);
    psFree(sparse->Sj);
    return;
}

// allocate a sparse matrix container for Nrows, with Nelem slots allocated
psSparse *psSparseAlloc(int Nrows, int Nelem)
{
    psSparse *sparse = (psSparse *)psAlloc(sizeof(psSparse));
    psMemSetDeallocator(sparse, (psFreeFunc)sparseFree);

    sparse->Aij = psVectorAlloc(Nelem, PS_DATA_F32);
    sparse->Si  = psVectorAlloc(Nelem, PS_DATA_S32);
    sparse->Sj  = psVectorAlloc(Nelem, PS_DATA_S32);

    sparse->Aij->n = 0;
    sparse->Si->n  = 0;
    sparse->Sj->n  = 0;
    sparse->Nelem = 0;

    sparse->Bfj = psVectorAlloc(Nrows, PS_DATA_F32);
    sparse->Qii = psVectorAlloc(Nrows, PS_DATA_F32);
    sparse->Bfj->n = Nrows;
    sparse->Qii->n = Nrows;

    sparse->Nrows = Nrows;

    return sparse;
}

// user should only add elements above the diagonal, but we don't check this
bool psSparseMatrixElement(psSparse *sparse, int i, int j, float value)
{
    PS_ASSERT_PTR_NON_NULL(sparse, false);
    PS_ASSERT_INT_NONNEGATIVE(i, false);
    PS_ASSERT_INT_NONNEGATIVE(j, false);

    if (i < j) {
        psLogMsg(__func__, PS_LOG_WARN, "i=%d, j=%d refers to a sub-diagonal element; values switched.\n", i, j);
        int temp = i;
        i = j;
        j = temp;
    }

    if (i == j) {
        // add to the diagonal
        sparse->Qii->data.F32[i] = value;

        // check vectors lengths and extend if needed
        if (sparse->Nelem >= sparse->Aij->nalloc) {
            psVectorRealloc(sparse->Aij, sparse->Aij->nalloc + BUFFER);
            psVectorRealloc(sparse->Si,  sparse->Si->nalloc + BUFFER);
            psVectorRealloc(sparse->Sj,  sparse->Sj->nalloc + BUFFER);
        }

        int k = sparse->Nelem;         // Index at which to add
        sparse->Aij->data.F32[k] = value;
        sparse->Si->data.S32[k]  = i;
        sparse->Sj->data.S32[k]  = j;

        sparse->Nelem ++;
        sparse->Aij->n ++;
        sparse->Si->n ++;
        sparse->Sj->n ++;
    } else {
        // check vectors lengths and extend if needed
        if (sparse->Nelem >= sparse->Aij->nalloc - 1) {
            psVectorRealloc(sparse->Aij, sparse->Aij->nalloc + BUFFER);
            psVectorRealloc(sparse->Si,  sparse->Si->nalloc + BUFFER);
            psVectorRealloc(sparse->Sj,  sparse->Sj->nalloc + BUFFER);
        }

        int k = sparse->Nelem;         // Index at which to add
        sparse->Aij->data.F32[k] = value;
        sparse->Si->data.S32[k]  = i;
        sparse->Sj->data.S32[k]  = j;
        k++;

        sparse->Aij->data.F32[k] = value;
        sparse->Si->data.S32[k]  = j;
        sparse->Sj->data.S32[k]  = i;

        sparse->Nelem  += 2;
        sparse->Aij->n += 2;
        sparse->Si->n  += 2;
        sparse->Sj->n  += 2;
    }

    return true;
}

void inline psSparseVectorElement(psSparse *sparse, int i, float value)
{

    sparse->Bfj->data.F32[i] = value;
    return;
}

// multiply A * x
psVector *psSparseMatrixTimesVector(psVector *output, const psSparse *matrix, const psVector *vector)
{
    PS_ASSERT_PTR_NON_NULL(matrix, NULL);
    PS_ASSERT_VECTOR_NON_NULL(vector, NULL);

    output = psVectorRecycle(output, vector->n, PS_TYPE_F32);

    int Nelem = 0;                     // Number of elements
    for (int j = 0; j < vector->n; j++) {
        double F = 0;                    // Running total
        while (matrix->Sj->data.S32[Nelem] == j) {
            int i = matrix->Si->data.S32[Nelem];
            F += vector->data.F32[i] * matrix->Aij->data.F32[Nelem];
            Nelem++;
        }
        output->data.F32[j] = F;
    }
    return output;
}

psVector *psSparseSolve(psVector *output, psSparseConstraint constraint, const psSparse *sparse, int Niter)
{
    PS_ASSERT_PTR_NON_NULL(sparse, NULL);
    PS_ASSERT_INT_POSITIVE(Niter, NULL);
    PS_ASSERT_FLOAT_LARGER_THAN(constraint.paramDelta, 0.0, NULL);

    // Dereference some vectors
    psVector *Qii = sparse->Qii;
    psVector *Bfj = sparse->Bfj;

    output = psVectorCopy(output, Bfj, PS_DATA_F32);

    // temporary storage for intermediate results
    psVector *dQ = psVectorAlloc(output->n, PS_DATA_F32);
    dQ->n = output->n;

    for (int j = 0; j < Niter; j++) {
        dQ = psSparseMatrixTimesVector(dQ, sparse, output);
        for (int i = 0; i < dQ->n; i++) {
            psF32 dG = (dQ->data.F32[i] - Bfj->data.F32[i]) / Qii->data.F32[i];
            if (fabs (dG) > constraint.paramDelta) {
                if (dG > 0) {
                    dG = +constraint.paramDelta;
                } else {
                    dG = -constraint.paramDelta;
                }
            }
            output->data.F32[i] -= dG;
            output->data.F32[i] = PS_MAX(output->data.F32[i], constraint.paramMin);
            output->data.F32[i] = PS_MIN(output->data.F32[i], constraint.paramMax);
        }
    }
    psFree(dQ);

    return output;
}

bool psSparseResort(psSparse *sparse)
{
    int Nelem = sparse->Nelem;

    psVector *index = psVectorSortIndex(NULL, sparse->Sj); // Index key for sorting
    if (!index) {
        psError(PS_ERR_UNKNOWN, false, "Unable to sort sparse matrix.\n");
        return false;
    }
    psVector *Aij = sparse->Aij;
    psVector *Si = sparse->Si;
    psVector *Sj = sparse->Sj;

    // allocate new temporary vectors
    psVector *tAij = psVectorAlloc(Nelem, PS_DATA_F32);
    psVector *tSi  = psVectorAlloc(Nelem, PS_DATA_S32);
    psVector *tSj  = psVectorAlloc(Nelem, PS_DATA_S32);
    tAij->n = tSi->n = tSj->n = Nelem;

    for (int i = 0; i < Nelem; i++) {
        int j = index->data.U32[i];
        tAij->data.F32[i] = Aij->data.F32[j];
        tSi->data.S32[i]  = Si->data.S32[j];
        tSj->data.S32[i]  = Sj->data.S32[j];
    }
    psFree(index);
    psFree(Aij);
    psFree(Si);
    psFree(Sj);

    sparse->Aij = tAij;
    sparse->Si = tSi;
    sparse->Sj = tSj;

    return true;
}
