/* $Id: hawki_bkg.cc,v 1.7 2013-09-02 14:36:56 cgarcia Exp $
 *
 * This file is part of the HAWKI Pipeline
 * Copyright (C) 2002,2003 European Southern Observatory
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

/*
 * $Author: cgarcia $
 * $Date: 2013-09-02 14:36:56 $
 * $Revision: 1.7 $
 * $Name: not supported by cvs2svn $
 */

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

/*-----------------------------------------------------------------------------
                                   Includes
 -----------------------------------------------------------------------------*/

#include <string.h>
#include <cpl.h>
#include <vector>
#include <algorithm>

#include "hawki_distortion.h"
#include "hawki_bkg.h"
#include "hawki_pfits.h"
#include "hawki_load.h"
#include "hawki_utils.h"
#include "hawki_image.h"

/*----------------------------------------------------------------------------*/
/**
 * @defgroup hawki_bkg   Background computation algorithms
 */
/*----------------------------------------------------------------------------*/

/**@{*/

hawki_bkg_frames_buffer * hawki_bkg_frames_buffer_init
(const cpl_frameset * fset)
{
    int i_ima;
    hawki_bkg_frames_buffer * frames_buffer;

    frames_buffer = 
          (hawki_bkg_frames_buffer *)cpl_malloc(sizeof(hawki_bkg_frames_buffer));

    frames_buffer->nframes = cpl_frameset_get_size(fset);
    frames_buffer->images  = 
           (cpl_image **)cpl_malloc(frames_buffer->nframes * 
                                        sizeof(cpl_image *));
    frames_buffer->medians = 
            (double *)cpl_malloc(frames_buffer->nframes * 
                                        sizeof(double));
    frames_buffer->frames  = cpl_frameset_duplicate(fset);

    for(i_ima=0; i_ima< frames_buffer->nframes; i_ima++)
        frames_buffer->images[i_ima] = NULL;

    return frames_buffer;
}

void hawki_bkg_frames_buffer_delete(hawki_bkg_frames_buffer * frames_buffer)
{
    int i_ima;
    for(i_ima=0; i_ima< frames_buffer->nframes; i_ima++)
    {
        if(frames_buffer->images[i_ima] != NULL)
            cpl_image_delete(frames_buffer->images[i_ima]);
    }
    cpl_free(frames_buffer->images);
    cpl_free(frames_buffer->medians);
    cpl_frameset_delete(frames_buffer->frames);
    cpl_free(frames_buffer);
}


/*----------------------------------------------------------------------------*/
/**
  @brief    Fill a propertylist with the association between objects and bkg 
  @param    objframes   The objframes associated to the background
  @param    proplist    The list of keywords to store the association
                        information to  
  @return   0 if everything is ok, -1 otherwise

  This function takes as an input a set of object frames. 
  For EACH input object frames it creates in the given propertylist a
  keyword "pointing" to the ARCFILE of that frame. 
  This propertylist can later be stored in the background image to establish 
  the association between this background and the object frames. 
 */
/*----------------------------------------------------------------------------*/
int hawki_bkg_fill_assoc(cpl_frameset * objframes, cpl_propertylist * proplist)
{
    int   iframe;
    int   nframes;
    
    /* Loop on the frames */
    nframes = cpl_frameset_get_size(objframes);
    for(iframe = 0 ; iframe < nframes ; ++iframe)
    {
        cpl_frame         * this_frame;
        cpl_propertylist  * this_proplist;
        const char        * arcfile;
        char                keystr[256];

        /* Getting the proplist for this frame */
        this_frame     = cpl_frameset_get_position(objframes, iframe);
        this_proplist  = cpl_propertylist_load
            (cpl_frame_get_filename(this_frame),0);
        arcfile        = hawki_pfits_get_arcfile(proplist);
        snprintf(keystr, 256, "ESO QC BKG ASSOC RAW%d",iframe+1);
        cpl_propertylist_append_string(proplist, keystr, arcfile);
        cpl_propertylist_delete(this_proplist);
    }
    
    return 0;
}

/*----------------------------------------------------------------------------*/
/**
  @brief    Computes the median of the object frames 
  @param    objframes   The objframes to get the background from
  @param    bkg         The background image
  @param    proplist    The list of keywords to store the association
                        information to  
  @return   0 if everything is ok, -1 otherwise

  This function takes as an input a set of object frames and computes
  the background using a median of these frames.
 */
