/*******************************************************************************
 * JMMC project ( http://www.jmmc.fr ) - Copyright (C) CNRS.
 ******************************************************************************/
/**
 * @file
 * Functions to handle frame selection structure/file. 
 */

/*
 * System Headers 
 */
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <sys/stat.h>
#include <time.h>

/*
 * Local Headers 
 */
#include "amdlib.h"
#include "amdlibProtected.h"

/**
 * Allocate memory for structure containing necessary information for frame
 * selection.
 *
 * @param selection structure to be allocated.
 * @param nbFrames total number of frames in current observation.
 * @param nbBases number of baselines in current observation.
 * @param errMsg error description message returned if function fails.
 *
 * @return 
 * amdlibSUCCESS on successful completion. Otherwise amdlibFAILURE is returned.
 */
amdlibCOMPL_STAT amdlibAllocateSelection(amdlibSELECTION *selection,
                                         int nbFrames, 
                                         int nbBases,
                                         amdlibERROR_MSG errMsg)
{
    int i;
    int band, base;

    amdlibLogTrace("amdlibAllocateSelection()");

    selection->nbFrames = nbFrames;
    selection->nbBases  = nbBases;

    /* For each band */
    for (band = 0; band < amdlibNB_BANDS; band++)
    {
        /* Flag related to selected frames */
        selection->band[band].isSelected = calloc(nbBases*nbFrames, sizeof(char));
        if (selection->band[band].isSelected == NULL)
        {
            amdlibSetErrMsg("%s selected frame flag array", 
                            amdlibERR_ALLOC_MEM);
            return amdlibFAILURE;
        }
        
        /* Wrap this array */
        selection->band[band].isSelectedPt = 
            amdlibWrap2DArrayUnsignedChar(selection->band[band].isSelected,
                                          nbFrames, nbBases, errMsg);
        if (selection->band[band].isSelectedPt == NULL)
        {
            return amdlibFAILURE;
        }
        
        /* Tag all frames as selected */
        for (base=0; base < selection->nbBases; base++)
        {
            selection->band[band].nbSelectedFrames[base] = nbFrames;
        }
        memset(selection->band[band].isSelected, 1, 
               nbBases * nbFrames * sizeof(char));

        /* Flag related to phase closures */
        selection->band[band].frameOkForClosure = calloc(nbFrames, sizeof(int));
        if (selection->band[band].frameOkForClosure == NULL)
        {
            amdlibSetErrMsg("%s phase closure flag array", amdlibERR_ALLOC_MEM);
            return amdlibFAILURE;
        }
        for (i = 0; i < nbFrames; i++)
        {
            selection->band[band].frameOkForClosure[i] = i;
        }
        selection->band[band].nbFramesOkForClosure = nbFrames;
    }

    return amdlibSUCCESS;
}

/**
 * Select/Unselect all frames for the given band.
 *
 * @param selection structure to be allocated.
 * @param band spectral band
 * @param nbBases number of baselines in current observation.
 *
 * @return 
 * amdlibSUCCESS on successful completion. Otherwise amdlibFAILURE is returned.
 */
amdlibCOMPL_STAT amdlibSetSelection(amdlibSELECTION *selection,
                                    amdlibBAND    band,
                                    amdlibBOOLEAN isSelected)
{
    int base, frame;
    
    amdlibLogTrace("amdlibSetSelection()");
    
    /* Flag related to frame selection */
    for (base = 0; base < selection->nbBases; base++)
    {
        for (frame = 0; frame < selection->nbFrames; frame++)
        {
            selection->band[band].isSelectedPt[base][frame] = isSelected;
        }
    }

    /* Update flags related to phase closures */
    amdlibUpdateSelection(selection);

    return amdlibSUCCESS;
}

/**
 * Update informations related to phase closures
 *
 * @param selection structure containing information on selection.
 *
 * @return
 * amdlibSUCCESS
 */
amdlibCOMPL_STAT amdlibUpdateSelection(amdlibSELECTION *selection)
{
    int band, frame, base;

    amdlibLogTrace("amdlibUpdateSelection()");
    
    /* For each band */
    for (band = 0; band < amdlibNB_BANDS; band++)
    {
        int nbFramesOkForClosure = 0;
        
        for (base = 0; base < selection->nbBases; base++)
        {    
            selection->band[band].nbSelectedFrames[base]   = 0;
            selection->band[band].firstSelectedFrame[base] = -1;
            for (frame = 0; frame < selection->nbFrames; frame++)
            {    
                if (selection->band[band].isSelectedPt[base][frame] == 
                                                                    amdlibTRUE)
                {
                    selection->band[band].nbSelectedFrames[base] += 1;
                    if (selection->band[band].firstSelectedFrame[base] == -1)
                    {
                        selection->band[band].firstSelectedFrame[base] = frame;
                    }
                }
            }
        }
        /* 3T case. We deem OK for closure if 1 frame has passed the selection test */
        if (selection->nbBases == 3)
        {
            int frame;
            for (frame = 0; frame < selection->nbFrames; frame++)
            {
                if ((selection->band[band].isSelectedPt[0][frame] == amdlibTRUE) ||
                    (selection->band[band].isSelectedPt[1][frame] == amdlibTRUE) ||
                    (selection->band[band].isSelectedPt[2][frame] == amdlibTRUE)   )
                {
                    selection->band[band].frameOkForClosure[nbFramesOkForClosure] = frame;
                    nbFramesOkForClosure++;
                }
            }
        }
        selection->band[band].nbFramesOkForClosure = nbFramesOkForClosure;
    }
    return amdlibSUCCESS;
}

