/*
 * integration_check.c
 *
 * Check peak integration
 *
 * Copyright © 2012 Thomas White <taw@physics.org>
 * Copyright © 2012 Andrew Martin <andrew.martin@desy.de>
 *
 * This file is part of CrystFEL.
 *
 * CrystFEL 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 3 of the License, or
 * (at your option) any later version.
 *
 * CrystFEL 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 CrystFEL.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

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


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

#include <image.h>
#include <peaks.h>
#include <utils.h>
#include <beam-parameters.h>


/* The third integration check draws a Poisson background and checks that, on
 * average, it gets subtracted by the background subtraction. */
static void third_integration_check(struct image *image, int n_trials,
                                    int *fail)
{
	double mean_intensity = 0.0;
	double mean_bg = 0.0;
	double mean_max = 0.0;
	double mean_sigma = 0.0;
	int i;
	int fs, ss;
	int nfail = 0;

	for ( i=0; i<n_trials; i++ ) {

		double intensity, sigma;
		double fsp, ssp;
		int r;

		for ( fs=0; fs<image->width; fs++ ) {
		for ( ss=0; ss<image->height; ss++ ) {
			image->data[fs+image->width*ss] = poisson_noise(10.0);
		}
		}

		r = integrate_peak(image, 64, 64, &fsp, &ssp,
		                   &intensity, &sigma);

		if ( r == 0 ) {
			mean_intensity += intensity;
			mean_sigma += sigma;
		} else {
			nfail++;
		}

	}
	mean_intensity /= n_trials;
	mean_bg /= n_trials;
	mean_max /= n_trials;
	mean_sigma /= n_trials;

	STATUS("  Third check (mean values): intensity = %.2f, sigma = %.2f,"
	       " integration failed %i/%i times\n",
	       mean_intensity, mean_sigma, nfail, n_trials);

	if ( fabs(mean_intensity) > 5.0 ) {
		ERROR("Mean intensity should be close to zero.\n");
		*fail = 1;
	}
	if ( fabs(mean_intensity) > mean_sigma/10.0 ) {
		ERROR("Mean intensity should be much less than mean sigma.\n");
		*fail = 1;
	}
}


/* The fourth integration check draws a Poisson background and draws a peak on
 * top of it, then checks that the intensity of the peak is correctly recovered
 * accounting for the background. */
static void fourth_integration_check(struct image *image, int n_trials,
                                     int *fail)
{
	double mean_intensity = 0.0;
	double mean_bg = 0.0;
	double mean_max = 0.0;
	double mean_sigma = 0.0;
	int i;
	int fs, ss;
	int pcount = 0;
	int nfail = 0;

	for ( i=0; i<n_trials; i++ ) {

		double intensity, sigma;
		double fsp, ssp;
		int r;

		for ( fs=0; fs<image->width; fs++ ) {
		for ( ss=0; ss<image->height; ss++ ) {
			int idx = fs+image->width*ss;
			image->data[idx] = 0.0;
			if ( (fs-64)*(fs-64) + (ss-64)*(ss-64) > 9*9 ) {
				image->data[idx] = poisson_noise(10.0);
			} else {
				image->data[idx] += 1000.0;
				pcount++;
			}
		}
		}

		r = integrate_peak(image, 64, 64, &fsp, &ssp,
		                   &intensity, &sigma);

		if ( r == 0 ) {
			mean_intensity += intensity;
			mean_sigma += sigma;
		} else {
			nfail++;
		}

	}
	mean_intensity /= n_trials;
	mean_bg /= n_trials;
	mean_max /= n_trials;
	mean_sigma /= n_trials;
	pcount /= n_trials;

	STATUS(" Fourth check (mean values): intensity = %.2f, sigma = %.2f,"
	       " integration failed %i/%i times\n",
	       mean_intensity, mean_sigma, nfail, n_trials);

	if ( fabs(mean_intensity - pcount*1000.0) > 4000.0 ) {
		ERROR("Mean intensity should be close to %f\n", pcount*1000.0);
		*fail = 1;
	}
	if ( fabs(mean_bg-10.0) > 0.3 ) {
		ERROR("Mean background should be close to ten.\n");
		*fail = 1;
	}
	if ( fabs(mean_intensity) < mean_sigma ) {
		ERROR("Mean intensity should be greater than mean sigma.\n");
		*fail = 1;
	}
}