/*----------------------------------------------------------------------------*/
int hawki_bkg_from_objects_median
(const cpl_frameset  *  objframes, cpl_imagelist * bkg)
{
    int             idet;

    /* Error state variables */
    cpl_errorstate  prestate = cpl_errorstate_get();
    
    /* Loop on detectors */
    for(idet = 0; idet < HAWKI_NB_DETECTORS ; ++idet)
    {
        cpl_imagelist * img_serie;
        cpl_image     * this_bkg_image;
        
        /* Loading the object frame */
        img_serie = hawki_load_detector(objframes, idet + 1, CPL_TYPE_FLOAT);
        if(img_serie== NULL)
        {
            cpl_msg_error(__func__, "Error reading object image") ;
            return -1;
        }
        
        /* Averaging */
        /* TODO: Comprobar que cpl_imagelist_collapse_median_create
         * admite las mascaras, para que cuando implemente lo de la mascara 
         * de objetos siga funcionando */ 
        if ((this_bkg_image = cpl_imagelist_collapse_median_create(img_serie))
                == NULL)
        {
            cpl_msg_error(__func__, "Cannot compute the median of obj images");
            cpl_imagelist_delete(img_serie);
            return -1;
        }

        /* Fill the corresponding chip */
        cpl_imagelist_set(bkg, this_bkg_image, idet);
        
        /* Freeing */
        cpl_imagelist_delete(img_serie);
    }

    /* Check error status and exit */
    if(!cpl_errorstate_is_equal(prestate))
        return -1;
    return 0;
}

/*----------------------------------------------------------------------------*/
/**
  @brief    Computes the median of the sky frames 
  @param    skyframes   The skyframes to get the background from
  @param    objframes   The objframes to which the association is made
  @param    bkg         The background image
  @param    proplist    The list of keywords to store the association
                        information to  
  @return   0 if everything is ok, -1 otherwise

  This function takes as an input a set of sky frames and computes
  the background using a median of these frames.
 */
/*----------------------------------------------------------------------------*/
int hawki_bkg_from_sky_median
(const cpl_frameset  *  skyframes, 
 cpl_imagelist       *  bkg)
{
    int             idet;

    /* Error state variables */
    cpl_errorstate  prestate = cpl_errorstate_get();
    
    /* Loop on detectors */
    for(idet = 0; idet < HAWKI_NB_DETECTORS ; ++idet)
    {
        cpl_imagelist * img_serie;
        cpl_image     * this_bkg_image;
        
        /* Loading the object frame */
        img_serie = hawki_load_detector(skyframes, idet + 1, CPL_TYPE_FLOAT);
        if(img_serie== NULL)
        {
            cpl_msg_error(__func__, "Error reading object image") ;
            return -1;
        }
        
        /* Averaging */
        if ((this_bkg_image = cpl_imagelist_collapse_median_create(img_serie))
                == NULL) 
        {
            cpl_msg_error(__func__, "Cannot compute the median of obj images");
            cpl_imagelist_delete(img_serie);
            return -1;
        }
        /* Fill the corresponding chip */
        cpl_imagelist_set(bkg, this_bkg_image, idet);
        
        /* Freeing */
        cpl_imagelist_delete(img_serie);
    }

    /* Check error status and exit */
    if(!cpl_errorstate_is_equal(prestate))
        return -1;
    return 0;
}

/*----------------------------------------------------------------------------*/
/**
  @brief  Computes the running mean for a set of object images using their masks  
  @param    objimages   The serie of object images
  @param    medians     A vector with the medians of each image
  @param    i_target    The index in the list for which we compute the bkg 
  @param    half_width  half with of the running window
  @param    rejlow      Number of pixels with lower values to reject
  @param    rejhigh     Number of pixels with higher values to reject
  @param    bkg         The background image (OUT)
  @return   0 if everything is ok, -1 otherwise

  This function takes as an input a set of object images and computes
  the background using a running median over this images.
  Given the target image i_target, a window in the imagelist is constructed
  from i_target - half_window to i_target + half_window. The value
  of the bkg for each pixel is constructed from the median of all the values
  in the window of the image list, rejecting those rejlow lower values and 
  rejhigh higher values (bad pixels are discarded before that). 
 */
