#include "qdldl.h"
#include <stdio.h>
#include <stdlib.h>

void print_arrayi(const QDLDL_int* data, QDLDL_int n, char* varName);
void print_arrayf(const QDLDL_float* data, QDLDL_int n, char* varName);
void print_line(void);

const QDLDL_int   An = 10;
const QDLDL_int   Ap[] = { 0, 1, 2, 4, 5, 6, 8, 10, 12, 14, 17 };
const QDLDL_int   Ai[] = { 0, 1, 1, 2, 3, 4, 1, 5, 0, 6, 3, 7, 6, 8, 1, 2, 9 };
const QDLDL_float Ax[] = { 1.0,        0.460641,  -0.121189, 0.417928,  0.177828,   0.1,
                           -0.0290058, -1.0,      0.350321,  -0.441092, -0.0845395, -0.316228,
                           0.178663,   -0.299077, 0.182452,  -1.56506,  -0.1 };
const QDLDL_float b[] = { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };

int main()
{
    QDLDL_int i; // Counter

    // data for L and D factors
    QDLDL_int    Ln = An;
    QDLDL_int*   Lp;
    QDLDL_int*   Li;
    QDLDL_float* Lx;
    QDLDL_float* D;
    QDLDL_float* Dinv;

    // data for elim tree calculation
    QDLDL_int* etree;
    QDLDL_int* Lnz;
    QDLDL_int  sumLnz;

    // working data for factorisation
    QDLDL_int*   iwork;
    QDLDL_bool*  bwork;
    QDLDL_float* fwork;

    // Data for results of A\b
    QDLDL_float* x;

    /*--------------------------------
     * pre-factorisation memory allocations
     *---------------------------------*/

    // These can happen *before* the etree is calculated
    // since the sizes are not sparsity pattern specific

    // For the elimination tree
    etree = (QDLDL_int*) malloc(sizeof(QDLDL_int) * An);
    Lnz = (QDLDL_int*) malloc(sizeof(QDLDL_int) * An);

    // For the L factors.   Li and Lx are sparsity dependent
    // so must be done after the etree is constructed
    Lp = (QDLDL_int*) malloc(sizeof(QDLDL_int) * (An + 1));
    D = (QDLDL_float*) malloc(sizeof(QDLDL_float) * An);
    Dinv = (QDLDL_float*) malloc(sizeof(QDLDL_float) * An);

    // Working memory.  Note that both the etree and factor
    // calls requires a working vector of QDLDL_int, with
    // the factor function requiring 3*An elements and the
    // etree only An elements.   Just allocate the larger
    // amount here and use it in both places
    iwork = (QDLDL_int*) malloc(sizeof(QDLDL_int) * (3 * An));
    bwork = (QDLDL_bool*) malloc(sizeof(QDLDL_bool) * An);
    fwork = (QDLDL_float*) malloc(sizeof(QDLDL_float) * An);

    /*--------------------------------
     * elimination tree calculation
     *---------------------------------*/
    sumLnz = QDLDL_etree(An, Ap, Ai, iwork, Lnz, etree);

    /*--------------------------------
     * LDL factorisation
     *---------------------------------*/

    // First allocate memory for Li and Lx
    Li = (QDLDL_int*) malloc(sizeof(QDLDL_int) * sumLnz);
    Lx = (QDLDL_float*) malloc(sizeof(QDLDL_float) * sumLnz);

    // now factor
    QDLDL_factor(An, Ap, Ai, Ax, Lp, Li, Lx, D, Dinv, Lnz, etree, bwork, iwork, fwork);

    /*--------------------------------
     * solve
     *---------------------------------*/
    x = (QDLDL_float*) malloc(sizeof(QDLDL_float) * An);

    // when solving A\b, start with x = b
    for(i = 0; i < Ln; i++) {
        x[i] = b[i];
    }
    QDLDL_solve(Ln, Lp, Li, Lx, Dinv, x);

    /*--------------------------------
     * print factors and solution
     *---------------------------------*/
    printf("\n");
    printf("A (CSC format):\n");
    print_line();
    print_arrayi(Ap, An + 1, "A.p");
    print_arrayi(Ai, Ap[An], "A.i");
    print_arrayf(Ax, Ap[An], "A.x");
    printf("\n\n");

    printf("elimination tree:\n");
    print_line();
    print_arrayi(etree, Ln, "etree");
    print_arrayi(Lnz, Ln, "Lnz");
    printf("\n\n");

    printf("L (CSC format):\n");
    print_line();
    print_arrayi(Lp, Ln + 1, "L.p");
    print_arrayi(Li, Lp[Ln], "L.i");
    print_arrayf(Lx, Lp[Ln], "L.x");
    printf("\n\n");

    printf("D:\n");
    print_line();
    print_arrayf(D, An, "diag(D)     ");
    print_arrayf(Dinv, An, "diag(D^{-1})");
    printf("\n\n");

    printf("solve results:\n");
    print_line();
    print_arrayf(b, An, "b");
    print_arrayf(x, An, "A\\b");
    printf("\n\n");

    /*--------------------------------
     * clean up
     *---------------------------------*/
    free(Lp);
    free(Li);
    free(Lx);
    free(D);
    free(Dinv);
    free(etree);
    free(Lnz);
    free(iwork);
    free(bwork);
    free(fwork);
    free(x);

    return 0;
}

void print_line(void) {
    printf("--------------------------\n");
}

void print_arrayi(const QDLDL_int* data, QDLDL_int n, char* varName) {
    QDLDL_int i;
    printf("%s = [", varName);

    for(i = 0; i < n; i++) {
        printf("%i,", (int) data[i]);
    }
    printf("]\n");
}

void print_arrayf(const QDLDL_float* data, QDLDL_int n, char* varName) {
    QDLDL_int i;
    printf("%s = [", varName);

    for(i = 0; i < n; i++) {
        printf("%.3g,", data[i]);
    }
    printf("]\n");
}
