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

#include <stdio.h>
#include <strings.h>
#include <string.h>
#include <pslib.h>

#include "pmHDU.h"
#include "pmFPA.h"
#include "pmFPALevel.h"
#include "pmHDUUtils.h"
#include "pmHDUGenerate.h"
#include "pmConcepts.h"

#include "pmFPAWrite.h"

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Definitions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Specify what to read
typedef enum {
    FPA_WRITE_TYPE_IMAGE,               // Write image
    FPA_WRITE_TYPE_MASK,                // Write mask
    FPA_WRITE_TYPE_WEIGHT               // Write weight map
} fpaWriteType;

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// File-static (private) functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////


// Update the FPA.NAME, CHIP.NAME and CELL.NAME in the FITS header, if required
bool pmFPAUpdateNames(pmFPA *fpa, pmChip *chip, pmCell *cell)
{
    pmHDU *hdu = pmHDUGetHighest(fpa, chip, cell); // Highest HDU, i.e., the PHU
    if (!hdu) {
        psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find HDU.\n");
        return false;
    }
    if (!hdu->header) {
        hdu->header = psMetadataAlloc();
    }
    bool mdok;                          // Status of MD lookup
    psMetadata *fileData = psMetadataLookupMetadata(&mdok, hdu->format, "FILE"); // File information
    if (!mdok || !fileData) {
        psError(PS_ERR_UNKNOWN, true, "Unable to find FILE information in camera format.\n");
        return false;
    }
    if (fpa) {
        const char *fpaNameHdr = psMetadataLookupStr(&mdok, fileData, "FPA.NAME");
        if (mdok && fpaNameHdr && strlen(fpaNameHdr) > 0) {
            const char *fpaName = psMetadataLookupStr(NULL, fpa->concepts, "FPA.NAME");
            psMetadataAddStr(hdu->header, PS_LIST_TAIL, fpaNameHdr, PS_META_REPLACE, "FPA name", fpaName);
        }
    }
    if (chip) {
        const char *chipNameHdr = psMetadataLookupStr(&mdok, fileData, "CHIP.NAME");
        if (mdok && chipNameHdr && strlen(chipNameHdr) > 0) {
            const char *chipName = psMetadataLookupStr(NULL, chip->concepts, "CHIP.NAME");
            psMetadataAddStr(hdu->header, PS_LIST_TAIL, chipNameHdr, PS_META_REPLACE, "Chip name", chipName);
        }
    }
    if (cell) {
        const char *cellNameHdr = psMetadataLookupStr(&mdok, fileData, "CELL.NAME");
        if (mdok && cellNameHdr && strlen(cellNameHdr) > 0) {
            const char *cellName = psMetadataLookupStr(NULL, cell->concepts, "CELL.NAME");
            psMetadataAddStr(hdu->header, PS_LIST_TAIL, cellNameHdr, PS_META_REPLACE, "Cell name", cellName);
        }
    }
    return true;
}

// Return the appropriate image array for the given type
static psArray **appropriateImageArray(pmHDU *hdu, // HDU containing the image arrays
                                       fpaWriteType type // Type to write
                                      )
{
    switch (type) {
    case FPA_WRITE_TYPE_IMAGE:
        return &hdu->images;
    case FPA_WRITE_TYPE_MASK:
        return &hdu->masks;
    case FPA_WRITE_TYPE_WEIGHT:
        return &hdu->weights;
    default:
        psAbort("Unknown write type: %x\n", type);
    }
    return NULL;
}

