/* $Id: hdrl_flat.c,v 1.29 2013-09-25 11:04:29 jtaylor Exp $
 *
 * This file is part of the HDRL
 * Copyright (C) 2013 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., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

/*
 * $Author: jtaylor $
 * $Date: 2013-09-25 11:04:29 $
 * $Revision: 1.29 $
 * $Name: not supported by cvs2svn $
 */

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

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

#include "hdrl_flat.h"

#include "hdrl_collapse.h"
#include "hdrl_combine.h"
#include "hdrl_elemop.h"
#include "hdrl_prototyping.h"
#include "hdrl_utils.h"

#include <cpl.h>

#include <assert.h>
#include <stdint.h>

/*-----------------------------------------------------------------------------
                                   Static
 -----------------------------------------------------------------------------*/

typedef enum {
    _HDRL_MEDIAN_SMOOTH,
    _HDRL_GAUSS_SMOOTH,
    _HDRL_MEDIAN_STEP_FIT
} _hdrl_internal_method;

static cpl_error_code hdrl_compute_flat(const cpl_imagelist *, 
        const cpl_imagelist *, hdrl_collapse_imagelist_to_image_t *,
        cpl_image *, cpl_image *, cpl_image **, cpl_image **, cpl_image **, 
        _hdrl_internal_method, void *);
static cpl_image * hdrl_medianclean_image(cpl_image *, cpl_size, cpl_size, 
        cpl_size, cpl_size);
static cpl_image * hdrl_smooth_image_median(const cpl_image *, int);

/** @cond PRIVATE */

/*----------------------------------------------------------------------------*/
/**
  @defgroup hdrl_flat   Flat Module
  This module provides functionality to compute a master flat frame out of
  several inputs. Several methods are available to deal with different flat
  characteristics

 */
/*----------------------------------------------------------------------------*/

/**@{*/

typedef struct {
    cpl_size filter_size;
} _hdrl_median_smooth_par;

/* ---------------------------------------------------------------------------*/
/**
 * @brief compute master flat with median filtering on exposure map
 *
 * @param data               input flats
 * @param errs               input flats errors
 * @param collapse_mode      method to collapse to master flat
 * @param filter_size        size of median filter
 * @param master_shape       optional initial master shape
 * @param master_shape_err   optional initial master shape error
 * @param flat_img           output flat
 * @param flat_err           output flat error
 * @param flat_ctr           output flat contribution map
 *
 * TODO describe algorithm (refer to called lower level function)
 * iterative + median filter on exposure map (= flat / masterflat)
 */
/* ---------------------------------------------------------------------------*/
cpl_error_code hdrl_compute_flat_median_smooth(
        const cpl_imagelist                 *   data,
        const cpl_imagelist                 *   errs,
        hdrl_collapse_imagelist_to_image_t  *   collapse_mode,
        cpl_size                                filter_size,
        cpl_image                           *   master_shape,
        cpl_image                           *   master_shape_err,
        cpl_image                           **  flat_img, 
        cpl_image                           **  flat_err,
        cpl_image                           **  flat_ctr)
{
    _hdrl_median_smooth_par par;
    par.filter_size = filter_size;
    return hdrl_compute_flat(data, errs, collapse_mode, master_shape,
                             master_shape_err, flat_img, flat_err, flat_ctr,
                             _HDRL_MEDIAN_SMOOTH, &par);
}

typedef struct {
    cpl_size filter_size;
    cpl_size mirror_x;
    cpl_size mirror_y;
} _hdrl_gauss_smooth_par;

/* ---------------------------------------------------------------------------*/
/**
 * @brief compute master flat with Gauss smoothing on exposure map
 *
 * @param data               input flats
 * @param errs               input flats errors
 * @param collapse_mode      method to collapse to master flat
 * @param filter_size        size of median filter
 * @param mirror_x           implemented as FFT with mirroring
 * @param mirror_y           implemented as FFT with mirroring
 * @param master_shape       optional initial master shape
 * @param master_shape_err   optional initial master shape error
 * @param flat_img           output flat
 * @param flat_err           output flat error
 * @param flat_ctr           output flat contribution map
 *
 * TODO describe algorithm (refer to called lower level function)
 * iterative + Gauss filter on exposure map (= flat / masterflat)
 * implemented as FFT thus needs mirroring values
 * should probably be implemented in normal space with decomposed Gauss instead
 */
