// Copyright (C) 2006  Paul A. Price (price@ifa.hawaii.edu)
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <assert.h>
#include "pslib.h"
#include "pois.h"

psVector *poisSolveEquation(psVector *solution, // Solution vector, or NULL
                            const psArray *stamps, // Array of stamps
                            const poisConfig *config
    )
{
    assert(stamps);
    assert(config);
    assert(!solution || solution->type.type == PS_TYPE_F64);

    int size = (2 * config->xKernel + 1) * (2 * config->yKernel + 1) * (config->spatialOrder + 1) *
        (config->spatialOrder + 2) / 2 + 1; // Size of matrix and vector

    assert(!solution || solution->n == size);

    psImage *matrix = psImageAlloc(size, size, PS_TYPE_F64);
    psVector *vector = psVectorAlloc(size, PS_TYPE_F64);
    for (int j = 0; j < size; j++) {
        for (int i = 0; i < size; i++) {
            matrix->data.F64[j][i] = 0.0;
        }
        vector->data.F64[j] = 0.0;
    }

    for (int s = 0; s < stamps->n; s++) {
        poisStamp *stamp = stamps->data[s]; // Stamp of interest
        psImage *stampMatrix = stamp->matrix; // Corresponding matrix
        psVector *stampVector = stamp->vector; // Corresponding vector

        // Check inputs
        assert(matrix->type.type == PS_TYPE_F64);
        assert(vector->type.type == PS_TYPE_F64);
        assert(matrix->numCols == size && matrix->numRows == size);
        assert(vector->n == size);

        if (stamp->status == POIS_STAMP_USED) {
            (void)psBinaryOp(matrix, matrix, "+", stampMatrix);
            (void)psBinaryOp(vector, vector, "+", stampVector);
        }
    }

    // Solution via LU Decomposition
    psVector *permutation = NULL; // Permutation vector for LU decomposition
    psImage *luMatrix = psMatrixLUD(NULL, &permutation, matrix); // LU decomposed matrix
    solution = psMatrixLUSolve(solution, luMatrix, vector, permutation); // Solution in x
    psTrace("pois.solveEquation", 1, "Background difference is %f\n", solution->data.F64[solution->n - 1]);

#if 0
    printf("Matrix:\n");
    for (int i = 0; i < matrix->numCols; i++) {
        for (int j = 0; j < matrix->numRows; j++) {
            printf("%f ", matrix->data.F64[i][j]);
        }
        printf("\n");
    }
#endif

    // Clean up
    psFree(luMatrix);
    psFree(permutation);
    psFree(matrix);
    psFree(vector);

    return solution;
}