/* The fifth check integrates a Poisson background with background subtraction
 * switched off, and checks that the result is what would be expected. */
static void fifth_integration_check(struct image *image, int n_trials,
                                    int *fail)
{
	double mean_intensity = 0.0;
	double mean_bg = 0.0;
	double mean_max = 0.0;
	double mean_sigma = 0.0;
	int i;
	int fs, ss;
	int pcount = 0;

	for ( i=0; i<n_trials; i++ ) {

		double intensity, bg, max, sigma;
		double fsp, ssp;

		for ( fs=0; fs<image->width; fs++ ) {
		for ( ss=0; ss<image->height; ss++ ) {
			int idx = fs+image->width*ss;
			image->data[idx] = poisson_noise(10.0);
			if ( (fs-64)*(fs-64) + (ss-64)*(ss-64) < 10*10 ) {
				pcount++;
			}
		}
		}
		integrate_peak(image, 64, 64, &fsp, &ssp, &intensity, &sigma);

		mean_intensity += intensity;
		mean_bg += bg;
		mean_max += max;
		mean_sigma += sigma;

	}
	mean_intensity /= n_trials;
	mean_bg /= n_trials;
	mean_max /= n_trials;
	mean_sigma /= n_trials;
	pcount /= n_trials;

	STATUS("  Fifth check (mean values): intensity = %.2f, bg = %.2f,"
	       " max = %.2f, sigma = %.2f\n",
	       mean_intensity, mean_bg, mean_max, mean_sigma);
	double s = pcount*10.0;
	if ( fabs(mean_intensity - s) > 5.0 ) {
		ERROR("Mean intensity should be close to %f.\n", pcount*10.0);
		*fail = 1;
	}
	if ( fabs(mean_bg-10.0) > 0.3 ) {
		ERROR("Mean background should be close to ten.\n");
		*fail = 1;
	}
}


/* The sixth check is like the fourth check, except that the background
 * subtraction is switched off */
static void sixth_integration_check(struct image *image, int n_trials,
                                    int *fail)
{
	double mean_intensity = 0.0;
	double mean_bg = 0.0;
	double mean_max = 0.0;
	double mean_sigma = 0.0;
	int i;
	int fs, ss;
	int pcount = 0;
	int npcount = 0;

	for ( i=0; i<n_trials; i++ ) {

		double intensity, bg, max, sigma;
		double fsp, ssp;

		for ( fs=0; fs<image->width; fs++ ) {
		for ( ss=0; ss<image->height; ss++ ) {
			int idx = fs+image->width*ss;
			double r = (fs-64)*(fs-64) + (ss-64)*(ss-64);
			image->data[idx] = poisson_noise(10.0);
			if ( r < 9*9 ) {
				image->data[idx] += 1000.0;
				pcount++;
			} else if ( r < 10*10 ) {
				npcount++;
			}
		}
		}
		integrate_peak(image, 64, 64, &fsp, &ssp, &intensity, &sigma);

		mean_intensity += intensity;
		mean_bg += bg;
		mean_max += max;
		mean_sigma += sigma;

	}
	mean_intensity /= n_trials;
	mean_bg /= n_trials;
	mean_max /= n_trials;
	mean_sigma /= n_trials;
	pcount /= n_trials;
	npcount /= n_trials;

	STATUS("  Sixth check (mean values): intensity = %.2f, bg = %.2f,"
	       " max = %.2f, sigma = %.2f\n",
	       mean_intensity, mean_bg, mean_max, mean_sigma);

	double s = pcount*1010.0 + npcount*10.0;
	if ( fabs(mean_intensity - s) > 4000.0 ) {
		ERROR("Mean intensity should be close to %f.\n", s);
		*fail = 1;
	}
	if ( fabs(mean_bg-10.0) > 0.3 ) {
		ERROR("Mean background should be close to ten.\n");
		*fail = 1;
	}

	STATUS("            (Absolute value of mean sigma not (yet) tested)\n");
}