/* ---------------------------------------------------------------------------*/
cpl_error_code hdrl_compute_flat_gauss_smooth(
        const cpl_imagelist                 *   data,
        const cpl_imagelist                 *   errs,
        hdrl_collapse_imagelist_to_image_t  *   collapse_mode,
        cpl_size                                filter_size,
        cpl_size                                mirror_x,
        cpl_size                                mirror_y,
        cpl_image                           *   master_shape,
        cpl_image                           *   master_shape_err,
        cpl_image                           **  flat_img, 
        cpl_image                           **  flat_err,
        cpl_image                           **  flat_ctr)
{
    _hdrl_gauss_smooth_par par;
    par.filter_size = filter_size;
    par.mirror_x = mirror_x;
    par.mirror_y = mirror_y;
    return hdrl_compute_flat(data, errs, collapse_mode, master_shape,
                             master_shape_err, flat_img, flat_err, flat_ctr,
                             _HDRL_GAUSS_SMOOTH, &par);
}

typedef struct {
    cpl_size    filter_size_x;
    cpl_size    filter_size_y;
    cpl_size    steps_x;
    cpl_size    steps_y;
    int         order_x;
    int         order_y;
} _hdrl_step_fit_par;

/* ---------------------------------------------------------------------------*/
/**
 * @brief compute master flat with polynomial fit to exposure map
 *
 * @param data               input flats
 * @param errs               input flats errors
 * @param collapse_mode      method to collapse to master flat
 * @param filter_size_x      size of median filter in x
 * @param filter_size_y      size of median filter in y
 * @param steps_x            number of steps in x to evaluate median on
 * @param steps_y            number of steps in y to evaluate median on
 * @param order_x            order of the polynomial in x
 * @param order_y            order of the polynomial in y
 * @param master_shape       optional initial master shape
 * @param master_shape_err   optional initial master shape error
 * @param flat_img           output flat
 * @param flat_err           output flat error
 * @param flat_ctr           output flat contribution map
 *
 * TODO describe algorithm (refer to called lower level function)
 * iterative + strided median filter on exposure map (= flat / masterflat) +
 * polyfit to filtered image
 */
/* ---------------------------------------------------------------------------*/
cpl_error_code hdrl_compute_flat_median_step_fit(
        const cpl_imagelist                 *   data,
        const cpl_imagelist                 *   errs,
        hdrl_collapse_imagelist_to_image_t  *   collapse_mode,
        cpl_size                                filter_size_x,
        cpl_size                                filter_size_y,
        cpl_size                                steps_x,
        cpl_size                                steps_y,
        int                                     order_x,
        int                                     order_y,
        cpl_image                           *   master_shape,
        cpl_image                           *   master_shape_err,
        cpl_image                           **  flat_img,
        cpl_image                           **  flat_err,
        cpl_image                           **  flat_ctr)
{
    _hdrl_step_fit_par par;
    /* TODO: should we check that input have proper input values ? */
    /* TODO: add proper comment, for example
     * init par structure parameters
     */
    par.filter_size_x = filter_size_x;
    par.filter_size_y = filter_size_y;
    par.steps_x = steps_x;
    par.steps_y = steps_y;
    par.order_x = order_x;
    par.order_y = order_y;
    return hdrl_compute_flat(data, errs, collapse_mode, master_shape, 
            master_shape_err, flat_img, flat_err, flat_ctr, 
            _HDRL_MEDIAN_STEP_FIT, &par);
}

/* ---------------------------------------------------------------------------*/
/**
 * @internal
 * @brief compute master flat
 * @param data             bias and dark corrected image data
 * @param errs             errors of data
 * @param collapse_mode    method to collapse data to flag
 * @param master_shape     shape where the single flats are bend towards
 * @param master_shape_err error of the shape where each flat is bend towards
 * @param flat_img         resulting flat
 * @param flat_err         error of flat
 * @param flat_ctr         contribution map of flat
 * @param method           method switch to choose image smoothing algorithm
 * @param par              image smoothing algorithm parameters
 * 
 * DONE (AMO): describe algorithm
 *
 * -check input frames are proper
 *
 * -before the final master frame is generate make sure each flat component
 * has the same level and shape. In order to compute properly each flat level
 * its shape need to be removed. As the shape information is contained in the
 * master flat that is also a product of this data reduction step we need to
 * perform iterations. The user may decide to provide an input shape (if
 * available) or let the function to find one. In this last case a master flat
 * is computed by simple mean-stack of each input frame (without error
 * propagation, as the error information is actually needed only from the last
 * iteration).
 *
 * For each iteration:
 *
 * -we divide each flat frame for the current master flat to remove the shape of
 * the flat (at last iteration the division is performed with error propagation)
 *
 * -we compute the flat estimator (a single number) using one of the possible
 * criteria (median-smooth over a rectangular box with interpolation of rejected
 * pixels, median-step over a rectangular box with fit of a Legendre polynomial
 * in each box, Gaussian smooth over rectangular box)
 *
 * Finally we divide each imagelist by the final master with variance
 * propagation and combine them into the final master flat using any of the
 * available methods(mean, weighted mean, median, sigma-clip)
 */
