/*----------------------------------------------------------------------------
   
   File name    :   pstest.c
   Author       :	Y. Jung
   Created on   :	June 2001
   Description  :	CONICA point source recipe test.

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

/*
   $Id: pstest.c,v 1.8 2001/11/23 15:02:33 yjung Exp $
   $Author: yjung $
   $Date: 2001/11/23 15:02:33 $
   $Revision: 1.8 $
 */

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

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "eclipse.h"
#include "conicap_lib.h"

/*----------------------------------------------------------------------------
                                Defines
 ---------------------------------------------------------------------------*/

#define CO_TYPE_POL0	1001
#define CO_TYPE_POL45	1002
#define CO_TYPE_POL90	1003
#define CO_TYPE_POL135	1004
#define CO_TYPE_DARK	1005
#define	CO_TYPE_UNKNOWN	1006

/*----------------------------------------------------------------------------
                            Private functions
 ---------------------------------------------------------------------------*/

static int conica_pstest_engine(char *, char *, double) ;
static framelist * conica_pstest_classify(char *) ;
static image_t * conica_pstest_dark(framelist *) ;
static cube_t * conica_pstest_load(framelist *, int, image_t *, image_t *) ;
static double3 * conica_pstest_findoffsets(cube_t *) ;
static int conica_pstest_select(cube_t *, double3 *, double) ;		
static image_t * conica_pstest_stack(cube_t *, intimage	*) ;

/*----------------------------------------------------------------------------
                            Main code
 ---------------------------------------------------------------------------*/

int conica_pstest_main(void * dict)	
{
	dictionary  *   d ;

	char            argname[10] ;
	char        *   name_i ;
	char        *   name_o ;
	double			select ;
	int				nfiles ;

	int				errors ;
	int				i ;
	
	d = (dictionary*)dict ;
	/* Get options */
	select =    dictionary_getdouble(d, "arg.select", 1.0) ;

	/* Get input/output file names */
	nfiles = dictionary_getint(d, "arg.n", -1) ;
	if (nfiles<0) {
		e_error("missing input file name(s): aborting");
		return -1 ;
	}

	/* Loop on input file names */
    errors = 0 ;
    for (i=1 ; i<nfiles ; i++) {
        sprintf(argname, "arg.%d", i);
        name_i = dictionary_get(d, argname) ;
        name_o = dictionary_get(d, "arg.output") ;
        if (name_o == NULL) name_o = strdup(get_rootname(get_basename(name_i)));
		
        else name_o = strdup(get_rootname(name_o)) ;

        /*
         * Once command-line options have been cleared out, call the main
         * computing function.
         */
        errors += conica_pstest_engine(name_i, name_o, select) ;
        free(name_o) ;
    }
    return errors ;
}

