/*
 * This file is part of the KMOS 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
 */

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

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

#include <math.h>
#include <cpl.h>

/*----------------------------------------------------------------------------*/
/**
  @defgroup kmos_oscan          Overscan Functions
 */
/*----------------------------------------------------------------------------*/
/** @{ */

static double kmos_oscan_biweight_mid_var(const cpl_vector *) ;
static double kmos_oscan_mid_var_tukey(const cpl_vector *, double) ;

/*----------------------------------------------------------------------------*/
/**
  @brief    Overscan Correction
  @param    in      The image to correct
  @return   The corrected image or NULL in error case.
  @see      kmos_oscan_mid_var_tukey()

  The correction applies in 2 steps: 
  - Column correction
    The 4 leftmost and rightmost columns are averaged to 1 column
    The column median* value is subtracted to the column. The result is
    the correction column, which is subtracted to all columns in the
    image.

  - Row correction
    The 4 top rows of the image are extracted (2048x4 ext image).
    This image is collapsed in the x direction (1x4 ext_coll image)
    The collapsed column median is subtracted to itself.
    It is then subtracted to all columns of the extracted image.

    The extracted image is then cut in 32 channels along x (64 pixels
    each). For each channel, the even (resp. odd) columns are extracted
    and their median* is computed. The result is then subtracted to all
    even (resp. odd) columns of the image. 

    The median* is computed by kmos_oscan_mid_var_tukey()
 */
