#include "psphotInternal.h"


typedef struct {
    int     numTrials;
    psF32   fRmajorMin;
    psF32   fRmajorMax;
    psF32   fRmajorDel;
    psF32   fRminorMin;
    psF32   fRminorMax;
    psF32   fRminorDel;
    psVector    *rMajor;
    psVector    *rMinor;
    psArray *zeroPt;
    psArray *exptime;
    psArray *cffTable;
} galaxyShapeOptions;

static pmSource *psphotFullForceSummarizeObject(pmConfig *config, pmPhotObj *obj, psVector *fluxScaleFactor, galaxyShapeOptions *options);
static pmPhotObj *findObjectForSource(psArray **pObjects, pmSource *source);

static bool setOptions(galaxyShapeOptions *options, pmReadout *readout, psMetadata *recipe, bool saveVectors);
static bool checkOptions(galaxyShapeOptions *options, pmReadout *readout, psMetadata *recipe);


bool psphotFullForceSummaryReadout (pmConfig *config, const pmFPAview *view) {

    bool status;
    psMetadata *recipe  = psMetadataLookupPtr (&status, config->recipes, PSPHOT_RECIPE);
    psAssert (recipe, "missing recipe?");

    int num = psphotFileruleCount(config, "PSPHOT.INPUT.CMF");

    pmFPAfile *output = psMetadataLookupPtr (&status, config->files, "PSPHOT.FULLFORCE.OUTPUT");
    pmCell *outputCell = pmFPAviewThisCell(view, output->fpa);
    pmReadout *outputReadout = pmFPAviewThisReadout(view, output->fpa);
    if (!outputReadout) {
        outputReadout = pmReadoutAlloc(outputCell);
    }

    // Get the exposure parameters for the output from recipe and set them on the output
    psF32 outputZeroPoint = psMetadataLookupF32(&status, recipe, "PSPHOT.FULLFORCE.ZERO_PT");
    psF32 outputExptime = psMetadataLookupF32(&status, recipe, "PSPHOT.FULLFORCE.EXPTIME");

    psMetadataAddF32(output->fpa->concepts, PS_LIST_TAIL, "FPA.ZP", PS_META_REPLACE, "Magnitude zero point",
        outputZeroPoint);
    psMetadataAddF32(outputCell->concepts, PS_LIST_TAIL, "CELL.EXPOSURE", PS_META_REPLACE, "Exposure time (sec)",
        outputExptime);
    // don't think this one matters
//    psMetadataAddF32(output->fpa->concepts, PS_LIST_TAIL, "FPA.EXPOSURE", PS_META_REPLACE, "Exposure time (sec)",
 //       outputExptime);

    // Create objects from the various input's sources
    // loop over the available readouts
//    psArray *readouts = psArrayAlloc(num);
    psVector *fluxScaleFactor = psVectorAlloc(num, PS_TYPE_F32);
    galaxyShapeOptions options;
    psArray *objects = NULL;
    for (int index = 0; index < num; index++) {
        // find the currently selected readout
        pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PSPHOT.INPUT.CMF", index); // File of interest
        psAssert (file, "missing file?");

        pmReadout *readout = pmFPAviewThisReadout(view, file->fpa);
        psAssert (readout, "missing readout?");

        if (index == 0) {
            // Get the galaxy shape recipe values, from the analysis if present
            // or from the recipe if not
            if (!setOptions(&options, readout, recipe, true)) {
                psError (PS_ERR_UNKNOWN, false, "problem determining galaxy shape options.");
                return false;
            }
            options.cffTable = psMetadataLookupPtr(NULL, config->arguments, "CFF_TABLE");
            if (!options.cffTable) {
                psError (PS_ERR_UNKNOWN, true, "Cannot find cff table in arguments.");
                return false;
            }
        } else { 
            // Make sure that this input was created with the same galaxy shapes recipe
            if (!checkOptions(&options, readout, recipe)) {
                psError (PS_ERR_UNKNOWN, false, "galaxy shape options do not match for input %d", index);
                return false;
            }
        }

        // look up zero point
        psF32 zero_point = psMetadataLookupF32(&status, readout->parent->parent->parent->concepts, "FPA.ZP");
        psF32 exptime = psMetadataLookupF32(&status, readout->parent->parent->parent->concepts, "FPA.EXPOSURE");

        psF32 scaleFactor = pow(10, 0.4 * (outputZeroPoint - zero_point)) * outputExptime / exptime;
        fluxScaleFactor->data.F32[index] = scaleFactor;

//        readouts->data[index] = readout;

        // find detections
        pmDetections *detections = psMetadataLookupPtr (&status, readout->analysis, "PSPHOT.DETECTIONS");
        psAssert (detections, "missing detections?");
        psArray *sources = detections->allSources;
        psAssert (sources, "missing sources?");
        sources = psArraySort (sources, pmSourceSortBySeq);

        if (objects == NULL) {
            pmSource *lastSource = sources->data[sources->n - 1];
            psAssert(lastSource, "last source is null!");
            objects = psArrayAlloc(lastSource->seq + 1);
        }

        for (int i=0; i < sources->n; i++) {
            pmSource *source = sources->data[i];
            source->imageID = index;
            findObjectForSource(&objects, source);
        }
    }

    pmDetections *outputDetections = pmDetectionsAlloc();
    if (!psMetadataAddPtr (outputReadout->analysis, PS_LIST_TAIL, "PSPHOT.DETECTIONS", 
            PS_META_REPLACE | PS_DATA_UNKNOWN, "psphot detections", outputDetections)) {
        psError (PSPHOT_ERR_CONFIG, false, "problem saving detections on readout");
        return false;
    }

    psArray *outputSources = psArrayAllocEmpty (objects->n);

    // Loop over objects and compute the summaries
    long nObjects = 0;
    long nSources = 0;
    for (int i=0 ; i<objects->n; i++) {
        pmPhotObj *obj = objects->data[i];
        if (!obj) continue;

        ++nObjects;
        pmSource *source = psphotFullForceSummarizeObject(config, obj, fluxScaleFactor, &options);
        if (source) {
            psArrayAdd (outputSources, 100, source);
            ++nSources;
        }
    }

    psLogMsg("psphot", PS_LOG_INFO, "constructed %ld output sources, from %ld objects.", nSources, nObjects);

    psFree(fluxScaleFactor);

    if (nSources) {
        // We have data
        outputDetections->allSources = outputSources;
        outputReadout->data_exists = true;
        outputReadout->parent->data_exists = true;
        outputReadout->parent->parent->data_exists = true;
    } else {
        // XXX: tooo set a quality or fault code
        return false;
    }

    return true;
}