static int conica_pstest_engine(
		char 	* 	name_i, 
		char 	* 	name_o, 
		double		select)
{
	framelist	*	tot_list ;
	image_t		*	dark ;
	image_t		*	flat ;
	
	cube_t		*	cube_pol ;
	image_t		*	stacked ;
	double3		*	offsets ;
	double3		*	purged_offsets ;
	intimage	*	contrib ;
	int			*	valid ;
	int				is_valid ;
	int				pol[4] ;
	
	char			outname[FILENAMESZ] ;
	int				i, j, k ;
	
	/* Initialize */
	flat = NULL ;
	pol[0] = CO_TYPE_POL0 ;
	pol[1] = CO_TYPE_POL45 ;
	pol[2] = CO_TYPE_POL90 ;
	pol[3] = CO_TYPE_POL135 ;
	
	/* Classify the input frames : create a framelist with types */
	/* CO_TYPE_POL0, CO_TYPE_POL45, CO_TYPE_POL90, CO_TYPE_POL135,  */
	/* CO_TYPE_DARK or CO_TYPE_UNKNOWN */
	if ((tot_list = conica_pstest_classify(name_i)) == NULL) {
		e_error("cannot classify input frames") ;
		return -1 ;
	}
	/* Compute the MASTER DARK */
	if ((dark = conica_pstest_dark(tot_list)) == NULL) {
		e_warning("cannot compute dark") ;
	}
	/* For each polarization */
	for (i=0 ; i<4 ; i++) {
		/* Load, dark, flat, for current polarization */
		if ((cube_pol = conica_pstest_load(tot_list, pol[i], dark, flat)) == NULL) {
			e_error("cannot load frames for polarization %d\n", i+1) ;
			continue ;
		}
		/* Find objects and offsets	 */
		if ((offsets=conica_pstest_findoffsets(cube_pol)) == NULL) {
			e_error("cannot find offsets") ;
			cube_del(cube_pol) ;
			continue ;
		}
		if (conica_pstest_select(cube_pol, offsets, select) == -1) {
			e_error("problem in selection") ;
			cube_del(cube_pol) ;
			double3_del(offsets) ;
			continue ;
		}
		/* Purge planes and offsets with offsets->z[j] == -1	 */
		valid = calloc(offsets->n, sizeof(int)) ;
		is_valid = 0 ;
		for (j=0 ; j<offsets->n ; j++) {
			if (fabs(offsets->z[j] + 1) > 1e-3) {
				valid[j] = 1 ;
				is_valid = 1 ;
			}
		}
		if (is_valid == 0) {
			e_error("no valid plane found") ;
			cube_del(cube_pol) ;
			double3_del(offsets) ;
			free(valid) ;
			continue ;
		}
		/* Purge planes */
		if ((cube_reject_planes(&cube_pol, valid)) == -1) {
			e_error("cannot purge the cube") ;
			cube_del(cube_pol) ;
			double3_del(offsets) ;
			free(valid) ;
			continue ;
		}
		/* Purge offsets */
		k = 0 ;
		purged_offsets = double3_new(cube_pol->np) ;
		for (j=0 ; j<offsets->n ; j++) {
			if (valid[j] == 1) {
				purged_offsets->x[k] = offsets->x[j] ;
				purged_offsets->y[k] = offsets->y[j] ;
				purged_offsets->z[k] = offsets->z[j] ;
				k++ ;
			}
		}
		double3_del(offsets) ;
		free(valid) ;
		/* Shift planes */
		if ((contrib = cube_shift_expand(&(cube_pol), purged_offsets)) == NULL) {
			e_error("cannot shift images") ;
			cube_del(cube_pol) ;
			double3_del(purged_offsets) ;
			continue ;
		}
		double3_del(purged_offsets) ;
		/* Stack planes */
		if ((stacked = conica_pstest_stack(cube_pol, contrib)) == NULL) {
			e_error("cannot stack images") ;
			cube_del(cube_pol) ;
			intimage_del(contrib) ;
			continue ;
		}
		intimage_del(contrib) ;
		cube_del(cube_pol) ;
		/* Output stacked image */
		sprintf(outname, "%s_%d.fits", name_o, i+1) ;
		image_save_fits(stacked, outname, BPP_DEFAULT) ;	
		image_del(stacked) ;
	}
	
	/* Free and return  */
	framelist_del(tot_list) ;
	if (dark != NULL) image_del(dark) ;
	if (flat != NULL) image_del(flat) ;
	
	return 0 ;
}

static framelist * conica_pstest_classify(char * ascii_list)
{
	framelist	*	flist ;
	char		*	sval ;
	int				i ;
	
	if ((flist = framelist_load(ascii_list)) == NULL) {
		e_error("cannot load input frames") ;
		return NULL ;
	}

	for (i=0 ; i<flist->n ; i++) {
		sval = conica_get_opti4_id(flist->name[i]) ;
		if (sval == NULL) flist->label[i] = CO_TYPE_UNKNOWN ;
		else if (!strcmp(sval, "Pol_00")) flist->label[i] = CO_TYPE_POL0 ;
		else if (!strcmp(sval, "Pol_45")) flist->label[i] = CO_TYPE_POL45 ;
		else if (!strcmp(sval, "Pol_90")) flist->label[i] = CO_TYPE_POL90 ;
		else if (!strcmp(sval, "Pol_135")) flist->label[i] = CO_TYPE_POL135 ;
		else if (!strcmp(sval, "closed")) flist->label[i] = CO_TYPE_DARK ;
		else flist->label[i] = CO_TYPE_UNKNOWN ;
	}
	return flist ;
}

static image_t * conica_pstest_dark(framelist * flist)
{
	framelist	*	d_list ;
	cube_t		*	d_frames ;
	image_t		*	dark ;
	
	if ((d_list = framelist_select(flist, CO_TYPE_DARK)) == NULL) {
		e_error("cannot select dark frames") ;
		return NULL ;
	}

	/* Load cube */
	if ((d_frames = cube_load_strings(d_list->name, d_list->n)) ==NULL) {
		e_error("loading dark frames") ;
		framelist_del(d_list) ;
		return NULL ;
	}
	framelist_del(d_list) ;

	/* Average cube */
	if (d_frames->np>1) {
		dark = cube_avg_linear(d_frames);
	} else {
		e_warning("only 1 frame used for dark") ;
		dark = image_copy(d_frames->plane[0]) ;
	}
	cube_del(d_frames);

	return dark ;
}

