/* $Id: fors_flat_normalise.cc,v 1.9 2013-10-24 16:44:34 cgarcia Exp $
 *
 * This file is part of the MOSES library
 * Copyright (C) 2002-2010 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: cgarcia $
 * $Date: 2013-10-24 16:44:34 $
 * $Revision: 1.9 $
 * $Name: not supported by cvs2svn $
 */

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

#include <cmath>
#include <fors_flat_normalise.h>
#include <moses.h>
#include <image_smooth.h>
#include <image_spline_fit.h>

#define STRETCH_FACTOR   (1.20)

/**
 * @brief 
 *   Normalise a flat field exposure
 *  
 * @param flat        Image containing the original flat field spectra
 * @param spatial     Spatial calibration image
 * @param slits       Table with slits positions
 * @param polytraces  Coefficients of spectral curvature polynomials
 * @param reference   Reference wavelength
 * @param blue        Start lambda to process
 * @param red         End lambda to process
 * @param dispersion  Mean spectral dispersion
 * @param sradius     Radius of smoothing box along the dispersion direction
 * @param polyorder   Order of fitting polynomial along the dispersion
 *  
 * @return The smoothed flat field exposure used for normalisation
 *
 * The input @em flat frame should be already bias subtracted, and should
 * be oriented so that the dispersion direction is horizontal with @em blue
 * on the left and @em red on the right. The flat field spectra are spatially 
 * rectified, heavily smoothed, and then mapped back on the CCD. The original 
 * @em flat image is divided IN PLACE by its smoothed counterpart, which is 
 * also returned. If the polynomial @em polyorder is set to a negative number 
 * the smoothing consists of a linear fit along the spatial direction 
 * (excluding 3+3 pixels at the spectral edges), and by a median filtering 
 * along the dispersion direction using a window with the specified 
 * @em sradius; alternatively, if @em polyorder is not negative, the smoothing
 * will consist of a polynomial fitting of the illumination profile along
 * the dispersion direction, performed independently for each row of the
 * spatially remapped spectra.
 */
 