static pmPhotObj *findObjectForSource(psArray **pObjects, pmSource *source) {
    int seq = source->seq;

    psArray *objects = *pObjects;
    if (seq >= objects->n) {
        // We need to expand the object array. Kind of suprising.
        objects = *pObjects = psArrayRealloc(objects, seq+1);
    }

    // Look up object for this seq
    pmPhotObj *obj = objects->data[seq];
    if (!obj) {
        // not found allocate one
        obj = pmPhotObjAlloc();
        objects->data[seq] = obj;
    }
    pmPhotObjAddSource(obj, source);

    return obj;
}

static pmSource *psphotFullForceSummarizeObject(pmConfig *config, pmPhotObj *obj, psVector *fluxScaleFactor, galaxyShapeOptions *options) {

    pmSource *outSrc = NULL;

    psVector *sumWeightedFlux = NULL;
    psVector *sumInvSig2 = NULL;
    psVector *numerator   = NULL;
//    psVector *tmp = NULL;
    psF32   totalNPix = 0;
    long    vectorLength = 0;

    for (int i=0; i < obj->sources->n; i++) {
        pmSource *source = obj->sources->data[i];

        // XXX: get cut from recipe
        if (source->pixWeightNotPoor < .9) continue;

        // For now just start the output source as a copy of the first input source that makes cuts
        if (!outSrc) {
            outSrc = pmSourceCopy(source);
            // This copies 
            //  the peak 
            //  the moments which are mostly nan except for Mrf Mx, My, and some of the kron parameters
            //
            // type, mode, flags
            // magnitudes
            outSrc->imageID = 0;
            outSrc->seq = source->seq;

            if (source->modelPSF) {
                outSrc->modelPSF = psMemIncrRefCounter(source->modelPSF);
            }
            if (source->extpars) {
                outSrc->extpars =  psMemIncrRefCounter(source->extpars);
            }
            if (source->modelEXT) {
                outSrc->modelEXT =  psMemIncrRefCounter(source->modelEXT);
            }
            if (source->modelFits) {
                outSrc->modelFits =  psMemIncrRefCounter(source->modelFits);
            }
        }

        if (source->galaxyFits && isfinite(source->galaxyFits->nPix) && source->galaxyFits->chisq->n) {
            if (numerator == NULL) {
                vectorLength = source->galaxyFits->chisq->n;
                // tmp = psVectorAlloc(vectorLength, PS_TYPE_F32);
                sumWeightedFlux = psVectorAlloc(vectorLength, PS_TYPE_F32);
                psVectorInit(sumWeightedFlux, 0.0);
                sumInvSig2 = psVectorAlloc(vectorLength, PS_TYPE_F32);
                psVectorInit(sumInvSig2, 0.0);
                numerator   = psVectorAlloc(vectorLength, PS_TYPE_F32);
                psVectorInit(numerator, 0.0);
            }

            psAssert(vectorLength == options->numTrials, "length of chisq vector %ld does not match options %d",
                vectorLength, options->numTrials);
            psAssert(source->galaxyFits->chisq->n == vectorLength, "length of chisq vectors do not match %ld %ld",
                    source->galaxyFits->chisq->n, vectorLength);


            psF32 scaleFactor = fluxScaleFactor->data.F32[source->imageID];

            for (int k = 0; k < vectorLength; k++) {

                psF32 chisq = source->galaxyFits->chisq->data.F32[k];
                psF32 flux  = source->galaxyFits->Flux->data.F32[k]  * scaleFactor;
                psF32 dFlux = source->galaxyFits->dFlux->data.F32[k] * scaleFactor;
                // for chisq 
                // Numerator = sum( nPix[i] * chisq[i] )
                // denominator will be sum( nPix )


//                tmp = (psVector *) psBinaryOp(tmp, 
 //                   source->galaxyFits->chisq, "*", psScalarAlloc(source->galaxyFits->nPix, PS_TYPE_F32));
//                numerator = (psVector *) psBinaryOp(numerator, numerator, "+", tmp);

                numerator->data.F32[k] += chisq * source->galaxyFits->nPix;
                totalNPix += source->galaxyFits->nPix;

                // sumInvSig2 = sum( 1 / dFlux**2 )
//                tmp = (psVector *) psBinaryOp(tmp, source->galaxyFits->dFlux, "*", source->galaxyFits->dFlux);
//                tmp = (psVector *) psBinaryOp(tmp, psScalarAlloc(1.0, PS_TYPE_F32), "/", tmp);
//                sumInvSig2 = (psVector *) psBinaryOp(sumInvSig2, sumInvSig2, "+", tmp);
//
                psF32 invSig2 = 1.0 / (dFlux * dFlux);
                sumInvSig2->data.F32[k] += invSig2;

                // sumWeightedFlux = sum ( Flux / dFlux**2 )
//                tmp = (psVector *) psBinaryOp(tmp, source->galaxyFits->Flux, "*", tmp);
//                sumWeightedFlux = (psVector *) psBinaryOp(sumWeightedFlux, sumWeightedFlux, "+", tmp);
                sumWeightedFlux->data.F32[k] += flux * invSig2;
            }
        }
    }

    if (outSrc) {
        if (vectorLength) {
            outSrc->galaxyFits = pmSourceGalaxyFitsAlloc();
            outSrc->galaxyFits->nPix = totalNPix;
            outSrc->galaxyFits->Flux  = psVectorRecycle(outSrc->galaxyFits->Flux, vectorLength, PS_TYPE_F32);
            outSrc->galaxyFits->dFlux = psVectorRecycle(outSrc->galaxyFits->dFlux, vectorLength, PS_TYPE_F32);
            outSrc->galaxyFits->chisq = psVectorRecycle(outSrc->galaxyFits->chisq, vectorLength, PS_TYPE_F32);
            for (int k = 0; k < vectorLength; k++) {
                outSrc->galaxyFits->Flux->data.F32[k]  = sumWeightedFlux->data.F32[k] / sumInvSig2->data.F32[k];
                outSrc->galaxyFits->dFlux->data.F32[k] = 1.0 / sumInvSig2->data.F32[k];
                outSrc->galaxyFits->chisq->data.F32[k] = numerator->data.F32[k] / totalNPix;

#ifdef nodef
                outSrc->galaxyFits->Flux = (psVector *) psBinaryOp(outSrc->galaxyFits->Flux,
                        sumWeightedFlux, "/", sumInvSig2);

                outSrc->galaxyFits->dFlux = (psVector *) psBinaryOp(outSrc->galaxyFits->dFlux,
                        psScalarAlloc(1.0, PS_TYPE_F32), "/", sumInvSig2);

                outSrc->galaxyFits->chisq = (psVector *) psBinaryOp(outSrc->galaxyFits->chisq,
                        numerator, "/", psScalarAlloc(totalNPix, PS_TYPE_F32));
#endif
            }

            psFree(numerator);
            psFree(sumInvSig2);
            psFree(sumWeightedFlux);
//            psFree(tmp);

            int min_j = -1;
            psF32 minChisq = NAN;
            psVector *chisq = outSrc->galaxyFits->chisq;
            for (int j=0; j < chisq->n; j++) {
                psF32 thischisq = chisq->data.F32[j];
                if (isfinite(thischisq)  && (!isfinite(minChisq) || thischisq < minChisq)) {
                    min_j = j;
                    minChisq = thischisq;
                }
            }
            if (min_j >= 0 && isfinite(minChisq) && outSrc->modelEXT) {
                // copy the best fit params to the model
                psEllipseAxes axes = pmPSF_ModelToAxes(outSrc->modelEXT->params->data.F32, outSrc->modelEXT->type);
                if (outSrc->seq >= options->cffTable->n) {
                    psError(PS_ERR_PROGRAMMING, true, "object sequence number %d is larger than cff table length %ld",
                        outSrc->seq, options->cffTable->n);
                    return NULL;
                }
                psMetadata *row = options->cffTable->data[outSrc->seq];
                if (!row) {
                    psError (PS_ERR_PROGRAMMING, true, "NO cff data for object %d", outSrc->seq);
                    psFree(outSrc);
                    return NULL;
                }
                bool mdok;
                psF32 rMajor = psMetadataLookupF32(&mdok, row, "R_MAJOR");
                if (!mdok) {
                    psError (PS_ERR_PROGRAMMING, true, "can't find R_MAJOR for object %d", outSrc->seq);
                    psFree(outSrc);
                    return NULL;
                }
                psF32 rMinor = psMetadataLookupF32(&mdok, row, "R_MINOR");
                if (!mdok) {
                    psError (PS_ERR_PROGRAMMING, true, "can't find R_MINOR for object %d", outSrc->seq);
                    psFree(outSrc);
                    return NULL;
                }
                axes.major = rMajor * options->rMajor->data.F32[min_j];
                axes.minor = rMinor * options->rMinor->data.F32[min_j];
                pmPSF_AxesToModel (outSrc->modelEXT->params->data.F32, axes, outSrc->modelEXT->type);
                outSrc->modelEXT->chisq = minChisq; 
                outSrc->modelEXT->mag = -2.5 * log10(outSrc->galaxyFits->Flux->data.F32[min_j]);
                outSrc->modelEXT->magErr = 1.0 / 
                    (outSrc->galaxyFits->Flux->data.F32[min_j]/outSrc->galaxyFits->dFlux->data.F32[min_j]); // 1 / SN
            }
        }
    }

    return outSrc;
}

