/* @(#)spatial_profile.c	17.1.1.1 (ESO-IPG) 01/25/02 17:52:00 */
/* 

   Otmar Stahl, Thomas Szeifert, Anton Malina
   
   spatial_profile.c        
   
   fit spatial profile for optimum extraction a la 
   Horne but with some minor modifications

   Polynom fit routine was changed - we fit only the pixels which are
   not masked in the earlier iterations and take the standard deviation
   down to the fitting routines. Secondly we check now if the masked
   pixels are still bad or only rejected by accident. In the very first
   iteration no fit is done, but a median filter to mask the worst
   pixels. Otherwise a lot of pixels will be rejected.

 */

/* system includes */

#include <math.h>

/* FEROS specific includes */

#include <proto_nrutil.h>
#include <proto_mutil.h>
#include <mutil.h>
#include <dlfit.h>

#define mymax(A,B) ((A)  > (B) ? (A) : (B))

static void select_fit_poly 
(
#ifdef __STDC__
 float *inimage, float *ximage, float *sigin, int npixi, 
 float *mask, float *outimage, float *ximageo, int npixo, 
 int fit_deg, double ad[]
#endif
 );

int spatial_profile
#ifdef __STDC__
(
 float frame[], float sky[], float outframe[], 
 float varframe[], float mask[], float profile[], 
 float variance[], float x[], float y[],
 float z[], int slitlen, int ntotal, int npix, 
 int niter, int fit_deg, 
 float v0, float gain, float thres, double p_coeff[]
 )
#else
     (
      frame, sky, outframe, varframe, mask, profile, variance, 
      x, y, z, slitlen, ntotal, npix, niter, fit_deg, v0, gain, 
      thres, p_coeff
      )
     float frame[], sky[], outframe[], varframe[];
     float mask[], profile[], variance[], x[], y[], z[], v0, gain, thres;
     int   slitlen, ntotal, npix, niter, fit_deg;
     double p_coeff[];
#endif

/* get spatial profile */

{

  int i, j, iter, index, kk;
  float ratio, diff, diff2, rmax, sum1, sum2, sum3, mp,*sigin;
  double *ad;



  if (niter > 0){
    sigin = (float *) fvector(0,npix-1);
    ad = (double *) dvector (1, fit_deg);
  }

  for (i = 0; i < ntotal; i++)
    {
      varframe[i] = 0.0;
      outframe[i] = 0.0;
    }

  for (j = 0; j < slitlen; j++)
     {
      for (i = 0; i < ntotal; i++)
	{
	  index = j * ntotal + i;
	  mask[index] = 1.0;
	  profile[index] = 0.0;
	  variance[index] = 0.0;
	  outframe[i] += (frame[index] - sky[index]);
	}
    }
  
  /*                construct spatial profile               */

  if (niter<=0) {

    for (i = 0; i < ntotal; i++)
      varframe[i] = (float) sqrt( (double) (v0 + outframe[i]/gain) );

    return 0;
  }

  for (kk=1; kk<=fit_deg; kk++)
    ad[kk] = 0.0;

  
  for (iter = 0; iter < niter; iter++)
    {/*iteration*/
      for (i = 0; i < npix; i++)
	outframe[i] = FMAX (outframe[i], 1.0);
      
      for (j = 0; j < slitlen; j++)
	{
	  for (i = 0; i < npix; i++)
	    {
	      index = j * ntotal + i;
	      y[i] = (frame[index] - sky[index]) / outframe[i];
	      sigin[i] = (float) sqrt(mymax((v0 + frame[index]/gain),v0)) / 
                outframe[i];
	    }

	  /*   filter spatial profile                              */

	  if (iter == 0)
	    for (i = 4; i < npix - 4; i++)
	      {
		index = j * ntotal + i;
		profile[index] = heap_median (9, &y[i-4]); 
	      }
	  else
	    for (i = 0; i < npix; i++)
	      profile[index] = y[i];
	}

      
      for (j = 0; j < slitlen; j++)
	{
	  
	  /*        fit a polynomial to spatial profile          */
	  
	  index = j * ntotal;

	  /*
	    printf("%d %d %d %f %f\n",j,index,npix,profile[1],profile[npix-1]);
	    */
	  /*
	    fit_poly (&profile[index], x, npix, z, x, npix, fit_deg, ad);
	  */

	  if (iter > 0){
	    select_fit_poly(&profile[index], x, sigin, npix, 
			    &mask[index], z, x, npix, fit_deg, ad);  

	    for (kk=1; kk<=fit_deg; kk++)
	      {
		p_coeff[j*fit_deg+kk] = ad[kk];
	      }
	  
	    for (i = 0; i < npix; i++)
	      {
		index = j * ntotal + i;
		profile[index] = z[i];
	      }
	  }
	}
	  
      /*      normalize spatial profile, enforce positivity         */
	  
      for (i = 0; i < npix; i++)
	z[i] = 0.0;

      for (j = 0; j < slitlen; j++)
	{
	  for (i = 0; i < npix; i++)
	    {
	      index = j * ntotal + i;
		  
	      /*               THIS HAS TO BE INVESTIGATED !!          */
		
	      /*
		profile[index] = FMAX (profile[index], 0.0);      
		*/

	      z[i] += profile[index]; /*cross order sum*/
	    }
	}
	  
      for (j = 0; j < slitlen; j++)
	{
	  for (i = 0; i < npix; i++)
	    {
	      index = j * ntotal + i;
	      if (z[i] != (float) 0.0)
		profile[index] = profile[index] / z[i]; /*normalize profile*/
	    }
	}
	  
      /*             compute  mask and variance                        */
	  
      for (j = 0; j < slitlen; j++)
	{
	  for (i = 0; i < npix; i++)
	    {
	      index = j * ntotal + i;
	      variance[index] = v0 + fabs (outframe[i] *
					   profile[index] + sky[index]) / gain;
		  
	    }
	}
      for (i = 0; i < npix; i++)
	{
	  rmax = 1.0;
	  kk = -1;
	  for (j = 0; j < slitlen; j++) /* filter the maximum of deviation*/
	    {
	      index = j * ntotal + i;
	      diff = (frame[index] - sky[index]) - outframe[i] *
		profile[index];
	      diff2 = diff * diff;
	      ratio = (diff2 / (thres * variance[index]));
              
	      if (mask[index] < (float) 1.0 && ratio < (float) 1.0)
                mask[index] = (float) 1.0;    /* heal the good pixels */

	      if (ratio  * mask[index] > rmax) 
		{
		  rmax = ratio;
		  kk = index;
		}
	    }
	      
	  /*      now extract the weighted spectrum               */
	      
	  if (kk >= 0 && rmax > (float) 1.0)
	    mask[kk] = 0.0;

	  sum1 = 0.0;
	  sum2 = 0.0;
	  sum3 = 0.0;
	  for (j = 0; j < slitlen; j++)
	    {
	      index = j * ntotal + i;
	      mp = mask[index] * profile[index];
	      sum1 += mp * (frame[index] - sky[index]) / variance[index];
	      sum2 += mp * (profile[index] / variance[index]);
	      sum3 += mp;
	    }


	  if (sum2 != 0)
	    outframe[i] = sum1 / sum2;
	  if (sum2 != 0)
	    varframe[i] = (float) sqrt( fabs( (double) sum3/sum2));
	}
    }

  free_dvector(ad, 1, fit_deg);
  free_fvector(sigin,0,npix-1);

  return 0;
}

