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

#include <stdio.h>
#include <stdlib.h>
#include <limits.h>
#include <string.h>

#include "psAbort.h"
#include "psMemory.h"
#include "psConstants.h"
#include "psImageConvolve.h"
#include "psTrace.h"
#include "psBinaryOp.h"

#include "psImageCovariance.h"

psKernel *psImageCovarianceNone(void)
{
    psKernel *covar = psKernelAlloc(0, 0, 0, 0); // Covariance pseudo-matrix
    covar->kernel[0][0] = 1.0;
    return covar;
}


psKernel *psImageCovarianceCalculate(const psKernel *kernel, const psKernel *covariance)
{
    PS_ASSERT_KERNEL_NON_NULL(kernel, NULL);

    // See http://en.wikipedia.org/wiki/Error_propagation
    //
    // If
    //     f_k = sum_i A_ik x_i
    // is a set of functions, then the covariance matrix for f is given by:
    //     M^f_ij = sum_k sum_l A_ik M^x_kl A_lj
    // where M^x is the covariance matrix for x.
    // Note that the errors in f are correlated (covariance) even if the errors in x are not.

    psKernel *covar;                    // Covariance matrix to use
    if (covariance) {
        covar = psMemIncrRefCounter((psKernel*)covariance); // Casting away const
    } else {
        covar = psImageCovarianceNone();
    }

    // Check for non-finite elements
    for (int y = kernel->yMin; y <= kernel->yMax; y++) {
        for (int x = kernel->xMin; x <= kernel->xMax; x++) {
            if (!isfinite(kernel->kernel[y][x])) {
                psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                        "Non-finite covariance matrix element in kernel at %d,%d", x, y);
                psFree(covar);
                return NULL;
            }
        }
    }
    for (int y = covar->yMin; y <= covar->yMax; y++) {
        for (int x = covar->xMin; x <= covar->xMax; x++) {
            if (!isfinite(covar->kernel[y][x])) {
                psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                        "Non-finite covariance matrix element in covariance matrix at %d,%d", x, y);
                psFree(covar);
                return NULL;
            }
        }
    }

    // The above (double) sum for the covariance matrix means that, for each point in the output covariance
    // matrix, we need to work out all combinations of getting to the central point via a kernel, input
    // covariance matrix and another kernel.  This means that the resultant covariance matrix has twice the
    // size of the kernel plus the size of the input covariance matrix.
    int xMin = kernel->xMin - kernel->xMax + covar->xMin, xMax = kernel->xMax - kernel->xMin + covar->xMax;
    int yMin = kernel->yMin - kernel->yMax + covar->yMin, yMax = kernel->yMax - kernel->yMin + covar->yMax;
    psKernel *out = psKernelAlloc(xMin, xMax, yMin, yMax); // Covariance matrix for output

    // Need to go:
    // (x,y) --kernel--> (u,v) --covar--> (p,q) --kernel--> (0,0)
    // All coordinates (x,y), (u,v), and (p,q) are "absolute", meaning that they are in x,y space, which is
    // the space of the output covariance matrix.
    //
    // So, the ranges in the coordinates are:
    // (x,y): -outSize:outSize
    // (u,v): (x,y)-kernelSize:(x,y)+kernelSize  AND  -kernelSize-inSize:kernelSize+inSize
    // (p,q): -kernelSize:kernelSize  AND  (u,v)-inSize:(u,v)+inSize
    // Here outSize is the size of the output covariance matrix, kernelSize is the size of the convolution
    // kernel, and inSize is the size of the input covariance matrix.
    // Since (u,v) and (p,q) have two ways of specifying the range (one from the target coordinate and one
    // from the source coordinate), we take the smallest possible (because everything else is zero outside).

    double total = 0.0;                 // Total covariance
    for (int y = yMin; y <= yMax; y++) {
        // Range for v
        int vMin = PS_MAX(kernel->yMin + covar->yMin, y + kernel->yMin);
        int vMax = PS_MIN(kernel->yMax + covar->yMax, y + kernel->yMax);
        for (int x = xMin; x <= xMax; x++) {
            // Range for u
            int uMin = PS_MAX(kernel->xMin + covar->xMin, x + kernel->xMin);
            int uMax = PS_MIN(kernel->xMax + covar->xMax, x + kernel->xMax);

            double sum = 0.0;           // Sum for value of covariance matrix at (x,y)
            for (int v = vMin; v <= vMax; v++) {
                // Range for q
                int qMin = PS_MAX(v + covar->yMin, kernel->yMin);
                int qMax = PS_MIN(v + covar->yMax, kernel->yMax);
                for (int u = uMin; u <= uMax; u++) {
                    // Range for p
                    int pMin = PS_MAX(u + covar->xMin, kernel->xMin);
                    int pMax = PS_MIN(u + covar->xMax, kernel->xMax);

                    double xyuvValue = kernel->kernel[v-y][u-x]; // Value for (x,y) --> (u,v)

                    double uvpqValue = 0.0; // Value for (u,v) --> (p,q) --> (0,0)
                    for (int q = qMin; q <= qMax; q++) {
                        for (int p = pMin; p <= pMax; p++) {
                            uvpqValue += (double)covar->kernel[q-v][p-u] * (double)kernel->kernel[q][p];
                        }
                    }
                    sum += xyuvValue * uvpqValue;
                }
            }
            out->kernel[y][x] = sum;
            total += sum;
        }
    }
    psTrace("psLib.imageops", 3, "Total covariance: %lf ; Central variance: %f\n", total, out->kernel[0][0]);

    psFree(covar);
    return out;
}

float psImageCovarianceFactor(const psKernel *covariance)
{
    return covariance ? covariance->kernel[0][0] : 1.0;
}