cpl_image *mos_mosflat_normalise(cpl_image *flat, cpl_image *spatial, 
                                 cpl_table *slits, cpl_table *polytraces, 
                                 double reference, double blue, double red, 
                                 double dispersion,
                                 int spa_smooth_radius, int disp_smooth_radius, 
                                 int spa_fit_nknots, int disp_fit_nknots,
                                 double fit_threshold)
{
    const char     *func = "mos_mosflat_normalise";

    const char *clab[6] = {"c0", "c1", "c2", "c3", "c4", "c5"};
                                                 /* Max order is 5 */

    cpl_image      *rectified;
    cpl_image      *smo_flat;
    cpl_image      *exslit;
    cpl_polynomial *polytop;
    cpl_polynomial *polybot;

    int            *slit_id;
    float          *sdata;
    float          *xdata;
    float          *wdata;
    double          vtop, vbot, value;
    double          top, bot;
    double          coeff;
    double          ytop, ybot;
    double          ypos;
    double          fvalue;
    int             ivalue;
    int             yint, yprev = 0;
    int             npseudo;

    int             pixel_above, pixel_below, refpixel, start_pixel, end_pixel;
    int             nx, ny;
    int             xlow, ylow, xhig, yhig;
    int             nslits;
    int            *position;
    int            *length;
    int             missing_top, missing_bot;
    int             order;
    int             null;
    int             i, j;
    cpl_size        k;

/*    int             exclude = 5;     Number of excluded pixels at edges */

    /* For debug puposes only: cpl_image      *smo_rectified; */

    if (flat == NULL || slits == NULL || polytraces == NULL) {
        cpl_error_set(func, CPL_ERROR_NULL_INPUT);
        return NULL;
    }
 
    if (dispersion <= 0.0) {
        cpl_error_set(func, CPL_ERROR_ILLEGAL_INPUT);
        return NULL;
    }

    if (red - blue < dispersion) {
        cpl_error_set(func, CPL_ERROR_ILLEGAL_INPUT);
        return NULL;
    }
    
    rectified = mos_spatial_calibration(flat, slits, polytraces, reference,
                                        blue, red, dispersion, 0, NULL);

    nx = cpl_image_get_size_x(rectified);
    ny = cpl_image_get_size_y(rectified);

    smo_flat = cpl_image_new(cpl_image_get_size_x(spatial), 
                             cpl_image_get_size_y(spatial), CPL_TYPE_FLOAT);
    wdata = cpl_image_get_data_float(smo_flat);

    nslits   = cpl_table_get_nrow(slits);
    order    = cpl_table_get_ncol(polytraces) - 2;
    position = cpl_table_get_data_int(slits, "position");
    length   = cpl_table_get_data_int(slits, "length");
    slit_id  = cpl_table_get_data_int(slits, "slit_id");

    /*
     * The spatial resampling is performed for a certain number of
     * pixels above and below the position of the reference wavelength:
     */
    
    pixel_above = (int)(STRETCH_FACTOR * (red - reference) / dispersion);
    pixel_below = (int)(STRETCH_FACTOR * (reference - blue) / dispersion);

    xlow = 1;
    xhig = nx;
    for (i = 0; i < nslits; i++) {

        if (length[i] == 0)
            continue;

        /*
         * We DON'T write:
         *
         * ylow = position[i];
         * yhig = ylow + length[i];
         *
         * because the cpl_image pixels are counted from 1, and because in 
         * cpl_image_extract() the coordinates of the last pixel are inclusive.
         */

        ylow = position[i] + 1;
        yhig = ylow + length[i] - 1;

        exslit = cpl_image_extract(rectified, xlow, ylow, xhig, yhig);
        mosca::image im_slit(exslit);

        if (spa_smooth_radius > 0) 
        {
            int  final_spa_smooth_radius;
            if (im_slit.size_spatial() / 2 < spa_smooth_radius)
            {
                final_spa_smooth_radius = im_slit.size_spatial() / 2;
                cpl_msg_warning(cpl_func, "Slit too narrow for requested "
                                "smoothing radius %d. Using %d", 
                                spa_smooth_radius, final_spa_smooth_radius);
            }
            else
                final_spa_smooth_radius = spa_smooth_radius;
            mosca::image_smooth_1d_median<float>(im_slit, 
                                                 final_spa_smooth_radius, 
                                                 mosca::SPATIAL_AXIS);
        }
        
        if (spa_fit_nknots > 0)
            mosca::image_cubicspline_1d_fit<float>(im_slit, spa_fit_nknots, 
                                                   fit_threshold,
                                                   mosca::SPATIAL_AXIS);

        if (disp_smooth_radius > 0)
            mosca::image_smooth_1d_median<float>(im_slit, 
                                                 disp_smooth_radius, 
                                                 mosca::DISPERSION_AXIS);
        if (disp_fit_nknots > 0)
            mosca::image_cubicspline_1d_fit<float>(im_slit, disp_fit_nknots, 
                                                   fit_threshold,
                                                   mosca::DISPERSION_AXIS);
        //image_smooth_fit_1d_pol_disp(exslit, disp_fit_polyorder);

        
        /*
         * Recover from the table of spectral curvature coefficients
         * the curvature polynomials.
         */

        refpixel = (int)(cpl_table_get_double(slits, "xtop", i, NULL));

        start_pixel = refpixel - pixel_below;
        if (start_pixel < 0)
            start_pixel = 0;

        end_pixel = refpixel + pixel_above;
        if (end_pixel > nx)
            end_pixel = nx;

        missing_top = 0;
        polytop = cpl_polynomial_new(1);
        for (k = 0; k <= order; k++) {
            coeff = cpl_table_get_double(polytraces, clab[k], 2*i, &null);
            if (null) {
                cpl_polynomial_delete(polytop);
                missing_top = 1;
                break;
            }
            cpl_polynomial_set_coeff(polytop, &k, coeff);
        }

        missing_bot = 0;
        polybot = cpl_polynomial_new(1);
        for (k = 0; k <= order; k++) {
            coeff = cpl_table_get_double(polytraces, clab[k], 2*i+1, &null);
            if (null) {
                cpl_polynomial_delete(polybot);
                missing_bot = 1;
                break;
            }
            cpl_polynomial_set_coeff(polybot, &k, coeff);
        }

        if (missing_top && missing_bot) {
            cpl_msg_debug(func, "Slit %d was not traced: no extraction!",
                          slit_id[i]);
            continue;
        }

        /*
         * In case just one of the two edges was not traced, the other
         * edge curvature model is duplicated and shifted to the other
         * end of the slit: better than nothing!
         */

        if (missing_top) {
            cpl_msg_debug(func, "Upper edge of slit %d was not traced: "
                          "the spectral curvature of the lower edge "
                          "is used instead.", slit_id[i]);
            polytop = cpl_polynomial_duplicate(polybot);
            ytop = cpl_table_get_double(slits, "ytop", i, NULL);
            ybot = cpl_table_get_double(slits, "ybottom", i, NULL);
            k = 0;
            coeff = cpl_polynomial_get_coeff(polybot, &k);
            coeff += ytop - ybot;
            cpl_polynomial_set_coeff(polytop, &k, coeff);
        }

        if (missing_bot) {
            cpl_msg_debug(func, "Lower edge of slit %d was not traced: "
                          "the spectral curvature of the upper edge "
                          "is used instead.", slit_id[i]);
            polybot = cpl_polynomial_duplicate(polytop);
            ytop = cpl_table_get_double(slits, "ytop", i, NULL);
            ybot = cpl_table_get_double(slits, "ybottom", i, NULL);
            k = 0;
            coeff = cpl_polynomial_get_coeff(polytop, &k);
            coeff -= ytop - ybot;
            cpl_polynomial_set_coeff(polybot, &k, coeff);
        }


        /*
         * Now map smoothed image to CCD.
         * Note that the npseudo value related to this slit is equal
         * to the number of spatial pseudo-pixels decreased by 1
         * (compare with function mos_spatial_calibration()).
         */

        nx = cpl_image_get_size_x(flat);
        ny = cpl_image_get_size_y(flat);

        sdata = cpl_image_get_data_float(spatial);
        xdata = cpl_image_get_data_float(exslit);
        npseudo = cpl_image_get_size_y(exslit) - 1;

        /*
         * Write interpolated smoothed values to CCD image
         */

        for (j = start_pixel; j < end_pixel; j++) {
            top = cpl_polynomial_eval_1d(polytop, j, NULL);
            bot = cpl_polynomial_eval_1d(polybot, j, NULL);
            for (k = 0; k <= npseudo; k++) {
                ypos = top - k*(top-bot)/npseudo;
                yint = (int)(ypos);

                /*
                 * The line:
                 *     value = sdata[j + nx*yint];
                 * should be equivalent to:
                 *     value = npseudo*(top-yint)/(top-bot);
                 */

                if (yint < 0 || yint >= ny-1) {
                    yprev = yint;
                    continue;
                }

                value = sdata[j + nx*yint];   /* Spatial coordinate on CCD */
                ivalue = (int)(value);        /* Nearest spatial pixels:   */
                fvalue = value - ivalue;      /* ivalue and ivalue+1       */
                if (ivalue < npseudo && ivalue >= 0) {
                    vtop = xdata[j + nx*(npseudo-ivalue)];
                    vbot = xdata[j + nx*(npseudo-ivalue-1)];
                    wdata[j + nx*yint] = vtop*(1-fvalue) + vbot*fvalue;

                    if (k) {

                        /*
                         * This is added to recover lost pixels on
                         * the CCD image (pixels are lost because
                         * the CCD pixels are less than npseudo+1).
                         */

                        if (yprev - yint > 1) {
                            value = sdata[j + nx*(yint+1)];
                            ivalue = (int)(value);
                            fvalue = value - ivalue;
                            if (ivalue < npseudo && ivalue >= 0) {
                                vtop = xdata[j + nx*(npseudo-ivalue)];
                                vbot = xdata[j + nx*(npseudo-ivalue-1)];
                                wdata[j + nx*(yint+1)] = vtop*(1-fvalue) 
                                                       + vbot*fvalue;
                            }
                        }
                    }
                }
                yprev = yint;
            }
        }
        cpl_polynomial_delete(polytop);
        cpl_polynomial_delete(polybot);
        cpl_image_delete(exslit);
    }
    
    //TODO: What happens with inter-slit flux conservation. In principle the overall level
    //difference across slits should be conserved.

    cpl_image_delete(rectified);

    cpl_image_divide(flat, smo_flat);

    return smo_flat;
}