// Run the appropriate HDU write function
static bool appropriateWriteFunc(pmHDU *hdu, // HDU to write
                                 psFits *fits, // FITS file to which to write
                                 fpaWriteType type // Type to write
                                )
{
    switch (type) {
    case FPA_WRITE_TYPE_IMAGE:
        return pmHDUWrite(hdu, fits);
    case FPA_WRITE_TYPE_MASK:
        return pmHDUWriteMask(hdu, fits);
    case FPA_WRITE_TYPE_WEIGHT:
        return pmHDUWriteWeight(hdu, fits);
    default:
        psAbort("Unknown write type: %x\n", type);
    }
    return false;
}
// Write a cell image/mask/weight
static bool cellWrite(pmCell *cell,     // Cell to write
                      psFits *fits,     // FITS file to which to write
                      psDB *db,         // Database handle for "concepts" update
                      bool blank,       // Write a blank PHU?
                      fpaWriteType type // Type to write
                     )
{
    assert(cell);
    assert(fits);

    psTrace ("pmFPAWrite", 5, "writing to Cell (%d)\n", blank);

    pmHDU *hdu = cell->hdu;             // The HDU
    if (!hdu || !cell->data_exists) {
        return true;                    // We wrote every HDU that exists
    }

    psArray **imageArray = appropriateImageArray(hdu, type); // Array of images in the HDU

    // Generate the HDU if needed --- this is required after a pmFPACopy, or similar, which does not
    // generate the HDU, but only copies the structure.
    if (!blank && !hdu->blankPHU && !*imageArray && (!pmHDUGenerateForCell(cell) || !*imageArray)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to generate HDU for cell --- likely programming error.\n");
        return false;
    }

    // We only write out a blank PHU if it's specifically requested.
    bool writeBlank = blank && hdu->blankPHU && !*imageArray; // Write a blank PHU?
    bool writeImage = !blank && !hdu->blankPHU && *imageArray; // Write an image?

    if (writeBlank || writeImage) {

        pmFPAUpdateNames(cell->parent->parent, cell->parent, cell);
        pmConceptSource source = PM_CONCEPT_SOURCE_HEADER | PM_CONCEPT_SOURCE_CELLS |
                                 PM_CONCEPT_SOURCE_DEFAULTS;
        if (!pmConceptsWriteCell(cell, source, true, NULL)) {
            psError(PS_ERR_IO, false, "Unable to write concepts for cell.\n");
            return false;
        }
        if (!appropriateWriteFunc(hdu, fits, type)) {
            psError(PS_ERR_IO, false, "Unable to write HDU for cell.\n");
            return false;
        }
    }
    // No lower levels to which to recurse

    return true;
}

// Write a chip image/mask/weight
static bool chipWrite(pmChip *chip,     // Chip to write
                      psFits *fits,     // FITS file to which to write
                      psDB *db,         // Database handle for "concepts" update
                      bool blank,       // Write a blank PHU?
                      bool recurse,     // Recurse to lower levels?
                      fpaWriteType type // Type to write
                     )
{
    assert(chip);
    assert(fits);

    pmHDU *hdu = chip->hdu;             // The HDU

    psTrace ("pmFPAWrite", 5, "writing to Chip (%d, %d)\n", blank, recurse);

    // If we have data at this level, try to write it out
    if (hdu && chip->data_exists) {
        psArray **imageArray = appropriateImageArray(hdu, type); // Array of images in HDU

        // Generate the HDU if needed --- this is required after a pmFPACopy, or similar, which does not
        // generate the HDU, but only copies the structure.
        if (!blank && !hdu->blankPHU && !*imageArray && (!pmHDUGenerateForChip(chip) || !*imageArray)) {
            psError(PS_ERR_UNKNOWN, false,
                    "Unable to generate HDU for chip --- likely programming error.\n");
            return false;
        }

        // We only write out a blank PHU if it's specifically requested.
        bool writeBlank = blank && hdu->blankPHU && !*imageArray; // Write a blank HDU?
        bool writeImage = !blank && !hdu->blankPHU && *imageArray; // Write an image?

        if (writeBlank || writeImage) {
            pmFPAUpdateNames(chip->parent, chip, NULL);
            pmConceptSource source = PM_CONCEPT_SOURCE_HEADER | PM_CONCEPT_SOURCE_CELLS |
                                     PM_CONCEPT_SOURCE_DEFAULTS;
            if (!pmConceptsWriteChip(chip, source, true, true, NULL)) {
                psError(PS_ERR_IO, false, "Unable to write concepts for chip.\n");
                return false;
            }
            if (!appropriateWriteFunc(hdu, fits, type)) {
                psError(PS_ERR_IO, false, "Unable to write HDU for chip.\n");
                return false;
            }
        }
    }

    // Recurse to lower level if specifically requested.
    // XXX recursion implies blank == false (must be called on correct level?)
    if (recurse) {
        psArray *cells = chip->cells;       // Array of cells
        for (int i = 0; i < cells->n; i++) {
            pmCell *cell = cells->data[i];  // The cell of interest
            if (!cellWrite(cell, fits, db, false, type)) {
                psError(PS_ERR_IO, false, "Unable to write Chip.\n");
                return false;
            }
        }
    }

    return true;
}