psKernel *psImageCovarianceSum(const psArray *array)
{
    PS_ASSERT_ARRAY_NON_NULL(array, NULL);
    PS_ASSERT_ARRAY_NON_EMPTY(array, NULL);

    int xMin = INT_MAX, xMax = INT_MIN, yMin = INT_MAX, yMax = INT_MIN; // Range for covariance
    int num = 0;                        // Number of good matrices to sum
    for (int i = 0; i < array->n; i++) {
        psKernel *covar = array->data[i]; // Covariance matrix
        if (!covar) {
            continue;
        }
        xMin = PS_MIN(xMin, covar->xMin);
        xMax = PS_MAX(xMax, covar->xMax);
        yMin = PS_MIN(yMin, covar->yMin);
        yMax = PS_MAX(yMax, covar->yMax);
        num++;
    }
    if (num == 0) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true, "No covariance matrices supplied for summation");
        return NULL;
    }

    psKernel *sum = psKernelAlloc(xMin, xMax, yMin, yMax); // Summed covariance
    for (int i = 0; i < array->n; i++) {
        psKernel *covar = array->data[i]; // Covariance matrix
        if (!covar) {
            continue;
        }
        for (int y = covar->yMin; y <= covar->yMax; y++) {
            for (int x = covar->xMin; x <= covar->xMax; x++) {
                if (!isfinite(covar->kernel[y][x])) {
                    psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                            "Non-finite covariance matrix element at %d,%d for input %d",
                            x, y, i);
                    psFree(sum);
                    return NULL;
                }
                sum->kernel[y][x] += covar->kernel[y][x];
            }
        }
    }

    return sum;
}


psKernel *psImageCovarianceAverage(const psArray *array)
{
    PS_ASSERT_ARRAY_NON_NULL(array, NULL);
    PS_ASSERT_ARRAY_NON_EMPTY(array, NULL);

    int num = 0;                        // Number of good matrices to average
    for (int i = 0; i < array->n; i++) {
        psKernel *covar = array->data[i]; // Covariance matrix
        if (covar) {
            num++;
        }
    }
    if (num == 0) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true, "No covariance matrices supplied for averaging.");
        return NULL;
    }

    psKernel *sum = psImageCovarianceSum(array); // Sum of covariances
    psBinaryOp(sum->image, sum->image, "/", psScalarAlloc(num, PS_TYPE_F32));

    return sum;
}


psKernel *psImageCovarianceTruncate(const psKernel *covar, float frac)
{
    PS_ASSERT_KERNEL_NON_NULL(covar, NULL);
    PS_ASSERT_FLOAT_WITHIN_RANGE(frac, 0, 1, NULL);

    int xMin = covar->xMin, xMax = covar->xMax, yMin = covar->yMin, yMax = covar->yMax; // Range
    int maxRadius = PS_MAX(PS_MAX(PS_MAX(xMax, yMax), -xMin), -yMin); // Maximum radius of covariance matrix

    double sum = 0.0;                   // Sum of covariance
    psVector *radiusSum = psVectorAlloc(maxRadius + 1, PS_TYPE_F64); // Totals within (square) radius
    psVectorInit(radiusSum, 0.0);
    for (int y = yMin; y <= yMax; y++) {
        for (int x = xMin; x <= xMax; x++) {
            int radius = PS_MAX(abs(x), abs(y)); // Squarish radius
            psAssert(radius <= maxRadius, "Radius doesn't fit");
            if (!isfinite(covar->kernel[y][x])) {
                psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Non-finite covariance matrix element at %d,%d",
                        x, y);
                return NULL;
            }
            radiusSum->data.F64[radius] += fabsf(covar->kernel[y][x]);
            sum += fabsf(covar->kernel[y][x]);
        }
    }

    // Determine radius for truncation
    double threshold = (1.0 - frac) * sum; // Threshold for truncation
    double enclosed = 0.0;              // Enclosed value
    int radius;                         // Radius at which to truncate
    for (radius = 0; radius <= maxRadius && enclosed < threshold; radius++) {
        enclosed += radiusSum->data.F64[radius];
    }
    psFree(radiusSum);

    if (radius >= maxRadius) {
        radius = maxRadius;
    }

    // Generate truncated version
    int xMinNew = PS_MAX(xMin, -radius), xMaxNew = PS_MIN(xMax, radius); // New range in x
    int yMinNew = PS_MAX(yMin, -radius), yMaxNew = PS_MIN(yMax, radius); // New range in y
    psKernel *trunc = psKernelAlloc(xMinNew, xMaxNew, yMinNew, yMaxNew); // Truncated covariance matrix
    int numBytes = (xMaxNew - xMinNew + 1) * PSELEMTYPE_SIZEOF(PS_TYPE_F32); // Number of bytes to copy
    for (int y = yMinNew; y <= yMaxNew; y++) {
        memcpy(&trunc->kernel[y][xMinNew], &covar->kernel[y][xMinNew], numBytes);
    }

    return trunc;
}


bool psImageCovarianceTransfer(psImage *variance, psKernel *covar)
{
    PS_ASSERT_IMAGE_NON_NULL(variance, NULL);
    PS_ASSERT_KERNEL_NON_NULL(covar, NULL);

    float factor = psImageCovarianceFactor(covar); // Factor to transfer
    if (factor == 1.0) {
        // No work required
        return true;
    }

    psBinaryOp(variance, variance, "*", psScalarAlloc(factor, PS_TYPE_F32));
    psBinaryOp(covar->image, covar->image, "/", psScalarAlloc(factor, PS_TYPE_F32));

    return true;
}