/**
 * @brief 
 *   Normalise a long slit flat field exposure
 *  
 * @param flat        Image containing the original flat field spectra
 * @param spa_smooth_radius     Radius of smoothing box along the spatial direction
 * @param disp_smooth_radius     Radius of smoothing box along the dispersion direction
 * @param spa_fit_polyorder   Order of fitting polynomial along the spatial direction
 * @param disp_fit_polyorder   Order of fitting polynomial along the dispersion direction
 *  
 * @return The smoothed flat field exposure used for normalisation
 *
 * The input @em flat frame should be already bias subtracted, and should
 * be oriented so that the dispersion direction is horizontal with @em blue
 * on the left and @em red on the right. The original @em flat image is 
 * divided IN PLACE by its smoothed counterpart, which is also returned. 
 * For any of the directions, if the polynomial @em polyorder is set to a 
 * negative number the smoothing consists of a median filtering box of specified
 * sizes. Alternatively, if @em polyorder is not negative, the smoothing will
 * consist of a polynomial fitting along the requested direction performed 
 * independently for each column/row of the image.
 */
cpl_image *mos_lssflat_normalise
(cpl_image *flat, int spa_smooth_radius, int disp_smooth_radius, 
 int spa_fit_nknots, int disp_fit_nknots,
 double fit_threshold)
{
    const char     *func = "mos_normalise_longflat";

    cpl_image      *smo_flat;

    if (flat == NULL) {
        cpl_error_set(func, CPL_ERROR_NULL_INPUT);
        return NULL;
    }
 
    if (spa_fit_nknots < 1 && spa_smooth_radius < 1) {
        cpl_error_set(func, CPL_ERROR_ILLEGAL_INPUT);
        return NULL;
    }
    if (disp_fit_nknots < 1 && disp_smooth_radius < 1) {
        cpl_error_set(func, CPL_ERROR_ILLEGAL_INPUT);
        return NULL;
    }

    smo_flat = cpl_image_duplicate(flat);
    mosca::image smoothed_flat(smo_flat);

    if (spa_fit_nknots < 0)
    {
        mosca::image_smooth_1d_median<float>(smoothed_flat, 
                                             spa_smooth_radius, 
                                             mosca::SPATIAL_AXIS);        
    }
    else
        mosca::image_cubicspline_1d_fit<float>(smoothed_flat, spa_fit_nknots, 
                                               fit_threshold,
                                               mosca::SPATIAL_AXIS);
        //image_smooth_fit_1d_pol_spa(smo_flat, spa_fit_polyorder);

    if (disp_fit_nknots < 0) 
    {
        mosca::image_smooth_1d_median<float>(smoothed_flat, 
                                             disp_smooth_radius, 
                                             mosca::DISPERSION_AXIS);
    }
    else 
        mosca::image_cubicspline_1d_fit<float>(smoothed_flat, disp_fit_nknots, 
                                               fit_threshold,
                                               mosca::DISPERSION_AXIS);
        //image_smooth_fit_1d_pol_disp(smo_flat, disp_fit_polyorder);

    cpl_image_divide(flat, smo_flat);

    return smo_flat;
}

