//# tProfileFit1D.cc: test the ProfileFit1D class
//# Copyright (C) 1995,1996,1998,1999,2000,2001,2002,2004
//# Associated Universities, Inc. Washington DC, USA.
//#
//# 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.,
//# 675 Massachusetts Ave, Cambridge, MA 02139, USA.
//#
//# Correspondence concerning AIPS++ should be addressed as follows:
//#        Internet email: aips2-request@nrao.edu.
//#        Postal address: AIPS++ Project Office
//#                        National Radio Astronomy Observatory
//#                        520 Edgemont Road
//#                        Charlottesville, VA 22903-2475 USA
//#
//# $Id: tProfileFit1D.cc 21451 2014-06-10 07:48:08Z gervandiepen $

#include <casa/aips.h>
#include <components/SpectralComponents/ProfileFit1D.h>

#include <casa/Containers/Record.h>
#include <casa/Arrays/ArrayLogical.h>
#include <casa/Arrays/ArrayMath.h>
#include <casa/Arrays/IPosition.h>
#include <casa/Arrays/Vector.h>
#include <casa/Utilities/Assert.h>
#include <components/SpectralComponents/GaussianMultipletSpectralElement.h>

#include <casa/iostream.h>

#include <casa/namespace.h>
void makeData (Vector<Double>& x, Vector<Double>& y, Vector<Bool>& m, 
		Double& amp, Double& cen, Double& sigma,
		Double& p0, Double& p1);
void check (Double amp, Double cen, Double sigma, Double p0, Double p1,
		const SpectralList& l);
void checkMasks (uInt n, const ProfileFit1D<Double>& fitter, Int start,
		Int end);

GaussianMultipletSpectralElement makeMultiplet (
	Vector<Double>& x, Vector<Double>& y,
	const Vector<Double>& amp, const Vector<Double>& cen,
	const Vector<Double>& sigma
);

vector<LorentzianSpectralElement> makeLorentzians (
	Vector<Double>& x, Vector<Double>& y,
	const Vector<Double>& amp, const Vector<Double>& cen,
	const Vector<Double>& fwhm
);

PowerLogPolynomialSpectralElement makePowerLogPoly(
	Vector<Double>&x, Vector<Double>& y, const Vector<Double>& coeffs
);

PolynomialSpectralElement makePoly(
	Vector<Double>&x, Vector<Double>& y, const Vector<Double>& coeffs
);

