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

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

#include "pmFPALevel.h"
#include "pmFPAUtils.h"
#include "pmShifts.h"

#define HASH_SIZE 100                   // Size of hash with cells; should be > numCells for efficiency
#define SHIFTS_BUFFER 100               // Buffer size for shifts
#define ANALYSIS_NAME "OT.KERNEL"       // Name for kernel on the analysis metadata
#define TRACE "psModules.detrend"       // Trace facility

// XXX To do:
// * Make the table column names configurable, by having a SHIFTS metadata in the camera format, with entries
//   specifying the column names.


// Shifts due to orthogonal transfer
typedef struct {
    psVector *x;                        // Shifts in x
    psVector *y;                        // Shifts in y
    psVector *t;                        // Times of shifts
    long num;                           // Number of values
} otShifts;

// Allocator for otShifts
static otShifts *otShiftsAlloc(void)
{
    otShifts *shifts = psAlloc(sizeof(otShifts));
    shifts->x = psVectorAllocEmpty(SHIFTS_BUFFER, PS_TYPE_S32);
    shifts->y = psVectorAllocEmpty(SHIFTS_BUFFER, PS_TYPE_S32);
    shifts->t = psVectorAllocEmpty(SHIFTS_BUFFER, PS_TYPE_F32);
    shifts->num = 0;
    return shifts;
}

// Look up the cell within a hash; supplement the hash with a new value if it doesn't exist
static otShifts *cellVectors(psHash *shifts, // Hash of shifts
                             const char *cellName // Key for hash
    )
{
    assert(shifts);
    assert(cellName);

    // Find the appropriate cell
    otShifts *vectors = psHashLookup(shifts, cellName);
    if (!vectors) {
        vectors = otShiftsAlloc();
        psHashAdd(shifts, cellName, vectors);
        psFree(vectors);            // Drop reference
    }
    return vectors;
}

// Generate a kernel and stuff it on the cell metadata
static bool stuffKernel(pmCell *cell,     // Cell to which the shifts belong
                        const otShifts *shifts // Shifts to apply
    )
{
    assert(cell);
    assert(shifts);

    psKernel *kernel = psKernelGenerate(shifts->t, shifts->x, shifts->y, false, true); // Shift kernel
    if (!kernel) {
        psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to generate kernel from OT shifts");
        return false;
    }
    psMetadataAddPtr(cell->analysis, PS_LIST_TAIL, ANALYSIS_NAME, PS_DATA_KERNEL,
                     "Orthogonal transfer kernel, calculated from shifts", kernel);
    psFree(kernel);

    return true;
}