/*----------------------------------------------------------------------------*/
int hawki_bkg_from_running_mean
(cpl_imagelist       * objimages,
 const cpl_vector    * medians,
 int                   i_target,
 int                   half_width,
 int                   rejlow,
 int                   rejhigh,
 cpl_image          *  bkg)
{
    float        *  bkg_p;
    float        *  curima_p;
    cpl_binary   *  curmask_p;
    float        ** objimages_p;
    cpl_binary   ** maskimages_p;
    const double *  medians_p;
    double          out;
    int             from_ima;
    int             to_ima;
    int             nima;
    int             nx;
    int             ny;
    
    int             pos_x;
    int             pos_y;
    int             pos;
    int             i_ima;
    int             i_win;
    int             n_win;
    
    std::vector<double> localwin;

    /* Get the pointer to the background image */
    bkg_p = cpl_image_get_data_float(bkg);
    
    /* Compute border indices */
    nima = cpl_imagelist_get_size(objimages);
    from_ima = i_target - half_width;
    to_ima   = i_target + half_width;
    if (from_ima<0) from_ima = 0 ;
    if (to_ima>(nima-1)) to_ima=nima-1 ;

    /* Get the size of images */
    nx = cpl_image_get_size_x(bkg);
    ny = cpl_image_get_size_y(bkg);
    
    /* Get the pointers for better performance */
    medians_p = cpl_vector_get_data_const(medians);

    /* Get fast accesors to input image pointers */
    objimages_p  = (float **)cpl_malloc(nima * sizeof(float *));
    maskimages_p = (cpl_binary **)cpl_malloc(nima * sizeof(cpl_binary *));
    for(i_ima=from_ima ; i_ima<=to_ima ; i_ima++)
    {
        objimages_p[i_ima]  = cpl_image_get_data_float
            (cpl_imagelist_get(objimages, i_ima));
        maskimages_p[i_ima] = cpl_mask_get_data
            (cpl_image_get_bpm(cpl_imagelist_get(objimages, i_ima)));
    }
    
    /* Initialize the bad pixel mask */
    cpl_image_accept_all(bkg);

    /* Loop over all pixels */
    for (pos_x=0 ; pos_x<nx ; pos_x++)
    {
        for (pos_y=0 ; pos_y<ny ; pos_y++)
        {
            /* Get the position in the array */
            pos = pos_x + pos_y * nx;
            
            /* Empty the local window */
            localwin.clear();
        
            /* Fill up local window */
            for (i_ima=from_ima ; i_ima<=to_ima ; i_ima++) 
            {
                if (i_ima != i_target) 
                {
                    curima_p = objimages_p[i_ima];
                    curmask_p = maskimages_p[i_ima];
                    if(!curmask_p[pos])
                        localwin.push_back
                            ((double)curima_p[pos]-medians_p[i_ima]);
                }
            }
            n_win = localwin.size();
            if(n_win - rejlow - rejhigh < 1)
            {
                /* Add it to the bad pixel mask */
                cpl_msg_debug(__func__,"Pixel %d %d added to the sky bpm",
                        pos_x, pos_y);
                if(cpl_image_reject(bkg, pos_x + 1, pos_y + 1) != CPL_ERROR_NONE)
                {
                    cpl_msg_error(__func__,"Cannot add pixel to sky bpm");
                    cpl_free(objimages_p);
                    cpl_free(maskimages_p);
                    return(-1);
                }
            }
            else
            {
                /* Sort window */
                std::sort(localwin.begin(), localwin.end());
                
                /* Reject min and max, accumulate other pixels */
                out = 0.0 ;
                for (i_win=rejlow ; i_win<n_win - rejhigh; i_win++) 
                    out += localwin[i_win];
                /* Take the mean */
                out /= (double)(n_win - rejlow - rejhigh);
        
                /* Assign value */
                bkg_p[pos] = (out+medians_p[i_target]);
            }
        }
    }

    /* Free */
    cpl_free(objimages_p);
    cpl_free(maskimages_p);

    return 0;
}

/*----------------------------------------------------------------------------*/
/**
  @brief  Computes the running mean for a set of object frames using their masks  
  @param    objframes   The serie of object frames
  @param    bkg         The background image (OUT)
  @return   0 if everything is ok, -1 otherwise

  This function takes as an input a set of object frames and computes
  the background using a running median of these frames. 
  It loads in memory only the images needed for the running window.
 */