#define GETVAL(member, key) \
    opt->member = psMetadataLookupF32(&status, md, key); \
    if (!status) { \
        psError (PSPHOT_ERR_CONFIG, true, "failed to looup value for %s in %s", key, \
            useAnalysis ? "readout->analysis" : "recipe"); \
    }

static bool setOptions(galaxyShapeOptions *opt, pmReadout *readout, psMetadata *recipe, bool makeVectors) {
    bool status;
    bool useAnalysis;

    psMetadataLookupF32(&useAnalysis, readout->analysis, "GALAXY_SHAPES_FR_MAJOR_MIN");
    psMetadata *md = useAnalysis ? readout->analysis : recipe;

    GETVAL(fRmajorMin, "GALAXY_SHAPES_FR_MAJOR_MIN");
    GETVAL(fRmajorMax, "GALAXY_SHAPES_FR_MAJOR_MAX");
    GETVAL(fRmajorDel, "GALAXY_SHAPES_FR_MAJOR_DEL");
    GETVAL(fRminorMin, "GALAXY_SHAPES_FR_MINOR_MIN");
    GETVAL(fRminorMax, "GALAXY_SHAPES_FR_MINOR_MAX");
    GETVAL(fRminorDel, "GALAXY_SHAPES_FR_MINOR_DEL");

    opt->numTrials = ceil((opt->fRmajorMax - opt->fRmajorMin + 0.5*opt->fRmajorDel) / opt->fRmajorDel) *
                         ceil((opt->fRminorMax - opt->fRminorMin + 0.5*opt->fRminorDel) / opt->fRminorDel) ;

    if (makeVectors) {
        opt->rMinor = psVectorAlloc(opt->numTrials, PS_TYPE_F32);
        opt->rMajor = psVectorAlloc(opt->numTrials, PS_TYPE_F32);
        int i = 0;
        for (float fRmajor = opt->fRmajorMin; fRmajor < opt->fRmajorMax + 0.5*opt->fRmajorDel;
                fRmajor += opt->fRmajorDel) {
            for (float fRminor = opt->fRminorMin; fRminor < opt->fRminorMax + 0.5*opt->fRminorDel;
                    fRminor += opt->fRminorDel) {
                opt->rMinor->data.F32[i] = fRminor;
                opt->rMajor->data.F32[i] = fRmajor;
                i++;
            }
        }
        psAssert(i == opt->numTrials, "Something's wrong with my loop got %d entries expected %d", i, opt->numTrials);
    } else {
        opt->rMinor = NULL;
        opt->rMajor = NULL;
    }
        
    return true;
}

