# include <pslib.h>
# include "psSparse.h"

static void psSparseFree (psSparse *sparse)
{
    if (sparse == NULL)
        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));
    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->Nrows = Nrows;
    psMemSetDeallocator(sparse, (psFreeFunc) psSparseFree);
    return (sparse);
}

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

    int k;

    if (i < j) {
        fprintf (stderr, "*** error: subdiagonal element ***\n");
        return;
    }

    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 + 100);
            psVectorRealloc (sparse->Si,  sparse->Si->nalloc + 100);
            psVectorRealloc (sparse->Sj,  sparse->Sj->nalloc + 100);
        }

        k = sparse->Nelem;
        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 + 100);
            psVectorRealloc (sparse->Si,  sparse->Si->nalloc + 100);
            psVectorRealloc (sparse->Sj,  sparse->Sj->nalloc + 100);
        }

        k = sparse->Nelem;
        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;
}

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

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

// multiple A * x
psVector *psSparseMatrixTimesVector (psVector *output, psSparse *matrix, psVector *vector)
{

    int i, Nelem;
    float F;

    if (output == NULL) {
        output = psVectorAlloc (vector->n, PS_DATA_F32);
    }

    Nelem = 0;
    for (int j = 0; j < vector->n; j++) {
        F = 0;
        while (matrix->Sj->data.S32[Nelem] == j) {
            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 *guess, psSparseConstraint constraint, psSparse *sparse, int Niter)
{

    psF32 dG;

    psVector *Qii = sparse->Qii;
    psVector *Bfj = sparse->Bfj;

    guess = psVectorCopy (guess, Bfj, PS_DATA_F32);

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

    for (int j = 0; j < Niter; j++) {
        dQ = psSparseMatrixTimesVector (dQ, sparse, guess);
        for (int i = 0; i < dQ->n; i++) {
            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;
                }
            }
            guess->data.F32[i] -= dG;
            guess->data.F32[i] = PS_MAX (guess->data.F32[i], constraint.paramMin);
            guess->data.F32[i] = PS_MIN (guess->data.F32[i], constraint.paramMax);
        }
    }
    psFree (dQ);
    return (guess);
}

void psSparseResort (psSparse *sparse)
{

    int Nelem = sparse->Nelem;

    psVector *index = psVectorSortIndex (NULL, sparse->Sj);
    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);
    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;
}

void psSparseMatrixTest ()
{

    // build a sparse matrix
    psSparse *sparse = psSparseAlloc (3, 9);

    psSparseMatrixElement (sparse, 0, 0, 3.0);
    psSparseMatrixElement (sparse, 1, 1, 2.0);
    psSparseMatrixElement (sparse, 2, 2, 1.0);

    psSparseMatrixElement (sparse, 1, 0, 0.1);
    psSparseMatrixElement (sparse, 2, 0, -0.1);

    psSparseResort (sparse);
    for (int i = 0; i < sparse->Nelem; i++) {
        fprintf (stderr, "%d %d %f\n",
                 sparse->Si->data.S32[i],
                 sparse->Sj->data.S32[i],
                 sparse->Aij->data.F32[i]);
    }

    psVector *x = psVectorAlloc (3, PS_DATA_F32);
    x->data.F32[0] = 3;
    x->data.F32[1] = 5;
    x->data.F32[2] = 7;
    fprintf (stderr, "x: %f %f %f\n", x->data.F32[0], x->data.F32[1], x->data.F32[2]);

    psVector *B = psSparseMatrixTimesVector (NULL, sparse, x);
    fprintf (stderr, "B: %f %f %f\n", B->data.F32[0], B->data.F32[1], B->data.F32[2]);

    sparse->Bfj->data.F32[0] = B->data.F32[0];
    sparse->Bfj->data.F32[1] = B->data.F32[1];
    sparse->Bfj->data.F32[2] = B->data.F32[2];

    psSparseConstraint constraint;
    constraint.paramMin   = -1e8;
    constraint.paramMax   = +1e8;
    constraint.paramDelta = +1e8;

    x = psSparseSolve (x, constraint, sparse, 0);
    fprintf (stderr, "x: %f %f %f\n", x->data.F32[0], x->data.F32[1], x->data.F32[2]);

    x = psSparseSolve (x, constraint, sparse, 1);
    fprintf (stderr, "x: %f %f %f\n", x->data.F32[0], x->data.F32[1], x->data.F32[2]);

    x = psSparseSolve (x, constraint, sparse, 2);
    fprintf (stderr, "x: %f %f %f\n", x->data.F32[0], x->data.F32[1], x->data.F32[2]);

    x = psSparseSolve (x, constraint, sparse, 3);
    fprintf (stderr, "x: %f %f %f\n", x->data.F32[0], x->data.F32[1], x->data.F32[2]);

    x = psSparseSolve (x, constraint, sparse, 4);
    fprintf (stderr, "x: %f %f %f\n", x->data.F32[0], x->data.F32[1], x->data.F32[2]);
    return;
}

