#include <stdio.h>
#include <assert.h>
#include <math.h>
#include <pslib.h>

#include "pmFlatNormalize.h"


// Estimate the flat-field normalisation; return the source flux in each integration
psVector *pmFlatNormalize(bool *converge, // Did we converge?
                          psVector *chipGains, // Initial guess of the chip gains; modified for return
                          const psImage *fluxLevels, // Fluxes for each integration (row) and chip (col)
                          unsigned int maxIter, // Maximum number of iterations
                          float tolerance   // Tolerance level before dying
                         )
{
    PS_ASSERT_PTR_NON_NULL(chipGains, NULL);
    PS_ASSERT_PTR_NON_NULL(fluxLevels, NULL);

    int numSources = fluxLevels->numRows; // Number of integrations
    int numChips = fluxLevels->numCols; // Number of chips with which each integration is made

    // Sanity checks
    assert(chipGains->n == numChips);
    assert(chipGains->type.type == PS_TYPE_F32);
    assert(maxIter >= 1);
    assert(tolerance > 0);

    // Take the logarithms
    psImage *flux = psImageCopy(NULL, fluxLevels, PS_TYPE_F32); // Copy of the input flux levels matrix
    psImage *fluxMask = psImageAlloc(numChips, numSources, PS_TYPE_U8); // Mask for bad measurements
    psImageInit(fluxMask, 0);
    psVector *gainMask = psVectorAlloc(numChips, PS_TYPE_U8); // Mask for bad gains
    gainMask->n = numChips;
    psVectorInit(gainMask, 0);
    psVector *sourceMask = psVectorAlloc(numSources, PS_TYPE_U8); // Mask for bad integrations
    sourceMask->n = numSources;
    psVectorInit(sourceMask, 0);
    for (int i = 0; i < numChips; i++) {
        if (isfinite(chipGains->data.F32[i]) && chipGains->data.F32[i] > 0) {
            chipGains->data.F32[i] = log(chipGains->data.F32[i]);
        } else {
            chipGains->data.F32[i] = 0.0; // Take a wild guess
            #if 0
            // Blank out this chip
            gainMask->data.U8[i] = 1;
            chipGains->data.F32[i] = NAN;
            #endif

        }

        for (int j = 0; j < numSources; j++) {
            if (isfinite(flux->data.F32[j][i]) && flux->data.F32[j][i] > 0) {
                flux->data.F32[j][i] = log(flux->data.F32[j][i]);
            } else {
                // Blank out this measurement
                fluxMask->data.U8[j][i] = 1;
                flux->data.F32[j][i] = NAN;
            }
        }
    }

    float diff = INFINITY;             // Difference from previous iteration
    psVector *sourceFlux = psVectorAlloc(numSources, PS_TYPE_F32); // The flux in each integration
    sourceFlux->n = numSources;
    psVector *oldSourceFlux = psVectorAlloc(numSources, PS_TYPE_F32); // The fluxes in the previous iteration
    oldSourceFlux->n = numSources;
    psVectorInit(oldSourceFlux, 0.0);
    for (int iter = 0; iter < maxIter && diff > tolerance; iter++) {
        diff = 0.0;


        // Improve on the fluxes
        float sumFlux = 0.0;           // Total fluxes
        long numFluxes = 0;             // Number of fluxes
        for (int i = 0; i < numSources; i++) {
            if (sourceMask->data.U8[i]) {
                psTrace(__func__, 7, "Flux for exposure %d is masked.\n", i);
                continue;
            }
            numFluxes++;
            float sum = 0.0;           // Sum of F_ij - G_j
            int number = 0;             // Number of chips contributing
            for (int j = 0; j < numChips; j++) {
                if (!gainMask->data.U8[j] && !fluxMask->data.U8[i][j]) {
                    sum += flux->data.F32[i][j] - chipGains->data.F32[j];
                    number++;
                }
            }
            if (number > 0) {
                sourceFlux->data.F32[i] = sum / (float)number;
            } else {
                sourceMask->data.U8[i] = 1;
                sourceFlux->data.F32[i] = NAN;
            }
            sumFlux += exp(sourceFlux->data.F32[i]);
            psTrace(__func__, 7, "Flux for exposure %d is %f\n", i, exp(sourceFlux->data.F32[i]));
        }
        // Normalise the mean to unity
        sumFlux /= (float)numFluxes;
        sumFlux = log(sumFlux);
        for (int i = 0; i < numSources; i++) {
            if (sourceMask->data.U8[i]) {
                continue;
            }
            sourceFlux->data.F32[i] -= sumFlux;
            diff += abs((sourceFlux->data.F32[i] - oldSourceFlux->data.F32[i]) / sourceFlux->data.F32[i]);
        }

        // Improve on the gains
        for (int i = 0; i < numChips; i++) {
            if (gainMask->data.U8[i]) {
                continue;
            }
            float sum = 0.0;           // Sum of F_ji - S_j
            int number = 0;             // Numer of sources contributing
            for (int j = 0; j < numSources; j++) {
                if (!fluxMask->data.U8[j][i]) {
                    sum += flux->data.F32[j][i] - sourceFlux->data.F32[j];
                    number++;
                }
            }
            if (number > 0) {
                chipGains->data.F32[i] = sum / (float)number;
            } else {
                gainMask->data.U8[i] = 1;
                chipGains->data.F32[i] = NAN;
            }
            psTrace(__func__, 7, "Gain for chip %d is %f\n", i, exp(chipGains->data.F32[i]));
        }

        psTrace(__func__, 2, "Iteration %d: difference is %f\n", iter, diff);

        // Switch the old and new
        psVector *temp = sourceFlux;
        sourceFlux = oldSourceFlux;
        oldSourceFlux = temp;
    }
    psFree(flux);
    psFree(fluxMask);
    psFree(oldSourceFlux);

    // Un-log the vectors
    for (int i = 0; i < numChips; i++) {
        if (!gainMask->data.U8[i]) {
            chipGains->data.F32[i] = exp(chipGains->data.F32[i]);
        }
    }
    for (int i = 0; i < numSources; i++) {
        if (!sourceMask->data.U8[i]) {
            sourceFlux->data.F32[i] = exp(sourceFlux->data.F32[i]);
        }
    }
    psFree(gainMask);
    psFree(sourceMask);

    if (converge) {
        *converge = (diff < tolerance); // Did we converge?
    }

    return sourceFlux;
}