#define CHECKVAL(left, right, val, message) \
    if (left->val != right.val) { \
        psError (PSPHOT_ERR_CONFIG, true, message); \
        return false; \
    }

static bool checkOptions(galaxyShapeOptions *options, pmReadout *readout, psMetadata *recipe) {
    galaxyShapeOptions thisReadoutsOptions;
    
    if (!setOptions(&thisReadoutsOptions, readout, recipe, false)) {
        psError (PS_ERR_UNKNOWN, false, "problem determining galaxy shape options for readout");
        return false;
    }
    CHECKVAL(options, thisReadoutsOptions, numTrials, "mismatched number of trials")
    CHECKVAL(options, thisReadoutsOptions, fRmajorMin, "mismatched fRmajorMin")
    CHECKVAL(options, thisReadoutsOptions, fRmajorMax, "mismatched fRmajorMax")
    CHECKVAL(options, thisReadoutsOptions, fRmajorDel, "mismatched fRmajorDel")
    CHECKVAL(options, thisReadoutsOptions, fRminorMin, "mismatched fRminorMin")
    CHECKVAL(options, thisReadoutsOptions, fRminorMax, "mismatched fRminorMax")
    CHECKVAL(options, thisReadoutsOptions, fRminorDel, "mismatched fRminorDel")

    return true;
}