/**
 * Smooth an image fitting a polynomial in spatial direction 
 * @param image       Image to be smoothed
 * @param polyorder   the order of the polynomial to fit
 * 
 * This function will smooth a MOS image in the spatial direction 
 * (here assumed to be the Y axis) fitting each column with a polynomial. 
 */
void image_smooth_fit_1d_pol_spa(cpl_image * image, cpl_size polyorder)
{
    /*
     * Fit with a polynomial the flat field trend column by column.
     */

    cpl_image_turn(image, -1);   /* For faster memory access */

    cpl_size nx  = cpl_image_get_size_x(image);
    cpl_size ny  = cpl_image_get_size_y(image);
    float * data = cpl_image_get_data_float(image);

    cpl_image * profile = cpl_image_collapse_median_create(image, 1, 0, 0);
    float * level = cpl_image_get_data_float(profile);

    for (cpl_size i = 0; i < ny; i++) {

        /*
         * First get the size of the vectors to allocate:
         * eliminate from fit any value more than 20% away
         * from median level in current column.
         */

        cpl_size npoints = 0;
        float * p = data + i*nx;
        for (cpl_size j = 0; j < nx; j++)
            if (std::fabs(p[j]/level[i] - 1) < 0.20)
                npoints++;

        if (npoints > polyorder + 1) {

            /*
             * Fill the vectors for the fitting
             */

            cpl_vector * flux = cpl_vector_new(npoints);
            double * fdata = cpl_vector_get_data(flux);
            cpl_vector * positions = cpl_vector_new(npoints);
            double * pdata = cpl_vector_get_data(positions);

            npoints = 0;
            p = data + i*nx;
            for (cpl_size j = 0; j < nx; j++) {
                if (fabs(p[j]/level[i] - 1) < 0.20) {
                    fdata[npoints] = p[j];
                    pdata[npoints] = j;
                    npoints++;
                }
            }

            cpl_polynomial * trend = cpl_polynomial_fit_1d_create
                    (positions, flux, polyorder, NULL);

            cpl_vector_delete(flux);
            cpl_vector_delete(positions);

            if (trend) {
                p = data + i*nx;
                for (cpl_size j = 0; j < nx; j++)
                    p[j] = cpl_polynomial_eval_1d(trend, j, NULL);
                cpl_polynomial_delete(trend);
            }
            else {
                cpl_msg_warning(cpl_func, 
                                "Invalid flat field flux fit (ignored)");
            }
        }
    }

    cpl_image_delete(profile);
    cpl_image_turn(image, 1);
}