bool pmShiftsRead(pmCell *cell, psFits *fits)
{
    PS_ASSERT_PTR_NON_NULL(fits, false);
    PS_ASSERT_PTR_NON_NULL(cell, false);

    bool mdok;                          // Status of MD lookup
    psKernel *check = psMetadataLookupPtr(&mdok, cell->analysis, ANALYSIS_NAME); // Kernel, or NULL
    if (check) {
        psTrace(TRACE, 2, "Cell already has OT kernel present.\n");
        return true;                    // No error
    }
    psTrace(TRACE, 2, "Reading FITS file for OT kernels.\n");

    pmChip *chip = cell->parent;        // The parent chip
    pmFPA *fpa = chip->parent;          // The parent FPA
    pmHDU *phu = NULL;                  // The primary header
    pmFPALevel phuLevel = PM_FPA_LEVEL_NONE; // Level of the PHU
    long numChips = 0;                  // Number of chips below the PHU; for setting efficient hash size
    long numCells = 0;                  // Number of cells below the PHU; for setting efficient hash size
    if (fpa->hdu) {
        phu = fpa->hdu;
        // Count the cells
        psArray *chips = fpa->chips;    // Array of chips
        numChips = chips->n;
        for (int i = 0; i < chips->n; i++) {
            pmChip *chip = chips->data[i]; // Chip of interest
            numCells += chip->cells->n;
        }
        numCells = (float)numCells / (float)numChips + 0.5; // Average number of cells per chip
        phuLevel = PM_FPA_LEVEL_FPA;
    }
    if (!phu && chip->hdu) {
        phu = chip->hdu;
        numChips = 1;
        numCells = chip->cells->n;
        phuLevel = PM_FPA_LEVEL_CHIP;
    }
    if (!phu && cell->hdu) {
        phu = cell->hdu;
        numChips = 0;
        numCells = 1;
        phuLevel = PM_FPA_LEVEL_CELL;
    }

    if (!phu || phuLevel == PM_FPA_LEVEL_NONE) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "Can't find the PHU.\n");
        return false;
    }

    // Find out what to read
    psMetadata *fileInfo = psMetadataLookupMetadata(&mdok, phu->format, "FILE"); // FILE in the format
    if (!mdok || !fileInfo) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to find FILE information in camera format.\n");
        return false;
    }
    const char *shiftsExt = psMetadataLookupStr(&mdok, fileInfo, "SHIFTS"); // Name of shifts extension
    if (!mdok || !shiftsExt) {
        // We read all the shifts that are present --- which is none.
        return true;
    }

    // Read the FITS file
    int origExt = psFitsGetExtNum(fits); // Original extension number; to preserve position
    if (!psFitsMoveExtName(fits, shiftsExt)) {
        psError(PS_ERR_IO, false, "Unable to move to shifts extension %s", shiftsExt);
        return false;
    }
    psArray *table = psFitsReadTable(fits); // The table of shifts
    if (!table) {
        psError(PS_ERR_IO, false, "Unable to read shifts table.\n");
        return false;
    }

    // More sensible storage
    otShifts *singleShifts = NULL;      // Shifts for a single cell
    psHash *multipleShifts = NULL;      // Shifts for multiple cells, stored by cell name
    switch (phuLevel) {
      case PM_FPA_LEVEL_FPA:
        multipleShifts = psHashAlloc(2 * numChips);
        break;
      case PM_FPA_LEVEL_CHIP:
        multipleShifts = psHashAlloc(2 * numCells);
        break;
      case PM_FPA_LEVEL_CELL:
        singleShifts = otShiftsAlloc();
        break;
      default:
        psAbort("Should never get here.\n");
    }

    // Pull values out of the table into something a bit more sensible
    for (int i = 0; i < table->n; i++) {
        psMetadata *row = table->data[i]; // The row of interest

        const char *chipName;           // Name of chip
        if (phuLevel == PM_FPA_LEVEL_FPA) {
            // Only care about the chip name if there's more than one chip
            chipName = psMetadataLookupStr(&mdok, row, "Chip");
            if (!mdok || !chipName) {
                psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to find Chip in row %d of shifts table\n", i);
                psFree(multipleShifts);
                psFree(table);
                return false;
            }
        }
        const char *cellName;           // Name of cell
        if (phuLevel <= PM_FPA_LEVEL_CHIP) {
            // Only care about the cell name if there's a chip
            cellName = psMetadataLookupStr(&mdok, row, "Cell");
            if (!mdok || !cellName) {
                psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to find Cell in row %d of shifts table\n", i);
                psFree(multipleShifts);
                psFree(table);
                return false;
            }
        }
        float x = psMetadataLookupS32(&mdok, row, "X"); // Shift in x
        if (!mdok) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to find X in row %d of shifts table\n", i);
            psFree(multipleShifts);
            psFree(table);
            return false;
        }
        float y = psMetadataLookupF32(&mdok, row, "Y"); // Shift in y
        if (!mdok) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to find Y in row %d of shifts table\n", i);
            psFree(multipleShifts);
            psFree(table);
            return false;
        }
        float t = psMetadataLookupF32(&mdok, row, "Time"); // Time of shift
        if (!mdok) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to find Time in row %d of shifts table\n", i);
            psFree(multipleShifts);
            psFree(table);
            return false;
        }

        otShifts *shifts = NULL;        // Shifts for the cell of interest
        switch (phuLevel) {
          case PM_FPA_LEVEL_FPA: {
              psHash *cells = psHashLookup(multipleShifts, chipName); // Hash of cells
              if (!cells) {
                  cells = psHashAlloc(numCells);
                  psHashAdd(multipleShifts, chipName, cells);
                  psFree(cells);          // Drop reference
              }
              shifts = cellVectors(cells, cellName);
              break;
          }
          case PM_FPA_LEVEL_CHIP:
            shifts = cellVectors(multipleShifts, cellName);
            break;
          case PM_FPA_LEVEL_CELL:
            shifts = singleShifts;
            break;
          default:
            psAbort("Should never get here.\n");
        }

        // Add the shift
        psVectorExtend(shifts->x, 1, SHIFTS_BUFFER);
        psVectorExtend(shifts->y, 1, SHIFTS_BUFFER);
        psVectorExtend(shifts->t, 1, SHIFTS_BUFFER);
        shifts->x->data.S32[shifts->num] = (int)x;
        shifts->y->data.S32[shifts->num] = (int)y;
        shifts->t->data.F32[shifts->num] = t;
        shifts->num++;
    }
    psFree(table);

    // Put the kernels into their own cells
    if (phuLevel == PM_FPA_LEVEL_CELL) {
        // Only a single cell
        if (!stuffKernel(cell, singleShifts)) {
            psFree(singleShifts);
            return false;
        }
        psFree(singleShifts);
    } else {
        psList *names = psHashKeyList(multipleShifts); // List of hash keys (chip/cell names)
        psListIterator *namesIter = psListIteratorAlloc(names, PS_LIST_HEAD, false); // Iterator for names
        const char *name;               // Name, from iteration
        while ((name = psListGetAndIncrement(namesIter))) {
            switch (phuLevel) {
              case PM_FPA_LEVEL_FPA: {
                  int chipNum = pmFPAFindChip(fpa, name); // Number of chip of interest
                  if (chipNum < 0) {
                      psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find chip %s", name);
                      psFree(namesIter);
                      psFree(names);
                      psFree(multipleShifts);
                      return false;
                  }
                  pmChip *chip = fpa->chips->data[chipNum]; // Chip of interest

                  // Loop over component cells
                  psHash *cells = psHashLookup(multipleShifts, name); // Hash of cells
                  psList *cellNames = psHashKeyList(multipleShifts); // List of hash keys (cell names)
                  psListIterator *cellNamesIter = psListIteratorAlloc(cellNames, PS_LIST_HEAD, false);
                  const char *cellName;       // Cell name, from iteration
                  while ((cellName = psListGetAndIncrement(cellNamesIter))) {
                      int cellNum = pmChipFindCell(chip, cellName); // Number of cell of interest
                      if (!cell) {
                          psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find cell %s", cellName);
                          psFree(cellNamesIter);
                          psFree(cellNames);
                          psFree(namesIter);
                          psFree(names);
                          psFree(multipleShifts);
                          return false;
                      }
                      pmCell *cell = chip->cells->data[cellNum]; // Cell of interest

                      otShifts *vectors = psHashLookup(cells, cellName); // Shifts for the cell of interest
                      if (!stuffKernel(cell, vectors)) {
                          psFree(cellNamesIter);
                          psFree(cellNames);
                          psFree(namesIter);
                          psFree(names);
                          psFree(multipleShifts);
                          return false;
                      }
                  }
                  psFree(cellNamesIter);
                  psFree(cellNames);
                  break;
              }
              case PM_FPA_LEVEL_CHIP: {
                  int cellNum = pmChipFindCell(chip, name); // Number of cell of interest
                  if (!cell) {
                      psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find cell %s", name);
                      psFree(namesIter);
                      psFree(names);
                      psFree(multipleShifts);
                      return false;
                  }
                  pmCell *cell = chip->cells->data[cellNum]; // Cell of interest

                  otShifts *vectors = psHashLookup(multipleShifts, name); // Shifts for this cell
                  if (!stuffKernel(cell, vectors)) {
                      psFree(namesIter);
                      psFree(names);
                      psFree(multipleShifts);
                      return false;
                  }
                  break;
              }
              default:
                psAbort("Should never get here.\n");
            }
        }
        psFree(namesIter);
        psFree(names);
        psFree(multipleShifts);
    }

    // Go back to the original position in the FITS file
    return psFitsMoveExtNum(fits, origExt, false);
}