/*----------------------------------------------------------------------------*/
cpl_image * kmos_oscan_correct(const cpl_image * in) 
{
    cpl_vector      *   col ;
    cpl_vector      *   sub_col ;
    double          *   pcol ;
    const float     *   pin ;
    double              sum, msig, mval, med, ne_med, no_med ;
    int                 nx, ny, nb_good ;
    cpl_image       *   in_colcorr ;
    float           *   pin_colcorr ;
    cpl_image       *   in_colrowcorr ;
    float           *   pin_colrowcorr ;
    cpl_image       *   extracted ;
    float           *   pextracted ;
    cpl_image       *   row_collapsed ;
    float           *   prow_collapsed ;
    cpl_vector      *   odd_vec ;
    double          *   podd ;
    cpl_vector      *   eve_vec ;
    double          *   peve ;
    cpl_vector      *   oddeven_vec ;
    double          *   poddeven ;
    int                 colcorr_skip_right, colcorr_skip_left,
                        colcorr_skip_top, colcorr_skip_bottom,
                        corr_nb_channels, corr_channel_size ;
    int                 i, j, k ;

    /* Check Inputs */
    if (in == NULL) return NULL ;
   
    /* Data access */
    pin = cpl_image_get_data_float_const(in) ;
    nx = cpl_image_get_size_x(in) ;
    ny = cpl_image_get_size_y(in) ;

    /* Check entries */
    if (nx < 10 || ny < 10) return NULL ;

    /* Initialise */
    colcorr_skip_right = 4 ;
    colcorr_skip_left = 4 ;
    colcorr_skip_top = 4 ;
    colcorr_skip_bottom = 4 ;
    corr_nb_channels = 32 ;
    corr_channel_size = (int)(nx/corr_nb_channels); 

    /* *** COLUMN CORRECTION *** */
    /* Compute column col by averaging first and last columns*/
    col = cpl_vector_new(ny) ;
    pcol = cpl_vector_get_data(col) ;

    /* Loop on the rows ie col vector elements */
    for (j=0 ; j<ny ; j++) {
        /* Compute the mean of the first and last pixels of the row */
        sum = 0.0 ;
        nb_good = 0 ;
        for (i=0 ; i<colcorr_skip_left ; i++) {
            if (!isnan(pin[i+j*nx]) && !isinf(pin[i+j*nx])) {
                nb_good ++ ;
                sum += pin[i+j*nx] ;
            }
        }
        for (i=nx-colcorr_skip_right ; i<nx ; i++) {
            if (!isnan(pin[i+j*nx]) && !isinf(pin[i+j*nx])) {
                nb_good ++ ;
                sum += pin[i+j*nx] ;
            }
        }
        /* Store it */
        if (nb_good != 0)   pcol[j] = sum/nb_good ;
        else                pcol[j] = 0.0 ;
    }
    /* cpl_plot_vector("","w lines","",col) ; */

    /* Compute msig / mval from the correction column */
    sub_col = cpl_vector_extract(col, colcorr_skip_bottom, 
            ny-colcorr_skip_top-1, 1) ;
    msig = kmos_oscan_biweight_mid_var(sub_col) ;
    mval = kmos_oscan_mid_var_tukey(sub_col, msig) ;
    cpl_vector_delete(sub_col) ;
    if (isnan(msig) || isnan(mval)) {
        cpl_vector_delete(col) ;
        return NULL ;
    }
    cpl_msg_debug(__func__, "COLCORR: msig / mval: %g / %g", msig, mval) ;
    
    /* Adjust Correction column col */
    cpl_vector_subtract_scalar(col, mval); 
    pcol = cpl_vector_get_data(col) ;
    /* cpl_plot_vector("","w lines","",col) ; */

    /* Use col to correct the image */
    in_colcorr = cpl_image_duplicate(in) ;
    pin_colcorr = cpl_image_get_data_float(in_colcorr) ;
    for (i=0 ; i<nx ; i++) {
        for (j=0 ; j<ny ; j++) {
            if (i<colcorr_skip_left || i>nx-colcorr_skip_right || 
                    j<colcorr_skip_bottom || j>ny-colcorr_skip_top) {
                pin_colcorr[i+j*nx] = pin[i+j*nx] - mval ;
            } else {
                pin_colcorr[i+j*nx] = pin[i+j*nx] + pcol[j] ;
            }
        }
    }
    cpl_vector_delete(col) ;

    /* *** ROW CORRECTION *** */
    /* Extract the top rows of the image */
    extracted = cpl_image_extract(in_colcorr, 1, ny-colcorr_skip_top+1, nx, ny);
    pextracted = cpl_image_get_data_float(extracted) ;

    /* Median collapse them along x */
    row_collapsed = cpl_image_collapse_median_create(extracted, 1, 0, 0) ;
    prow_collapsed = cpl_image_get_data_float(row_collapsed) ;

    /* Median of the colcorr_skip_top values */
    med = cpl_image_get_median(row_collapsed) ;

    /* extracted = extracted - row_collapsed + med */
    for (i=0 ; i<nx ; i++) {
        for (j=0 ; j<colcorr_skip_top ; j++) {
            pextracted[i+j*nx] -= prow_collapsed[j] - med ;
        }
    }
    cpl_image_delete(row_collapsed) ;

    /* Create Even and Odd vectors */
    eve_vec = cpl_vector_new((int)(colcorr_skip_top*corr_channel_size/2)) ;
    peve = cpl_vector_get_data(eve_vec) ;
    odd_vec = cpl_vector_new((int)(colcorr_skip_top*corr_channel_size/2)) ;
    podd = cpl_vector_get_data(odd_vec) ;
    oddeven_vec = cpl_vector_new(nx) ;
    poddeven = cpl_vector_get_data(oddeven_vec) ;
    /* Loop on 32 channels of 64x4 pixels along the 2048x4 extracted ima */
    for (k=0 ; k<corr_nb_channels ; k++) {
        /* From the 64 4-pixels-columns, put 32 in eve, and 32 in odd */
        for (i=0 ; i<(int)(corr_channel_size/2) ; i++) {
            for (j=0 ; j<colcorr_skip_top ; j++) {
                peve[colcorr_skip_top*i+j] =
                    pextracted[k*corr_channel_size+2*i+nx*j] ;
                podd[colcorr_skip_top*i+j] =
                    pextracted[k*corr_channel_size+2*i+1+nx*j] ;
            }
        }

        /* Compute ne_med, no_med  */
        msig = kmos_oscan_biweight_mid_var(eve_vec) ;
        ne_med = kmos_oscan_mid_var_tukey(eve_vec, msig) ;
        msig = kmos_oscan_biweight_mid_var(odd_vec) ;
        no_med = kmos_oscan_mid_var_tukey(odd_vec, msig) ;

        if (isnan(msig) || isnan(ne_med) || isnan(no_med)) {
            cpl_image_delete(extracted) ;
            cpl_vector_delete(eve_vec) ;
            cpl_vector_delete(odd_vec) ;
            cpl_vector_delete(oddeven_vec) ;
            cpl_image_delete(in_colcorr) ;
            return NULL ;
        }
        cpl_msg_debug(__func__, "ROWCORR: msig / ne_med / no_med: %g / %g / %g",
                msig, ne_med, no_med) ;
        /* Fill oddeven_vec */
        for (i=0 ; i<(int)(corr_channel_size/2) ; i++) {
            poddeven[k*corr_channel_size+2*i] = ne_med ;
            poddeven[k*corr_channel_size+2*i+1] = no_med ;
        }
    }
    cpl_image_delete(extracted) ;
    cpl_vector_delete(eve_vec) ;
    cpl_vector_delete(odd_vec) ;
    /* cpl_plot_vector("","w lines","",oddeven_vec) ; */

    /* Correct in_colcorr */
    in_colrowcorr = cpl_image_duplicate(in_colcorr) ;
    pin_colrowcorr = cpl_image_get_data_float(in_colrowcorr) ;
    for (i=0 ; i<nx ; i++) {
        for (j=0 ; j<ny ; j++) {
            pin_colrowcorr[i+j*nx] = pin_colcorr[i+j*nx] - poddeven[i] ;
        }
    }
    cpl_vector_delete(oddeven_vec) ;
    cpl_image_delete(in_colcorr) ;
    
    return in_colrowcorr ;
}

/** @} */

/*----------------------------------------------------------------------------*/
/**
  @brief    Compute the biWeightMidVar
  @param    in      The input vector
  @return   The computed value

  The median mu of the input vector is computed.
  MAD = median(fabs(in-mu))
  udat = (in-mu)/(9*MAD)
out=sqrt((size(in)*SUM(in-mu)^2*((1-udat^2)^4)/(SUM(1-udat^2)*(1-5*udat^2)^2))))
  for all a where fabs(a) < 1
  Return out
 */