/**
 * Releases memory allocated to store information relative to frame selection.
 *
 * @param selection structure where selection information is stored.
 */
void amdlibReleaseSelection(amdlibSELECTION *selection)
{
    int band;

    amdlibLogTrace("amdlibReleaseSelection()");

    /* For each band */
    for (band = 0; band < amdlibNB_BANDS; band++)
    {
        if (selection->band[band].isSelected != NULL)
        {
            free(selection->band[band].isSelected);
        }

        if (selection->band[band].frameOkForClosure != NULL)
        {
            free(selection->band[band].frameOkForClosure);
        }
        amdlibFree2DArrayUnsignedCharWrapping(selection->band[band].isSelectedPt);
    }
    memset(selection, '\0', sizeof(amdlibSELECTION));    
}

/**
 * This function reads given fits file and extracts selection information.
 * 
 * @param fileName name of the file to read.
 * @param selKeywList list of keywords related to frame selection
 * @param selectionBandJ selection structure for spectral band 'J'.
 * @param selectionBandH selection structure for spectral band 'H'.
 * @param selectionBandK selection structure for spectral band 'K'.
 * @param errMsg error description message returned if function fails.
 *
 * @return 
 * amdlibSUCCESS on successful completion. Otherwise amdlibFAILURE is returned.
 */
amdlibCOMPL_STAT amdlibReadSelectionFile(const char      *fileName,
                                         amdlibINS_CFG   *selKeywList,
                                         amdlibSELECTION *selection,
                                         amdlibERROR_MSG errMsg)
{
    fitsfile *filePtr;
    int status = 0;
    char fitsioMsg[256];
    int band, i;
    int nbFrames;
    long nbBases;
    char keyName[36];
    char keyVal[36];
    char comment[amdlibKEYW_CMT_LEN+1];
    char nullLogical = amdlibFALSE;
    int anynull;
    int keysExist = 0;
    int moreKeys = 0;
    amdlibKEYW_LINE record;
    
    amdlibLogTrace("amdlibReadSelectionFile ()");
   
    /* Check a file name has been specified */
    if (strlen(fileName) == 0)
    {
        amdlibSetErrMsg("No input file name specified");
        return amdlibFAILURE;
    }
    
    /* Open FITS file */
    if (fits_open_file(&filePtr, fileName, READONLY, &status))
    {
        amdlibReturnFitsError(fileName);
    }    

    /* Get primary header's keywords */
    keysExist = 0;
    moreKeys = 0;
    if (fits_get_hdrspace(filePtr, &keysExist, &moreKeys, &status) != 0)
    {
        status = 0;
    }
    for (i = 1; i <= keysExist; i++)
    {
        /* Read keyword line #i */
        if (fits_read_record(filePtr, i, record, &status) != 0)
        {
            status = 0;
        }
        else
        {
            /* Only retrieve HIERARCH ESO keywords, excepted PRO CATG
             * keyword */
            if ((strstr(record, "HIERARCH ESO ") != NULL) &&
                (strstr(record, "HIERARCH ESO PRO CATG") == NULL))
            {
                /* Store keyword */
                if (amdlibAddInsCfgKeyword(selKeywList, 
                                           record, errMsg) == amdlibFAILURE)
                {
                    return amdlibFAILURE;
                }
            }
        }
    }
    
    /* Get all rows data */
    if (fits_movnam_hdu(filePtr, BINARY_TBL, "FRAME_SELECTION", 0, &status))
    {
        amdlibGetFitsError("FRAME_SELECTION table");
        fits_close_file(filePtr, &status);
        return amdlibFAILURE;
    }  
    
    /* Get number of bases */
    if (fits_get_num_rows(filePtr, &nbBases, &status))
    {
        amdlibGetFitsError("FRAME_SELECTION number of rows");
        fits_close_file(filePtr, &status);
        return amdlibFAILURE;
    }

    /* Get number of frames */
    int naxis;
    long naxes[2];
    if (fits_read_tdim(filePtr, 1, 2, &naxis, naxes, &status))
    {
        amdlibGetFitsError("FRAME_SELECTION dimensions");
        fits_close_file(filePtr, &status);
        return amdlibFAILURE;
    }
    nbFrames = naxes[0];
    
    /* Allocate frame selection structure */
    if (amdlibAllocateSelection(selection, nbFrames,
                                nbBases, errMsg) == amdlibFAILURE)
    {
        return amdlibFAILURE;
    }

    for (band=amdlibJ_BAND; band <= amdlibK_BAND; band++)
    {
        sprintf(keyName, "TFORM%d", band+1);
        if (fits_read_key(filePtr, TSTRING, keyName, keyVal, 
                          comment, &status))
        {
            amdlibGetFitsError(keyName);
            fits_close_file(filePtr, &status);
            return amdlibFAILURE;
        }
        sscanf(keyVal, "%dL", &nbFrames);

        if (fits_read_col(filePtr, TLOGICAL, band+1, 1, 1, 
                          nbFrames * nbBases,
                          &nullLogical, selection->band[band].isSelected,
                          &anynull, &status))
        {
            amdlibGetFitsError("FRAME_SELECTION read column");
            fits_close_file(filePtr, &status);
            return amdlibFAILURE;
        }
    }

    /* Close file */
    if (fits_close_file(filePtr, &status))
    {
        amdlibReturnFitsError(fileName);
    }
    
    /* Update selection; i.e. compute number of selected frames, find the first
     * ones, ...*/
    amdlibUpdateSelection(selection);

    return amdlibSUCCESS;
}