static void select_fit_poly 
#ifdef __STDC__
(
 float *inimage, float *ximage, float *sigin, int npixi, 
 float *mask, float *outimage, float *ximageo, int npixo, 
 int fit_deg, double ad[]
 )
#else
     ( 
      inimage, ximage, sigin, npixi, mask, outimage, ximageo, 
      npixo, fit_deg, ad
      )
     float *inimage, *ximage, *outimage, *ximageo,*mask,*sigin;
     int npixi, npixo, fit_deg;
     double ad[];
#endif
     
/*      fit a polynomial (Numerical Recipes), modified for echelle!  */

{

  int i,isel=0;
  float *xsel,*ysel,*sigsel;

  xsel = (float *) fvector(0,npixi-1);
  ysel = (float *) fvector(0,npixi-1);
  sigsel = (float *) fvector(0,npixi-1);

  for (i = 0 ; i < npixi ; i++)
    if (mask[i] > 0.5 && inimage[i] != (float) 0.0 && sigin[i] > (float) 0.0){
      xsel[isel] = ximage[i];
      ysel[isel] = inimage[i];
      sigsel[isel] = sigin[i];
      isel++;
    }
  if (isel > fit_deg)
    fit_poly_weight (ysel, xsel, sigsel, isel, outimage, ximageo, npixo, 
                     fit_deg, ad); 
  else
    for (i = 0 ; i < npixi ; i++)
      outimage[i] = inimage[i]; 

  free_fvector(xsel,0,npixi-1);
  free_fvector(ysel,0,npixi-1);
  free_fvector(sigsel,0,npixi-1);

}
/*---------------------------------------------------------------------------*/
void fit_poly_weight
#ifdef __STDC__
 (
  float *inimage, float *ximage, float *sigin, int npixi, 
  float *outimage, float *ximageo, int npixo, 
  int fit_deg, double *ad
  )
#else
     ( 
      inimage, ximage, sigin, npixi, outimage, ximageo, npixo, 
      fit_deg, ad 
      )
     float *inimage, *ximage, *outimage, *ximageo,*sigin;
     int npixi, npixo, fit_deg;
     double *ad;
#endif

/*      fit a polynomial (Numerical Recipes), modified for echelle!  */

{
  /*  float xout;*/
  double xout;
  double *xin, *yin;
  int i;

  /* new - dlfit not lfit*/
  double *sig, chisq, **covar;
  int    *ia;

  covar = (double **) dmatrix (1, fit_deg, 1, fit_deg);
  sig = (double *) dvector (1, npixi);
  ia = (int *) ivector (1, fit_deg);

  xin = (double *) dvector (1, npixi);
  yin = (double *) dvector (1, npixi);



  for (i = 0; i < npixi; i++)
    {
      xin[i+1] = (double) ximage[i];
      yin[i+1] = (double) inimage[i];
      sig[i+1] = (double) sigin[i];;
    }
  for (i = 1 ; i <= fit_deg ; i++)
    ia[i] = 1;


  dlfit (xin, yin, sig, npixi, ad, ia, fit_deg, covar, &chisq, dpoly);

  for (i = 0; i < npixo; i++)
    {
      xout = (double) ximageo[i];
      outimage[i] = (float) eval_dpoly (xout, ad, fit_deg);
    }


  free_dvector(xin, 1, npixi); 
  free_dvector(yin, 1, npixi);
  free_dvector(sig, 1, npixi);
  free_dmatrix(covar, 1, fit_deg, 1, fit_deg);
  free_ivector(ia, 1, fit_deg);
}