/* ---------------------------------------------------------------------------*/
static cpl_error_code hdrl_compute_flat(
        const cpl_imagelist                 *   data, 
        const cpl_imagelist                 *   errs,
        hdrl_collapse_imagelist_to_image_t  *   collapse_mode,
        cpl_image                           *   master_shape, 
        cpl_image                           *   master_shape_err,
        cpl_image                           **  flat_img, 
        cpl_image                           **  flat_err,
        cpl_image                           **  flat_ctr, 
        _hdrl_internal_method                   method,
        void                                *   par)
{
    cpl_image * comb_img = NULL;
    cpl_image * comb_err = NULL;
    cpl_image * comb_ctr = NULL;
    /* TODO: should we check that input have proper input values ? */
    /* TODO: quite long function implementation. Should we re-factor it */
    /* TODO should the following niter_max a hard-coded variable?
     * better a const or a define if not an algorithm parameter
     */
    size_t niter_max = 3;

    /* TODO: check that the size of the master_shape and master_shape_err
     * is the same as the size of the image in the imagelists*/

    cpl_msg_info(cpl_func, "Creating first Master Flat");
    cpl_msg_indent_more();

    /*If we pass a master-shape image to the function use this as
     * first master flat*/
    if (master_shape != NULL) {
        comb_img = cpl_image_duplicate(master_shape);
        /*If we pass a master-shape image to the function do only one
         * iteration*/
        niter_max = 1;
        /* optional error of the shape */
        if (master_shape_err != NULL) {
            comb_err = cpl_image_duplicate(master_shape_err);
        }
    } else {
        /* Generate 1st master flat */
        /*No error propagation needed in the first iteration*/
        cpl_imagelist * data_loc = cpl_imagelist_duplicate(data);
        cpl_msg_info(cpl_func, "Scale all flats to unity by the mean");
        cpl_imagelist_normalise(data_loc, CPL_NORM_MEAN);
        cpl_msg_info(cpl_func, "Combining all scaled flats by the median");
        comb_img = cpl_imagelist_collapse_median_create(data_loc);
        cpl_imagelist_delete(data_loc);
        comb_err = cpl_image_duplicate(comb_img);
        cpl_image_multiply_scalar(comb_err, 0);
    }
    cpl_msg_indent_less();

    cpl_msg_info(cpl_func, "Start iterative process");
    cpl_msg_indent_more();
    /* TODO: should we factor-out the following long loop ? May be not as it is
     * de facto the full job of this function */
    /* Here we perform several iterations to optimize the master flat shape */
    for (size_t i = 0; i < niter_max; i++) {
        cpl_msg_indent_less();
        cpl_msg_info(cpl_func, "Iteration %zu out of %zu", i + 1, niter_max);
        cpl_msg_indent_more();

        /*Restore original flat-fields but use the new master-flat*/
        cpl_msg_info(cpl_func, "Restore the original flatfields");
        cpl_imagelist * imglist = cpl_imagelist_duplicate(data);
        cpl_imagelist * errlist = cpl_imagelist_duplicate(errs);

        /*Divide flat-fields through the masterflat -> esposure-time map*/
        cpl_msg_info(cpl_func, "Divide single flatfields through the (new) "
                        "masterflat to get an exposure-time map for each "
                        "flatfield");
        /* only propagate error of master flat in last iteration */
        if (i == niter_max - 1) {
            /* TODO: function name. Here we perform division of two imagelists
             * with proper error (and bad pixel) propagation. This is done via
             * auxilliary arrays. If we follow cpl naming convention may be
             * better (clear but long) name is:
             * hdrl_imagelist_divide_errorpropagate ?
             */
            hdrl_elemop_imagelist_div_image(imglist, errlist, comb_img, comb_err);
        } else {
            cpl_imagelist_divide_image(imglist, comb_img);
        }

        /*  delete old masterframe and compute the new master frame */
        cpl_image_delete(comb_img);
        cpl_image_delete(comb_err);
        cpl_image_delete(comb_ctr);

        /* smooth residuals  */
        cpl_msg_info(cpl_func, "Smooth the exposure-time maps");

        for(cpl_size j = 0; j < cpl_imagelist_get_size(imglist); j++){
            cpl_msg_info(cpl_func, "working on image %d", (int)j+1);
            cpl_image * imgtmp = cpl_imagelist_get(imglist, j);
            cpl_image * r;
            if (method == _HDRL_MEDIAN_SMOOTH) {
                /* Method 1*/
                /* removing bad pixels allows use of a much faster algorithm */
                cpl_size fsize = ((_hdrl_median_smooth_par *)par)->filter_size;
                cpl_detector_interpolate_rejected(imgtmp);
                /* TODO: naming conventions. Should we name functions as
                 * hdrl_object_operation_specifier or
                 * hdrl_operation_object_specifier?
                 */
                r = hdrl_smooth_image_median(imgtmp, fsize);
            } else if (method == _HDRL_GAUSS_SMOOTH) {
                _hdrl_gauss_smooth_par * p = (_hdrl_gauss_smooth_par *)par;
                /* TODO: naming convention should the function name contain the
                 * object it operates on==>hdrl_image_get_spatial_freq()
                 */
                r = hdrl_get_spatial_freq(imgtmp, p->filter_size,
                                          p->mirror_x, p->mirror_y);
            } else if (method == _HDRL_MEDIAN_STEP_FIT) {
                _hdrl_step_fit_par * p = (_hdrl_step_fit_par *)par;
                cpl_size nx = cpl_image_get_size_x(imgtmp);
                cpl_size ny = cpl_image_get_size_y(imgtmp);
                cpl_size sx = CX_MAX(nx / p->steps_x, 1);
                cpl_size sy = CX_MAX(ny / p->steps_y, 1);

                /* fit to stepped grid */
                cpl_matrix * x = hdrl_matrix_linspace(sx / 2, nx, sx);
                cpl_matrix * y = hdrl_matrix_linspace(sy / 2, ny, sy);
                cpl_image * imgtmp_mod =
                    hdrl_medianfilter_image_grid(imgtmp, x, y,
                                                 p->filter_size_x,
                                                 p->filter_size_y);
                cpl_matrix * coeffs = hdrl_fit_legendre(imgtmp_mod,
                                                        p->order_x, p->order_y,
                                                        x, y, nx, ny);
                /* TODO: naming conventions: hdrl_matrix_legendre_to_image() */
                r = hdrl_legendre_to_image(coeffs,
                                           p->order_x, p->order_y,
                                           nx, ny);

                if (cpl_msg_get_level() == CPL_MSG_DEBUG)
                    cpl_matrix_dump(coeffs, stdout);

                cpl_matrix_delete(coeffs);
                cpl_matrix_delete(x);
                cpl_matrix_delete(y);
                cpl_image_delete(imgtmp_mod);
            } else {
                abort();
            }
            cpl_imagelist_set(imglist, r, j);
        }

        cpl_msg_info(cpl_func, "Divide the original flatfields by the smoothed "
                        "exposure-time map");

        /* Divide flatfields through the smoothed exposure-time map*/
        cpl_imagelist * div_imglist = cpl_imagelist_duplicate(data);
        cpl_imagelist * div_errlist = cpl_imagelist_duplicate(errs);
        hdrl_elemop_imagelist_div_imagelist(div_imglist, div_errlist,
                                            imglist, errlist);
        cpl_imagelist_delete(imglist);
        cpl_imagelist_delete(errlist);

        cpl_msg_info(cpl_func, "Combine the resulting normalized flatfields by "
                        "your preferred algorithm and generate a new(!)"
                        " masterflat");

        hdrl_imagelist_combine(div_imglist, div_errlist, collapse_mode,
                               &comb_img, &comb_err, &comb_ctr);

        cpl_imagelist_delete(div_imglist);
        cpl_imagelist_delete(div_errlist);

        if (cpl_error_get_code()) break;
    }

    if (cpl_error_get_code()) {
        cpl_image_delete(comb_img);
        cpl_image_delete(comb_err);
        cpl_image_delete(comb_ctr);
    } else {
        *flat_img = comb_img;
        *flat_err = comb_err;
        *flat_ctr = comb_ctr;
    }

    cpl_msg_indent_less();
    return cpl_error_get_code();
}