// Write an FPA image/mask/weight
static bool fpaWrite(pmFPA *fpa,        // FPA to write
                     psFits *fits,      // FITS file to which to write
                     psDB *db,          // Database handle for "concepts" update
                     bool blank,        // Write a blank PHU?
                     bool recurse,      // Recurse to lower levels?
                     fpaWriteType type  // Type to write
                    )
{
    assert(fpa);
    assert(fits);

    pmHDU *hdu = fpa->hdu;              // The HDU

    psTrace ("pmFPAWrite", 5, "writing to FPA (%d, %d)\n", blank, recurse);

    // If we have data at this level, try to write it out
    if (hdu) {
        psArray **imageArray = appropriateImageArray(hdu, type); // Array of images in HDU

        // Generate the HDU if needed --- this is required after a pmFPACopy, or similar, which does not
        // generate the HDU, but only copies the structure.
        if (!blank && !hdu->blankPHU && !*imageArray && (!pmHDUGenerateForFPA(fpa) || !*imageArray)) {
            psError(PS_ERR_UNKNOWN, false,
                    "Unable to generate HDU for FPA --- likely programming error.\n");
            return false;
        }

        // We only write out a blank PHU if it's specifically requested.
        bool writeBlank = blank && hdu->blankPHU && !*imageArray; // Write a blank PHU?
        bool writeImage = !blank && !hdu->blankPHU && *imageArray; // Write an image?

        if (writeBlank || writeImage) {
            pmFPAUpdateNames(fpa, NULL, NULL);
            pmConceptSource source = PM_CONCEPT_SOURCE_HEADER | PM_CONCEPT_SOURCE_CELLS |
                                     PM_CONCEPT_SOURCE_DEFAULTS;
            if (!pmConceptsWriteFPA(fpa, source, true, NULL)) {
                psError(PS_ERR_IO, false, "Unable to write concepts for FPA.\n");
                return false;
            }
            if (!appropriateWriteFunc(hdu, fits, type))  {
                psError(PS_ERR_IO, false, "Unable to write HDU for FPA.\n");
                return false;
            }
        }
    }

    // Recurse to lower levels if requested
    if (recurse) {
        psArray *chips = fpa->chips;        // Array of chips
        for (int i = 0; i < chips->n; i++) {
            pmChip *chip = chips->data[i];  // The chip of interest
            if (!chipWrite(chip, fits, db, false, true, type)) {
                psError(PS_ERR_IO, false, "Unable to write FPA.\n");
                return false;
            }
        }
    }

    return true;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Public functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

bool pmReadoutWriteNext(pmReadout *readout, psFits *fits, int z)
{
    PS_ASSERT_PTR_NON_NULL(readout, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    pmHDU *hdu = pmHDUFromReadout(readout); // The HDU to which to write
    if (!hdu) {
        psError(PS_ERR_IO, false, "Unable to find HDU for readout.\n");
        return false;
    }
    psMetadata *header = hdu->header;   // The FITS header
    if (!header) {
        psError(PS_ERR_IO, true, "No FITS header available in the HDU.\n");
        return false;
    }

    // We have to rely to a great extent on the FITS header, because otherwise we simply don't know how many
    // image planes there are (NAXIS3) or the size of the original image (NAXIS1, NAXIS2).
    bool mdok = true;                   // Status of MD lookup
    int naxis1 = psMetadataLookupS32(&mdok, header, "NAXIS1"); // Number of columns
    if (!mdok || naxis1 <= 0) {
        psError(PS_ERR_IO, true, "Can't find NAXIS1 in header.\n");
        return false;
    }
    int naxis2 = psMetadataLookupS32(&mdok, header, "NAXIS2"); // Number of rows
    if (!mdok || naxis2 <= 0) {
        psError(PS_ERR_IO, true, "Can't find NAXIS2 in header.\n");
        return false;
    }
    int naxis3 = psMetadataLookupS32(&mdok, header, "NAXIS3"); // Number of image planes
    if (!mdok || naxis3 <= 0) {
        naxis3 = 1;
    }
    if (z >= naxis3) {
        psError(PS_ERR_IO, true, "Specified a plane number (%d) greater than NAXIS3 allows.\n", z);
        return false;
    }

    if (!hdu->images) {
        psError(PS_ERR_IO, true, "No images allocated in HDU.\n");
        return false;
    }
    psImage *image = readout->image; // The image from the HDU to write
    if (readout->row0 == 0 && readout->col0 == 0 && z == 0) {
        // Then we can assume that nothing has been written to the FITS file for now
        if (naxis1 == image->numCols && naxis2 == image->numRows) {
            // We can write the whole lot at once
            return psFitsWriteImage(fits, header, image, z, hdu->extname);
        }
        // Create a dummy image so we can write something larger than we actually have
        psImage *dummy = psImageAlloc(naxis1, naxis2, image->type.type); // Dummy image
        psImageInit(dummy, 0);
        psImageOverlaySection(dummy, image, 0, 0, "=");
        bool result = psFitsWriteImage(fits, header, dummy, z, hdu->extname);
        psFree(dummy);
        return result;
    }

    // We can simply update an existing HDU
    if (hdu->blankPHU && !psFitsMoveExtNum(fits, 0, false)) {
        psError(PS_ERR_IO, false, "Unable to move to PHU\n");
        return false;
    }
    return psFitsUpdateImage(fits, image, readout->col0, readout->row0, z);
}


bool pmCellWrite(pmCell *cell, psFits *fits, psDB *db, bool blank)
{
    PS_ASSERT_PTR_NON_NULL(cell, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    return cellWrite(cell, fits, db, blank, FPA_WRITE_TYPE_IMAGE);
}

bool pmChipWrite(pmChip *chip, psFits *fits, psDB *db, bool blank, bool recurse)
{
    PS_ASSERT_PTR_NON_NULL(chip, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    return chipWrite(chip, fits, db, blank, recurse, FPA_WRITE_TYPE_IMAGE);
}

bool pmFPAWrite(pmFPA *fpa, psFits *fits, psDB *db, bool blank, bool recurse)
{
    PS_ASSERT_PTR_NON_NULL(fpa, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    return fpaWrite(fpa, fits, db, blank, recurse, FPA_WRITE_TYPE_IMAGE);
}


bool pmCellWriteMask(pmCell *cell, psFits *fits, psDB *db, bool blank)
{
    PS_ASSERT_PTR_NON_NULL(cell, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    return cellWrite(cell, fits, db, blank, FPA_WRITE_TYPE_IMAGE);
}

bool pmChipWriteMask(pmChip *chip, psFits *fits, psDB *db, bool blank, bool recurse)
{
    PS_ASSERT_PTR_NON_NULL(chip, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    return chipWrite(chip, fits, db, blank, recurse, FPA_WRITE_TYPE_MASK);
}

bool pmFPAWriteMask(pmFPA *fpa, psFits *fits, psDB *db, bool blank, bool recurse)
{
    PS_ASSERT_PTR_NON_NULL(fpa, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    return fpaWrite(fpa, fits, db, blank, recurse, FPA_WRITE_TYPE_MASK);
}


bool pmCellWriteWeight(pmCell *cell, psFits *fits, psDB *db, bool blank)
{
    PS_ASSERT_PTR_NON_NULL(cell, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    return cellWrite(cell, fits, db, blank, FPA_WRITE_TYPE_WEIGHT);
}

bool pmChipWriteWeight(pmChip *chip, psFits *fits, psDB *db, bool blank, bool recurse)
{
    PS_ASSERT_PTR_NON_NULL(chip, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    return chipWrite(chip, fits, db, blank, recurse, FPA_WRITE_TYPE_WEIGHT);
}

bool pmFPAWriteWeight(pmFPA *fpa, psFits *fits, psDB *db, bool blank, bool recurse)
{
    PS_ASSERT_PTR_NON_NULL(fpa, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    return fpaWrite(fpa, fits, db, blank, recurse, FPA_WRITE_TYPE_WEIGHT);
}


int pmCellWriteTable(psFits *fits, const pmCell *cell, const char *name)
{
    PS_ASSERT_PTR_NON_NULL(cell, 0);
    PS_ASSERT_PTR_NON_NULL(fits, 0);
    PS_ASSERT_STRING_NON_EMPTY(name, 0);

    const char *chipName = psMetadataLookupStr(NULL, cell->parent->concepts, "CHIP.NAME"); // Name of chip
    const char *cellName = psMetadataLookupStr(NULL, cell->concepts, "CELL.NAME"); // Name of cell

    psArray *table = psMetadataLookupPtr(NULL, cell->analysis, name); // The FITS table
    if (!table) {
        // We wrote everything we could find
        return 0;
    }

    psString headerName = NULL;         // Name for header in analysis metadata
    psStringAppend(&headerName, "%s.HEADER", name);
    psMetadata *header = psMetadataLookupMetadata(NULL, cell->analysis, headerName); // The FITS header
    psFree(headerName);

    psString extname = NULL;            // Extension name
    psStringAppend(&extname, "%s_%s_%s", name, chipName, cellName);

    // XXX Could do a table lookup from the camera format, in case the input file isn't laid out with
    // NAME_CHIP_CELL extension names --- use these as keys, and the value as the proper extension name.
    // Allow interpolation of concepts, e.g., "{CHIP.NAME}" --> "ccd13".

    if (!psFitsWriteTable(fits, header, table, extname)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to write table from chip %s, cell %s to extension %s\n",
                chipName, cellName, extname);
        psFree(extname);
        return 0;
    }

    psFree(extname);
    return 1;
}


int pmChipWriteTable(psFits *fits, const pmChip *chip, const char *name)
{
    PS_ASSERT_PTR_NON_NULL(chip, 0);
    PS_ASSERT_PTR_NON_NULL(fits, 0);
    PS_ASSERT_STRING_NON_EMPTY(name, 0);

    int numWrite = 0;                    // Number of reads
    psArray *cells = chip->cells;       // Array of cells
    for (int i = 0; i < cells->n; i++) {
        pmCell *cell = cells->data[i];  // Cell of interest
        numWrite += pmCellWriteTable(fits, cell, name);
    }

    return numWrite;
}


int pmFPAWriteTable(psFits *fits, const pmFPA *fpa, const char *name)
{
    PS_ASSERT_PTR_NON_NULL(fpa, 0);
    PS_ASSERT_PTR_NON_NULL(fits, 0);
    PS_ASSERT_STRING_NON_EMPTY(name, 0);

    int numWrite = 0;                    // Number of reads
    psArray *chips = fpa->chips;        // Array of chips
    for (int i = 0; i < chips->n; i++) {
        pmChip *chip = chips->data[i];  // Chip of interest
        numWrite += pmChipWriteTable(fits, chip, name);
    }

    return numWrite;
}