/*----------------------------------------------------------------------------*/
int hawki_bkg_from_running_mean_frame_extension
(hawki_bkg_frames_buffer * frames_buffer,
 cpl_vector              * offsets_x,
 cpl_vector              * offsets_y,
 cpl_image               * globalmask,
 double                    mask_off_x,
 double                    mask_off_y,
 cpl_image               * distortion_x,
 cpl_image               * distortion_y,
 int                       iextension,
 int                       i_target,
 int                       half_width,
 int                       rejlow,
 int                       rejhigh,
 cpl_image               * bkg)
{
    cpl_imagelist   * imagelist_window;
    cpl_vector      * medians_window;

    cpl_size        from_ima;
    cpl_size        to_ima;
    cpl_size        i_ima;
    cpl_size        nima;
    
    nima = frames_buffer->nframes;
    from_ima = i_target - half_width;
    to_ima   = i_target + half_width;
    if (from_ima<0) from_ima = 0 ;
    if (to_ima>(nima-1)) to_ima=nima-1;

    /* Assume that the previous run of this function was called
     * with i_target -1, and therefore we deallocate the first image of the
     * previous window to free space.
     * If the assumption is wrong, then probably there is a performance
     * penalty reading again that image.
     * The second line is asuming that the previous run was for i_target + 1
     * (running the loop in the opposite direction)
     * If the upper loop which calls this function runs in steps of 1, then
     * there is no performance penalty.
     */
    if(from_ima > 0)
    {
        if(frames_buffer->images[from_ima - 1] != NULL)
          cpl_image_delete(frames_buffer->images[from_ima - 1]);
        frames_buffer->images[from_ima - 1] = NULL;
    }
    if(to_ima < nima - 1)
    {
        if(frames_buffer->images[to_ima] != NULL)
          cpl_image_delete(frames_buffer->images[to_ima]);
        frames_buffer->images[to_ima] = NULL;
    }

    imagelist_window = cpl_imagelist_new();
    medians_window   = cpl_vector_new(to_ima - from_ima + 1);

    for(i_ima=from_ima; i_ima<=to_ima ; i_ima++)
    {
        if(frames_buffer->images[i_ima] == NULL)
        {
        	cpl_image * new_image;
        	cpl_frame * frame;

        	frame = cpl_frameset_get_position(frames_buffer->frames, i_ima);
        	new_image = hawki_load_frame_extension
        			(frame, iextension + 1, CPL_TYPE_FLOAT);
        	frames_buffer->images[i_ima] = new_image;

        	/* Creating a mask for this object, using the offsets 
             * and the distortion (if it applies) */
        	if(globalmask != NULL)
        	{
                double target_off_x;
                double target_off_y;
        	    target_off_x = cpl_vector_get(offsets_x, i_ima);
        	    target_off_y = cpl_vector_get(offsets_y, i_ima);
        	    hawki_bkg_set_obj_mask(new_image, globalmask,
        	                           distortion_x, distortion_y,
        	                           target_off_x, target_off_y,
        	                           mask_off_x, mask_off_y);
        	}
        	
        	/* Computing the median */
            frames_buffer->medians[i_ima] = cpl_image_get_median(new_image);
        }
        cpl_imagelist_set(imagelist_window, frames_buffer->images[i_ima],
                          i_ima - from_ima);
        cpl_vector_set(medians_window, i_ima - from_ima,
        		       frames_buffer->medians[i_ima]);
    }

    hawki_bkg_from_running_mean(imagelist_window, medians_window,
    		 i_target - from_ima, half_width, rejlow, rejhigh, bkg);
    
    /* Free. The images in imagelist_window are unset but not deallocated */
    for(i_ima=from_ima; i_ima<=to_ima ; i_ima++)
        cpl_imagelist_unset(imagelist_window, 0);
    cpl_imagelist_delete(imagelist_window);
    cpl_vector_delete(medians_window);

    return 0;
}

int hawki_bkg_set_obj_mask
(cpl_image  * target_image,
 cpl_image  * globalmask,
 cpl_image  * distor_x,
 cpl_image  * distor_y,
 double       target_off_x,
 double       target_off_y,
 double       mask_off_x,
 double       mask_off_y)
{
    cpl_image  * mask_intersection;
    cpl_mask   * mask_final;

    /* The new mask is in float because the distortion correction will
     * work on float data */
    mask_intersection = cpl_image_new(cpl_image_get_size_x(target_image),
                                      cpl_image_get_size_y(target_image),
                                      CPL_TYPE_FLOAT);
    
    /* The offsets are rounded to pixels here */
    hawki_image_copy_to_intersection(mask_intersection,
                                     globalmask,
                                     (cpl_size)(target_off_x - mask_off_x),
                                     (cpl_size)(target_off_y - mask_off_y));
    if(distor_x != NULL && distor_y != NULL)
    {
        cpl_image * mask_distcorr;

        /* Dedistort the mask */
        mask_distcorr = hawki_distortion_correct_detector
                (mask_intersection, distor_x, distor_y);
        if(mask_distcorr == NULL)
        {
            cpl_msg_error(__func__, "Cannot correct the distortion") ;
            cpl_image_delete(mask_intersection);
            cpl_msg_indent_less();
            cpl_msg_indent_less();
            return -1 ;
        }
        cpl_image_delete(mask_intersection);
        mask_intersection = mask_distcorr;
    }
    mask_final = 
            cpl_mask_threshold_image_create(mask_intersection, 0.5, FLT_MAX);
    /* TODO: Add the current bpm to this mask? */
    cpl_image_reject_from_mask(target_image, mask_final);
    
    cpl_image_delete(mask_intersection);
    cpl_mask_delete(mask_final);
}

/**@}*/