/*----------------------------------------------------------------------------*/
static double kmos_oscan_biweight_mid_var(const cpl_vector * in) 
{
    const double    *   pin ;
    cpl_vector      *   mad_vec ;
    cpl_vector      *   udat_vec ;
    double          *   pmad_vec ;
    double          *   pudat_vec ;
    double              out, med, mad, nomi, denomi ;
    int                 i, nelem ;

    /* Check entries */
    if (in == NULL) return -1.0 ;

    /* Access input */
    pin = cpl_vector_get_data_const(in) ;
    nelem = cpl_vector_get_size(in) ;

    /* cpl_plot_vector("","w lines","",in) ; */
    /* Compute MAD */
    med = cpl_vector_get_median_const(in) ;
    mad_vec = cpl_vector_new(nelem) ;
    pmad_vec = cpl_vector_get_data(mad_vec) ;
    for (i=0 ; i<nelem ; i++) pmad_vec[i] = fabs(pin[i] - med) ;
    mad = cpl_vector_get_median_const(mad_vec) ;
    cpl_vector_delete(mad_vec);
    /* printf("MAD: %g \n", mad) ; */

    /* Compute udat */
    udat_vec = cpl_vector_new(nelem) ;
    pudat_vec = cpl_vector_get_data(udat_vec) ;
    for (i=0 ; i<nelem ; i++) {
        pudat_vec[i] = (pin[i]-med) / (9*mad) ;
        /* Skip values outside [-1.0, 1.0]*/
        if (fabs(pudat_vec[i]) > 1.0) pudat_vec[i] = 2e6 ;
    }
    /* cpl_plot_vector("","w lines","",udat_vec) ; */

    /* Compute out */
    out = nomi = denomi = 0.0 ;
    for (i=0 ; i<nelem ; i++) {
        if (pudat_vec[i] <1e6) {
            nomi += pow(pin[i]-med, 2) * pow(1-pow(pudat_vec[i], 2), 4);
            denomi += (1-pow(pudat_vec[i], 2)) * (1-5*pow(pudat_vec[i], 2)) ;
        }
    }
    cpl_vector_delete(udat_vec) ;
    /* printf("nelem, nomi, denomi: %d %g %g\n", nelem, nomi, denomi) ; */
    out = sqrt(nelem*nomi/pow(denomi, 2)) ;
    return out ;
}

/*----------------------------------------------------------------------------*/
/**
  @brief    Compute the tukeyBiWeight
  @param    in      The input vector
  @param    scale   
  @return   The computed value

  The median mu of the input vector is computed.
  We iterate not more than 50 times:
    in2 = (in-mu)/scale
    in3 is elements of in2 that are <= 4.826
    W = (1-(in3/4.826)^2)^2
    new_mu = (sum(W.in)) / (sum(W))

    if fabs(mu-new_mu) < scale.1e-6
        stop
    else 
        mu = new_mu
  Return mu
 */
/*----------------------------------------------------------------------------*/
static double kmos_oscan_mid_var_tukey(
        const cpl_vector    *   in,
        double                  scale) 
{
    cpl_vector      *   weights_vec ;
    const double    *   pin ;
    double          *   pwei ;
    double              tol, c, mu, nmu, sum1, sum2 ;
    int                 maxiter ;
    int                 i, j, nelem, found ;

    /* Check Inputs */
    if (in == NULL) return -1.0 ;

    /* Initialise */
    tol = 1e-6 ;
    maxiter = 50 ;
    c = 4.826 ;
    found = 0 ;
    pin = cpl_vector_get_data_const(in) ;
    nelem = cpl_vector_get_size(in) ;
    
    /* Compute Median */
    mu = cpl_vector_get_median_const(in) ;

    /* Create weights_vec */
    weights_vec = cpl_vector_duplicate(in) ;
    pwei = cpl_vector_get_data(weights_vec) ;

    /* Loop maxiter times */
    for (i=0 ; i<maxiter ; i++) {
        /* Compute weights_vec */
        for (j=0 ; j<nelem ; j++) {
            if (fabs(pin[j])>c) pwei[j] = 2e6 ;
            else {
                found = 1 ;
                pwei[j] = pow(1-pow((pin[j]-mu)/(c*scale),2),2);
            }
        }
        if (found == 0) {
            /* Special case of failure- give the median */
            cpl_msg_warning(__func__, "All values above 4.826 - return median");
            cpl_vector_delete(weights_vec) ; 
            return cpl_vector_get_median_const(in) ;
        }
        /* nmu */
        sum1 = sum2 = 0 ;
        for (j=0 ; j<nelem ; j++) {
            if (pwei[j] < 1e6) {
                sum1 += pwei[j] * pin[j] ;
                sum2 += pwei[j] ;
            }
        }
        nmu = sum1 / sum2 ;

        /* Convergence Condition  */
        if (fabs(nmu-mu) < scale*tol)   break ;
        else                            mu = nmu ; 
    }
    cpl_vector_delete(weights_vec) ;
    return nmu ;
}