/**
 * Smooth an image fitting a polynomial in dispersion direction 
 * @param image       Image to be smoothed
 * @param polyorder   the order of the polynomial to fit
 * 
 * This function will smooth a MOS image in the dispersion direction 
 * (here assumed to be the X axis) fitting each row with a polynomial.
 * TODO: merge this with image_smooth_fit_1d_pol_spa() and place it in MOSCA.  
 */
void image_smooth_fit_1d_pol_disp(cpl_image * image, cpl_size polyorder)
{
    float          *p;
    float          *data;
    double         *fdata;
    double         *pdata;
    int             npoints;

    cpl_vector     *positions;
    cpl_vector     *flux;
    cpl_polynomial *trend;

    /*
     * Fit with a polynomial the flat field trend row by row.
     */

    cpl_size nx = cpl_image_get_size_x(image);
    cpl_size ny = cpl_image_get_size_y(image);
    data = cpl_image_get_data_float(image);

    for (cpl_size j = 0; j < ny; j++) {

        /*
         * First get the size of the vectors to allocate:
         */

        npoints = 0;
        p = data + j*nx;
        for (cpl_size k = 0; k < nx; k++)
            if (p[k] > 1.0)
                npoints++;

        if (npoints > polyorder + 1) {

            /*
             * Fill the vectors for the fitting
             */

            flux = cpl_vector_new(npoints);
            fdata = cpl_vector_get_data(flux);
            positions = cpl_vector_new(npoints);
            pdata = cpl_vector_get_data(positions);

            npoints = 0;
            p = data + j*nx;
            for (cpl_size k = 0; k < nx; k++) {
                if (p[k] > 1.0) {
                    fdata[npoints] = p[k];
                    pdata[npoints] = k;
                    npoints++;
                }
            }

            trend = cpl_polynomial_fit_1d_create(positions, flux, 
                                                 polyorder, NULL);

            cpl_vector_delete(flux);
            cpl_vector_delete(positions);

            if (trend) {
                p = data + j*nx;
                for (cpl_size k = 0; k < nx; k++)
                    if (p[k] > 1.0)
                        p[k] = cpl_polynomial_eval_1d(trend, k, NULL);
                cpl_polynomial_delete(trend);
            }
            else {
                cpl_msg_warning(cpl_func, "Invalid flat field flux fit "
                                "(ignored)");
            }
        }
    }
}