int main() {

	try {
		{
			// Data
			Vector<Double> x,y;
			Vector<Bool> m;
			Double amp, cen, sig, p0, p1;
			makeData(x, y, m, amp, cen, sig, p0, p1);
			const uInt n = x.nelements();

			// Make fitter, set data and fit
			ProfileFit1D<Double> fitter;
			fitter.setData (x,y,m);
			fitter.setGaussianElements (1);
			const SpectralElement *firstEl = fitter.getList(False)[0];
			AlwaysAssert(
				firstEl->getType() == SpectralElement::GAUSSIAN, AipsError
			);
			PolynomialSpectralElement p(1);
			fitter.addElement(p);
			AlwaysAssert(fitter.fit(), AipsError);
			// Check ok

			AlwaysAssert(fitter.getDataMask().nelements()==n, AipsError);
			AlwaysAssert(allEQ(fitter.getDataMask(), True), AipsError);
			AlwaysAssert(fitter.getRangeMask().nelements()==0, AipsError);
			AlwaysAssert(fitter.getTotalMask().nelements()==n, AipsError);
			AlwaysAssert(allEQ(fitter.getTotalMask(), True), AipsError);
			{
				const SpectralList& fitList = fitter.getList(True);
				check (amp, cen, sig, p0, p1, fitList);
			}
			{
				ProfileFit1D<Double> fitter2(fitter);
				const SpectralList& fitList = fitter2.getList(True);
				check (amp, cen, sig, p0, p1, fitList);
			}
			{
				ProfileFit1D<Double> fitter2;
				fitter2 = fitter;
				const SpectralList& fitList = fitter2.getList(True);
				check (amp, cen, sig, p0, p1, fitList);
			}
			// Set a range mask via indices

			{
				Vector<uInt> start(1), end(1);
				start(0) = n/2; end(0) = start(0) + n/10;
				fitter.setRangeMask (start, end, True);

				// Check masks

				checkMasks (n, fitter, start(0), end(0));

				// Now set range mask via abcissa values

				Vector<Double> startF(1), endF(1);
				startF(0) = x(start(0)); endF(0) = x(end(0));
				fitter.setRangeMask (startF, endF, True);

				// Check masks

				checkMasks (n, fitter, start(0), end(0));
			}
		}

		{
			Vector<Double> x, y, amp(2), cen(2), sigma(2);
			amp[0] = 1;
			amp[1] = 1;
			cen[0] = 10;
			cen[1] = 60;
			sigma[0] = 6;
			sigma[1] = 6;
			GaussianMultipletSpectralElement gm0 = makeMultiplet (x, y, amp, cen, sigma);
			ProfileFit1D<Double> fitter;
			Vector<Bool> m (x.size(), True);
			fitter.setData (x,y,m);
			fitter.addElement(gm0);
			AlwaysAssert(
				fitter.getList(False).nelements() == 1, AipsError
			);
			const SpectralElement *firstEl = fitter.getList(False)[0];
			AlwaysAssert(
				firstEl->getType() == SpectralElement::GMULTIPLET, AipsError
			);
			AlwaysAssert(fitter.fit(), AipsError);
			// Check ok
			AlwaysAssert(fitter.getDataMask().nelements() == x.size(), AipsError);
			AlwaysAssert(allEQ(fitter.getDataMask(), True), AipsError);
			AlwaysAssert(fitter.getRangeMask().nelements() == 0, AipsError);
			AlwaysAssert(fitter.getTotalMask().nelements() == x.size(), AipsError);
			AlwaysAssert(allEQ(fitter.getTotalMask(), True), AipsError);
			AlwaysAssert(
				*dynamic_cast<const GaussianMultipletSpectralElement*>(
					fitter.getList(True)[0]
				)
				== *dynamic_cast<const GaussianMultipletSpectralElement*>(
						fitter.getList(False)[0]
					),
				AipsError
			);
			Matrix<Double> r(1, 3);
			GaussianMultipletSpectralElement gm = gm0;

			cout << "amplitude ratio constrained" << endl;
			fitter = ProfileFit1D<Double>();
			fitter.setData (x,y,m);
			r = 0;
			r(0, 0) = 0.9;
			gm = GaussianMultipletSpectralElement(gm0.getGaussians(), r);
			fitter.addElement(gm);
			AlwaysAssert(fitter.fit(), AipsError);
			cout << *dynamic_cast<const GaussianMultipletSpectralElement *>(
					fitter.getList(True)[0]
				)
				<< endl;
			cout << "niter " << fitter.getNumberIterations() << endl;
			cout << endl;

			cout << "amplitude ratio constrained, sigma of reference fixed" << endl;
			fitter = ProfileFit1D<Double>();
			fitter.setData (x,y,m);
			r = 0;
			r(0, 0) = 0.9;
			gm = GaussianMultipletSpectralElement(gm0.getGaussians(), r);
			Vector<Bool> fixed(5, False);
			fixed[2] = True;
			gm.fix(fixed);
			fitter.addElement(gm);
			AlwaysAssert(fitter.fit(), AipsError);
			cout << *dynamic_cast<const GaussianMultipletSpectralElement *>(
					fitter.getList(True)[0]
				)
				<< endl;
			cout << "niter " << fitter.getNumberIterations() << endl;
			cout << endl;

			cout << "center offset constrained" << endl;
			fitter = ProfileFit1D<Double>();
			fitter.setData (x,y,m);
			r = 0;
			r(0, 1) = 48.5;
			gm = GaussianMultipletSpectralElement(gm0.getGaussians(), r);
			fitter.addElement(gm);
			AlwaysAssert(fitter.fit(), AipsError);
			cout << *dynamic_cast<const GaussianMultipletSpectralElement *>(
					fitter.getList(True)[0]
				)
				<< endl;
			cout << "niter " << fitter.getNumberIterations() << endl;
			cout << endl;

			cout << "sigma ratio constrained" << endl;
			fitter = ProfileFit1D<Double>();
			fitter.setData (x,y,m);
			r = 0;
			r(0, 2) = 0.9;
			gm = GaussianMultipletSpectralElement(gm0.getGaussians(), r);
			fitter.addElement(gm);
			AlwaysAssert(fitter.fit(), AipsError);
			cout << *dynamic_cast<const GaussianMultipletSpectralElement *>(
					fitter.getList(True)[0]
				)
				<< endl;
			cout << "niter " << fitter.getNumberIterations() << endl;
			cout << endl;
		}
		{
			Vector<Double> x, y, amp(2), cen(2), fwhm(2);
			amp[0] = 1.5;
			amp[1] = 4;
			cen[0] = 10;
			cen[1] = 60;
			fwhm[0] = 6;
			fwhm[1] = 6.5;
			vector<LorentzianSpectralElement> lse = makeLorentzians (x, y, amp, cen, fwhm);

			ProfileFit1D<Double> fitter;
			Vector<Bool> m (x.size(), True);
			fitter.setData (x,y,m);
			for (uInt i=0; i<lse.size(); i++) {
				// perturb the initial estimates
				LorentzianSpectralElement z = lse[i];
				z.setAmpl(amp[i] + (i*0.25));
				z.setCenter(cen[i] + 3 + 0.5*i);
				z.setFWHM(fwhm[i] + i*2);
				fitter.addElement(z);
			}
			AlwaysAssert(
				fitter.getList(False).nelements() == lse.size(), AipsError
			);
			const SpectralElement *firstEl = fitter.getList(False)[0];
			AlwaysAssert(
				firstEl->getType() == SpectralElement::LORENTZIAN, AipsError
			);
			AlwaysAssert(fitter.fit(), AipsError);
			// Check ok
			AlwaysAssert(fitter.getDataMask().nelements() == x.size(), AipsError);
			AlwaysAssert(allEQ(fitter.getDataMask(), True), AipsError);
			AlwaysAssert(fitter.getRangeMask().nelements() == 0, AipsError);
			AlwaysAssert(fitter.getTotalMask().nelements() == x.size(), AipsError);
			AlwaysAssert(allEQ(fitter.getTotalMask(), True), AipsError);
			for (uInt i=0; i<lse.size(); i++) {
				const LorentzianSpectralElement *got = dynamic_cast<
					const LorentzianSpectralElement*
				>(
					fitter.getList(True)[i]
				);
				LorentzianSpectralElement exp = lse[i];
				AlwaysAssert(nearAbs(*got, exp, 1e-15), AipsError);
			}

			cout << "niter " << fitter.getNumberIterations() << endl;
			cout << endl;
		}
		cout << "*** Fit a power log polynomial" << endl;
		{
			ProfileFit1D<Double> fitter;
			Vector<Double> x, y;
			Vector<Double> estimates(2);
			estimates[0] = 0.5;
			estimates[1] = 2;
			makePowerLogPoly(x, y, estimates);
			Vector<Bool> mask(x.size(), True);
			fitter.setData(x, y, mask);
			SpectralList list;
			estimates[1] = 1;
			estimates[0] = 1;
			list.add(PowerLogPolynomialSpectralElement(estimates));
			fitter.setElements(list);
			AlwaysAssert(fitter.fit(), AipsError);
			Vector<Double> parms = fitter.getList(True)[0]->get();
			cout << "parms " << parms << endl;
			AlwaysAssert(near(parms[0], 0.5) && near(parms[1], 2.0), AipsError);
		}
		{
			ProfileFit1D<Double> fitter;
			Vector<Double> x, y;
			Vector<Double> estimates(3);
			estimates[0] = 0.5;
			estimates[1] = 2;
			estimates[2] = 0;
			makePowerLogPoly(x, y, estimates);
			Vector<Bool> mask(x.size(), True);
			fitter.setData(x, y, mask);
			SpectralList list;
			estimates[0] = 0.55;
			estimates[1] = 1.93;
			estimates[2] = 0.4;
			list.add(PowerLogPolynomialSpectralElement(estimates));
			fitter.setElements(list);
			AlwaysAssert(fitter.fit(), AipsError);
			Vector<Double> parms = fitter.getList(True)[0]->get();
			cout << "parms " << parms << endl;
			AlwaysAssert(near(parms[0], 0.5) && near(parms[1], 2.0), AipsError);
		}
		{
			ProfileFit1D<Double> fitter;
			Vector<Double> x, y;
			Vector<Bool> mask(x.size(), True);
			Vector<Double> estimates(3);
			estimates[0] = 0.5;
			estimates[1] = 2;
			estimates[2] = 1;
			makePowerLogPoly(x, y, estimates);
			fitter.setData(x, y, mask);
			estimates[1] = 1.99;
			estimates[2] = 0.999;
			SpectralList list;
			list.add(PowerLogPolynomialSpectralElement(estimates));
			fitter.setElements(list);
			AlwaysAssert(fitter.fit(), AipsError);
			Vector<Double> parms = fitter.getList()[0]->get();
			cout << "parms " << parms << endl;
			AlwaysAssert(
				near(parms[0], 0.5, 1e-4)
				&& near(parms[1], 2.0, 1e-4)
				&& nearAbs(parms[2], 1.0, 1e-4),
				AipsError
			);
		}
		{
			cout << "*** Fit a polynomial" << endl;
			ProfileFit1D<Double> fitter;
			Vector<Double> x, y;
			Vector<Bool> mask(x.size(), True);
			Vector<Double> estimates(3);
			estimates[0] = 0.5;
			estimates[1] = 2;
			estimates[2] = 1;
			makePoly(x, y, estimates);
			fitter.setData(x, y, mask);
			SpectralList list;
			list.add(PolynomialSpectralElement(estimates));
			fitter.setElements(list);
			AlwaysAssert(fitter.fit(), AipsError);
			Vector<Double> parms = fitter.getList()[0]->get();
			cout << "parms " << parms << endl;
			cout << "niter " << fitter.getNumberIterations() << endl;
			AlwaysAssert(allNear(parms, estimates, 1e-5), AipsError);
			Vector<Double> bad(3);
			bad[0] = 2;
			bad[1] = 5;
			bad[2] = 3;
			list.clear();
			list.add(PolynomialSpectralElement(bad));
			fitter.clearList();
			fitter.setElements(list);
			AlwaysAssert(fitter.fit(), AipsError);
			parms = fitter.getList()[0]->get();
			cout << "niter " << fitter.getNumberIterations() << endl;
			cout << "parms " << parms << endl;
			AlwaysAssert(allNear(parms, estimates, 1e-5), AipsError);

		}
		cout << "OK" << endl;
		return 0;
	} catch (const AipsError& err) {
		cerr << err.getMesg() << endl;
	}
	return 1;
}