static cube_t * conica_pstest_load(
		framelist	*	flist,
		int				type,
		image_t		*	dark,
		image_t		*	flat)
{
	framelist	*	sublist ;
	cube_t		*	outcube ;
	
	/* Select frames */
	if ((sublist = framelist_select(flist, type)) == NULL) {
		e_error("cannot select frames") ;
		return NULL ;
	}

	/* Load cube */
	if ((outcube = cube_load_strings(sublist->name, sublist->n)) ==NULL) {
        e_error("loading frames") ;
        framelist_del(sublist) ;
        return NULL ;
    }
    framelist_del(sublist) ;

	/* Correct dark */
	if (dark != NULL) {
		if (cube_sub_im(outcube, dark) == -1) {
			e_warning("cannot subtract dark") ;
		}
	}
	
	/* Correct flat */
	if (flat != NULL) {
		if (cube_div_im(outcube, flat) == -1) {
			 e_warning("cannot divide by flat") ;
		}
	}

	return outcube ;
}

static double3 * conica_pstest_findoffsets(
		cube_t 	* 	incube)
{
	detected    *   det ;
	
	double3		*	positions ;
	double3		*	offsets ;
	int				obj_npix ;
	int				star_id ;
	int				i, j ;
	
	
	/* Test input cube */
	if (incube == NULL) return NULL ;
	positions = double3_new(incube->np) ;
	
	/* Detect stars */
	for (i=0 ; i<incube->np ; i++) {
    	if ((det = detected_ks_withstats(incube->plane[i], 
						DETECTED_KAPPA)) == NULL) {
    	    e_warning("cannot find any object") ;
    		positions->x[i] = -1.0 ;
    		positions->y[i] = -1.0 ;
    		positions->z[i] = -1.0 ;
		} else {
			/* Keep the largest object 	 */
			star_id = 0 ;
			obj_npix = det->obj_nbpix[0] ;
			for (j=1 ; j<det->nbobj ; j++) {
				if (det->obj_nbpix[j] > obj_npix) {
					obj_npix = det->obj_nbpix[j] ;
					star_id = j ;
				}
			}
			positions->x[i] = det->fine_x[star_id] + 1 ;
			positions->y[i] = det->fine_y[star_id] + 1 ;
			positions->z[i] = 0.0 ;
			detected_del(det) ;
		}
	}

	/* Compute the offsets */
	offsets = double3_new(incube->np) ;
	for (i=0 ; i<incube->np ; i++) {
		offsets->x[i] = positions->x[i] - positions->x[0] ;
		offsets->y[i] = positions->y[i] - positions->y[0] ;
		offsets->z[i] = positions->z[i] ;
	}
	double3_del(positions) ;
	return offsets ;
}

static int conica_pstest_select(
		cube_t	*	incube,
		double3	*	offsets,
		double		select)
{
	int				nb_good_frames ;
	pixelvalue	*	maxima ;
	int			*	classified_frames ;
	pixelvalue		current_max ;
	int				current_id ;
	int				i, j ;

	/* No selection */
	if (fabs(select-1.0) < 1e-3) return 0 ;	
	
	/* How many frames are kept */
	nb_good_frames = (int)(incube->np * select) ;

	/* Compute maxima */
	maxima = malloc(incube->np * sizeof(pixelvalue)) ;
	for (i=0 ; i<incube->np ; i++) {
		maxima[i] = image_getmax(incube->plane[i]) ;	
	}

	/* Classify frames */
	classified_frames = malloc(incube->np * sizeof(int)) ;
	for (i=0 ; i<incube->np ; i++) {
		current_max = -1 ;
		current_id  = 0 ;
		for (j=0 ; j<incube->np ; j++) {
			if (current_max < maxima[j]) {
				current_max = maxima[j] ;
				current_id = j ;
			}
		}
		maxima[current_id] = -1 ;
		classified_frames[i] = current_id ;
	}
	free(maxima) ;

	/* Switch off non selected frames */
	for (i=nb_good_frames ; i<incube->np ; i++) {
		offsets->z[classified_frames[i]] = -1 ;
	}
	free(classified_frames) ;
	
	return 0 ;
}

static image_t * conica_pstest_stack(
		cube_t 		*	incube,
		intimage	*	contributions)
{
	image_t		*	averaged ;

	if ((averaged = cube_avg_linear(incube)) == NULL) {
		e_error("cannot average cube") ;
		return NULL ;
	}
	image_cst_op_local(averaged, incube->np, '*') ;
	image_div_intimage_local(averaged, contributions) ;
	return averaged ;
}