int main(int argc, char *argv[])
{
	struct image image;
	double fsp, ssp, intensity, sigma;
	int fs, ss;
	FILE *fh;
	unsigned int seed;
	int fail = 0;
	const int n_trials = 1000;
	int r, npx;

	fh = fopen("/dev/urandom", "r");
	fread(&seed, sizeof(seed), 1, fh);
	fclose(fh);
	srand(seed);

	image.data = malloc(128*128*sizeof(float));
	image.flags = NULL;
	image.beam = NULL;
	image.lambda = ph_eV_to_lambda(1000.0);

	image.det = calloc(1, sizeof(struct detector));
	image.det->n_panels = 1;
	image.det->panels = calloc(1, sizeof(struct panel));

	image.det->panels[0].min_fs = 0;
	image.det->panels[0].max_fs = 128;
	image.det->panels[0].min_ss = 0;
	image.det->panels[0].max_ss = 128;
	image.det->panels[0].fsx = 1.0;
	image.det->panels[0].fsy = 0.0;
	image.det->panels[0].ssx = 0.0;
	image.det->panels[0].ssy = 1.0;
	image.det->panels[0].xfs = 1.0;
	image.det->panels[0].yfs = 0.0;
	image.det->panels[0].xss = 0.0;
	image.det->panels[0].yss = 1.0;
	image.det->panels[0].cnx = -64.0;
	image.det->panels[0].cny = -64.0;
	image.det->panels[0].clen = 1.0;
	image.det->panels[0].res = 1.0;
	image.det->panels[0].integr_radius = 10.0;
	image.det->panels[0].adu_per_eV = 1.0;

	image.width = 128;
	image.height = 128;
	memset(image.data, 0, 128*128*sizeof(float));

	/* First check: no intensity -> no peak, or very low intensity */
	r = integrate_peak(&image, 64, 64, &fsp, &ssp, &intensity, &sigma);
	STATUS("  First check: integrate_peak() returned %i", r);
	if ( r == 0 ) {

		STATUS("\n");
		STATUS("Intensity = %.2f, sigma = %.2f\n", intensity, sigma);

		if ( fabs(intensity) > 0.01 ) {
			ERROR("Intensity should be very close to zero.\n");
			fail = 1;
		}

	} else {
		STATUS(" (correct)\n");
	}

	/* Second check: uniform peak gives correct I and low sigma(I) */
	npx = 0;
	for ( fs=0; fs<image.width; fs++ ) {
	for ( ss=0; ss<image.height; ss++ ) {
		if ( (fs-64)*(fs-64) + (ss-64)*(ss-64) > 9*9 ) continue;
		image.data[fs+image.width*ss] = 1000.0;
		npx++;
	}
	}

	r = integrate_peak(&image, 64, 64, &fsp, &ssp, &intensity, &sigma);
	if ( r ) {
		ERROR(" Second check: integrate_peak() returned %i (wrong).\n",
		      r);
		fail = 1;
	} else {

		STATUS(" Second check: intensity = %.2f, sigma = %.2f\n",
		       intensity, sigma);

		if ( fabs(intensity - M_PI*9.0*9.0*1000.0) > 4000.0 ) {
			ERROR("Intensity should be close to %f\n", npx*1000.0);
			fail = 1;
		}
		if ( fabs(sigma-sqrt(intensity)) > 2.0 ) {
			ERROR("Sigma should be roughly %f.\n", sqrt(intensity));
			fail = 1;
		}

	}

	/* Third check: Poisson background should get mostly subtracted */
	third_integration_check(&image, n_trials, &fail);

	/* Fourth check: peak on Poisson background */
	fourth_integration_check(&image, n_trials, &fail);

	/* Fifth check: Like third check but without background subtraction */
	fifth_integration_check(&image, n_trials, &fail);

	/* Sixth check: Like fourth check but without background subtraction */
	sixth_integration_check(&image, n_trials, &fail);

	free(image.beam);
	free(image.det->panels);
	free(image.det);
	free(image.data);

	if ( fail ) return 1;

	return 0;
}