void makeData (Vector<Double>& x, Vector<Double>& y, Vector<Bool>& m,
		Double& amp, Double& cen, Double& sigma,
		Double& p0, Double& p1)
{
	Int n = 256;
	x.resize(n);
	y.resize(n);
	m.resize(n);
	indgen(x);
	x *= (2.3);
	x += (1.0);
	m = True;
	//
	amp = 10.0;
	cen = x(n/2);
	sigma = (x[n-1] - x[0]) / 50.0;
	p0 = 0.15;
	p1 = 1.2;
	GaussianSpectralElement g(amp, cen, sigma);
	cerr << "Gaussian: " << amp << ", " << cen << ", " << sigma << endl;
	cerr << "Polynomial: " << p0 << ", " << p1 << endl;
	//
	Vector<Double> pars(2);
	pars(0) = p0;
	pars(1) = p1;
	PolynomialSpectralElement p(pars);
	for (uInt i=0; i<x.nelements(); i++) {
		y(i) = g(x[i]) + p(x[i]);
	}
}

void check (Double amp, Double cen, Double sig, Double p0, Double p1,
		const SpectralList& list)
{
	Double tol(1e-4);
	Vector<Double> p;
	const SpectralElement *elG = list[0];
	const SpectralElement *elP = list[1];

	elG->get(p);
	cout << "p " << p << " amp " << amp << " tol " << tol << endl;
	AlwaysAssert(near(amp, p[0], tol), AipsError);
	AlwaysAssert(near(cen, p[1], tol), AipsError);
	AlwaysAssert(near(sig, p[2], tol), AipsError);
	p.resize(0);
	elP->get(p);
	AlwaysAssert(near(p0, p[0], tol), AipsError);
	AlwaysAssert(near(p1, p[1], tol), AipsError);
}