/**
 * This function stores given selection structures into a binary table.
 * 
 * @param fileName name of the file to be written.
 * @param selectionBandJ selection structure for spectral band 'J'.
 * @param selectionBandH selection structure for spectral band 'H'.
 * @param selectionBandK selection structure for spectral band 'K'.
 * @param errMsg error description message returned if function fails.
 *
 * @return 
 * amdlibSUCCESS on successful completion. Otherwise amdlibFAILURE is returned.
 */
amdlibCOMPL_STAT amdlibWriteSelectionFile(const char      *fileName,
                                          amdlibSELECTION *selection,
                                          amdlibERROR_MSG errMsg)
{
    struct stat statBuf;
    fitsfile *filePtr;
    int status = 0;
    char fitsioMsg[256];
    int  tfields = 3; 
    char *ttype[] = {"J", "H", "K"};
    char *tform[3];    
    char *tunit[] = {"", "", ""};
    int band;
    time_t      timeSecs;
    struct tm   *timeNow;
    char        strTime[amdlibKEYW_VAL_LEN+1];
   
    /* Check a file name has been specified */
    if (strlen(fileName) == 0)
    {
        amdlibSetErrMsg("No input file name specified");
        return amdlibFAILURE;
    }
    
    /* First remove previous IO file (if any) */
    if (stat(fileName, &statBuf) == 0)
    {
        if (remove(fileName) != 0)
        {
            amdlibSetErrMsg("Could not overwrite file %s", fileName); 
            return amdlibFAILURE;
        }
    }

    /* Create file */
    if (fits_create_file(&filePtr, fileName, &status) != 0)
    {
        amdlibReturnFitsError(fileName);
    }

    for (band=amdlibJ_BAND; band <= amdlibK_BAND; band++)
    {
        tform[band] = calloc(36, sizeof(char));
        sprintf(tform[band], "%dL", selection->nbFrames);
    }
    
    /* Create binary table */
    if (fits_create_tbl(filePtr, BINARY_TBL, 1, tfields, ttype, tform, tunit, 
                        "FRAME_SELECTION", &status) != 0)
    {
        amdlibGetFitsError("FRAME_SELECTION");
        fits_close_file(filePtr, &status);
        return amdlibFAILURE;        
    }

    for (band=amdlibJ_BAND; band <= amdlibK_BAND; band++)
    {
        free(tform[band]);
    }

    /* Write columns */
    for (band=amdlibJ_BAND; band <= amdlibK_BAND; band++)
    {
        if (fits_write_col(filePtr, TLOGICAL, band+1, 1, 1, 
                           selection->nbBases * selection->nbFrames, 
                           selection->band[band].isSelected, 
                           &status) != 0)
        {
            amdlibGetFitsError("FRAME_SELECTION - write column");
            fits_close_file(filePtr, &status);
            return amdlibFAILURE;
        }
    }

    /* Write keywords in primary header */
    if (fits_movabs_hdu(filePtr, 1,0, &status) != 0)
    {
        amdlibGetFitsError("Main header");
        fits_close_file(filePtr, &status);
        return amdlibFAILURE;        
    }
    timeSecs = time(NULL);
    timeNow = gmtime(&timeSecs);
    strftime(strTime, sizeof(strTime), "%Y-%m-%dT%H:%M:%S", timeNow);
    if (fits_write_key(filePtr, TSTRING, "DATE", strTime, 
                       "Date this file was written", &status) != 0)
    {
        amdlibGetFitsError("DATE");
        fits_close_file(filePtr, &status);
        return amdlibFAILURE;
    }
    if (fits_write_key(filePtr, TSTRING, "HIERARCH ESO PRO CATG", 
                       "FRAME_SELECTION",
                       "Category of product frames", &status) != 0)
    {
        amdlibGetFitsError("HIERARCH ESO PRO CATG");
        fits_close_file(filePtr, &status);
        return amdlibFAILURE;
    }

    /* Close file */
    if (fits_close_file(filePtr, &status))
    {
        amdlibReturnFitsError(fileName);
    }
    
    return amdlibSUCCESS;
}