/* ---------------------------------------------------------------------------*/
/**
 * @brief computes the median in predefined regions
 *
 * @param ima           image with the data
 * @param steps_x       number of steps in x where the median is calculated
 * @param steps_y       number of steps in y where the median is calculated
 * @param filtersize_x  half-size median-box in x where the median is calculated
 * @param filtersize_y  half-size median-box in y where the median is calculated
 *
 * @return  median filtered image
 *
 * DONE (AMo): describe algorithm
 * This function computes an image where each pixel value is obtained by a median
 * every x and y steps in an image of half-box size filtersize_x, filtersize_y.
 * It is assumed that the resulting bad pixel is good.
 * TODO: what happens if some of the full boxes over which the median is
 * computed is full of bad pixels?
 */
/* ---------------------------------------------------------------------------*/
static inline cpl_image * hdrl_medianclean_image(
        cpl_image   *   ima, 
        cpl_size        steps_x, 
        cpl_size        steps_y,
        cpl_size        filtersize_x, 
        cpl_size        filtersize_y) 
{
    cpl_error_ensure(ima != NULL, CPL_ERROR_NULL_INPUT,
            return NULL, "NULL input image");
    cpl_error_ensure(steps_x > 0 && steps_y > 0 && filtersize_x > 0 &&
            filtersize_y > 0 , CPL_ERROR_INCOMPATIBLE_INPUT,
            return NULL, "All function parameters must be greater then Zero");

    const cpl_size nx = cpl_image_get_size_x(ima);
    const cpl_size ny = cpl_image_get_size_y(ima);
    const cpl_size stepsize_x = (cpl_size)(nx / steps_x);
    cpl_size stepsize_y = (cpl_size)(ny / steps_y);
    cpl_image * ima_local = cpl_image_new(nx, ny, CPL_TYPE_DOUBLE);
    cpl_mask_not(cpl_image_get_bpm(ima_local));

    for (cpl_size iy = 0; iy < steps_y; iy++) {
        for (cpl_size ix = 0; ix < steps_x; ix++) {
            cpl_size middlep_x = (cpl_size)(stepsize_x * ix + stepsize_x / 2.);
            cpl_size middlep_y = (cpl_size)(stepsize_y * iy + stepsize_y / 2.);

            cpl_size lowerlimit_x = CX_MAX(middlep_x - filtersize_x, 1);
            cpl_size lowerlimit_y = CX_MAX(middlep_y - filtersize_y, 1);
            cpl_size upperlimit_x = CX_MIN(middlep_x + filtersize_x, nx);
            cpl_size upperlimit_y = CX_MIN(middlep_y + filtersize_y, ny);

            double median = cpl_image_get_median_window(ima, lowerlimit_x,
                            lowerlimit_y, upperlimit_x, upperlimit_y);

            cpl_image_set(ima_local, middlep_x, middlep_y, median);
            cpl_image_accept(ima_local, middlep_x, middlep_y);

            cpl_msg_debug(cpl_func, "middlep_x: %lld, middlep_y: %lld, median: "
                          "%g", middlep_x, middlep_y, median);
        }
    }
    return ima_local;
}

/*---------------------------------------------------------------------------*/
/**
  @brief    Smooth an image using a median filter
  @param    in      image to be smoothed
  @param    size    size of kernel matrix to use in filtering
  @return   out     smoothed image
 */
/*---------------------------------------------------------------------------*/
static cpl_image * hdrl_smooth_image_median(
        const cpl_image     *   in, 
        int                     size)
{
    cpl_image * out = cpl_image_duplicate(in);
    cpl_mask * kernel = cpl_mask_new(size, size); 
    cpl_mask_not(kernel);

    if(cpl_image_filter_mask(out, in, kernel, CPL_FILTER_MEDIAN,
                             CPL_BORDER_FILTER) != CPL_ERROR_NONE) {
        cpl_image_delete(out);
        cpl_mask_delete(kernel);
        return NULL;
    }

    cpl_mask_delete(kernel);
    return (out);
}



/**@}*/

/** @endcond */