void checkMasks (uInt n, const ProfileFit1D<Double>& fitter, Int start,
		Int end)
{
	Vector<Bool> rangeMask = fitter.getRangeMask();
	Vector<Bool> totalMask = fitter.getTotalMask();
	//
	AlwaysAssert(rangeMask.nelements()==n, AipsError);
	AlwaysAssert(totalMask.nelements()==n, AipsError);
	AlwaysAssert(allEQ(rangeMask, totalMask), AipsError);
	//
	IPosition iStart(1), iEnd(1);
	{
		iStart(0) = 0;
		iEnd(0) = start-1;
		Vector<Bool> tmp = rangeMask(iStart, iEnd);
		AlwaysAssert(allEQ(tmp, False), AipsError);
	}
	{
		iStart(0) = start;
		iEnd(0) = end;
		Vector<Bool> tmp = rangeMask(iStart, iEnd);
		AlwaysAssert(allEQ(tmp, True), AipsError);
	}
	{
		iStart(0) = end+1;
		iEnd(0) = n-1;
		Vector<Bool> tmp = rangeMask(iStart, iEnd);
		AlwaysAssert(allEQ(tmp, False), AipsError);
	}
}

GaussianMultipletSpectralElement makeMultiplet (
	Vector<Double>& x, Vector<Double>& y,
	const Vector<Double>& amp, const Vector<Double>& cen,
	const Vector<Double>& sigma
) {
	Double minx = cen[0] - 5*sigma[0];
	Double maxx = cen[0] + 5*sigma[0];
	for (uInt i=1; i<amp.size(); i++) {
		minx = min(minx, cen[i] - 5*sigma[i]);
		maxx = max(maxx, cen[i] + 5*sigma[i]);
	}
	minx = (int)minx;
	maxx = (int)maxx + 1;
	x.resize((Int)(maxx-minx+1));
	indgen(x);
	x += minx;
	y.resize(x.size());
	vector<GaussianSpectralElement> g(amp.size());
	Matrix<Double> r(g.size() - 1, 3, 0);
	for (uInt i=0; i<amp.size(); i++) {
		g[i] = GaussianSpectralElement(amp[i], cen[i], sigma[i]);
		if(i > 0) {
			r(i-1, 0) = amp[i]/amp[0];
		}
	}
	GaussianMultipletSpectralElement gm(g, r);
	for (uInt i=0; i<x.size(); i++) {
		y[i] = gm(x[i]);
	}
	return gm;
}

vector<LorentzianSpectralElement> makeLorentzians (
	Vector<Double>& x, Vector<Double>& y,
	const Vector<Double>& amp, const Vector<Double>& cen,
	const Vector<Double>& fwhm
) {
	Double minx = cen[0] - 5*fwhm[0];
	Double maxx = cen[0] + 5*fwhm[0];
	for (uInt i=1; i<amp.size(); i++) {
		minx = min(minx, cen[i] - 5*fwhm[i]);
		maxx = max(maxx, cen[i] + 5*fwhm[i]);
	}
	minx = (int)minx;
	maxx = (int)maxx + 1;
	x.resize((Int)(maxx-minx+1));
	indgen(x);
	x += minx;
	y.resize(x.size());
	vector<LorentzianSpectralElement> lse;
	for (uInt i=0; i<amp.size(); i++) {
		lse.push_back(LorentzianSpectralElement(amp[i], cen[i], fwhm[i]));
	}
	for (uInt i=0; i<x.size(); i++) {
		y[i] = 0;
		for (uInt j=0; j<lse.size(); j++) {
			y[i] += lse[j](x[i]);
		}
	}
	return lse;
}


PowerLogPolynomialSpectralElement makePowerLogPoly(
	Vector<Double>&x, Vector<Double>& y, const Vector<Double>& coeffs
) {
	x.resize(10);
	x[0] = 1;
	x[1] = 3;
	x[2] = 5;
	x[3] = 6;
	x[4] = 8;
	x[5] = 10;
	x[6] = 12;
	x[7] = 15;
	x[8] = 20;
	x[9] = 25;
	y.resize(x.size());
	PowerLogPolynomialSpectralElement plp(coeffs);
	for (uInt i=0; i<x.size(); i++) {
		y[i] = plp(x[i]);
	}
	return plp;

}

PolynomialSpectralElement makePoly(
	Vector<Double>&x, Vector<Double>& y, const Vector<Double>& coeffs
) {
	x.resize(10);
	x[0] = 1;
	x[1] = 3;
	x[2] = 5;
	x[3] = 6;
	x[4] = 8;
	x[5] = 10;
	x[6] = 12;
	x[7] = 15;
	x[8] = 20;
	x[9] = 25;
	y.resize(x.size());
	PolynomialSpectralElement poly(coeffs);
	for (uInt i=0; i<x.size(); i++) {
		y[i] = poly(x[i]);
	}
	return poly;

}