//#---------------------------------------------------------------------------
//# SDFitter.cc: A Fitter class for spectra
//#--------------------------------------------------------------------------
//# Copyright (C) 2004
//# ATNF
//#
//# 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 this software should be addressed as follows:
//#        Internet email: Malte.Marquarding@csiro.au
//#        Postal address: Malte Marquarding,
//#                        Australia Telescope National Facility,
//#                        P.O. Box 76,
//#                        Epping, NSW, 2121,
//#                        AUSTRALIA
//#
//# $Id:
//#---------------------------------------------------------------------------
#include <casa/aips.h>
#include <casa/Arrays/ArrayMath.h>
#include <casa/Arrays/ArrayLogical.h>
#include <scimath/Fitting.h>
#include <scimath/Fitting/LinearFit.h>
#include <scimath/Functionals/CompiledFunction.h>
#include <scimath/Functionals/CompoundFunction.h>
#include <scimath/Functionals/Gaussian1D.h>
#include <scimath/Functionals/Polynomial.h>
#include <scimath/Mathematics/AutoDiff.h>
#include <scimath/Mathematics/AutoDiffMath.h>
#include <scimath/Fitting/NonLinearFitLM.h>
#include <components/SpectralComponents/SpectralEstimate.h>

#include "SDFitter.h"
using namespace asap;
using namespace casa;

SDFitter::SDFitter()
{
}

SDFitter::~SDFitter()
{
    reset();
}

void SDFitter::clear()
{
    for (uInt i=0;i< funcs_.nelements();++i) {
        delete funcs_[i]; funcs_[i] = 0;
    }
    funcs_.resize(0, True);
    parameters_.resize();
    error_.resize();
    thefit_.resize();
    estimate_.resize();
    chisquared_ = 0.0;
}
void SDFitter::reset()
{
    clear();
    x_.resize();
    y_.resize();
    m_.resize();
}


bool SDFitter::computeEstimate() {
    if (x_.nelements() == 0 || y_.nelements() == 0)
        throw (AipsError("No x/y data specified."));

    if (dynamic_cast<Gaussian1D<Float>* >(funcs_[0]) == 0)
        return false;
    uInt n = funcs_.nelements();
    SpectralEstimate estimator(n);
    estimator.setQ(5);
    Int mn,mx;
    mn = 0;
    mx = m_.nelements()-1;
    for (uInt i=0; i<m_.nelements();++i) {
      if (m_[i]) {
	mn = i;
	break;
      }
    }
    for (uInt j=m_.nelements()-1; j>=0;--j) {
      if (m_[j]) {
	mx = j;
	break;
      }
    }
    //mn = 0+x_.nelements()/10;
    //mx = x_.nelements()-x_.nelements()/10;
    estimator.setRegion(mn,mx);
    //estimator.setWindowing(True);
    SpectralList listGauss = estimator.estimate(x_, y_);
    Gaussian1D<Float>* g;
    parameters_.resize(n*3);
    uInt count = 0;
    for (uInt i=0; i<n;i++) {
        g = dynamic_cast<Gaussian1D<Float>* >(funcs_[i]);
        if (g) {
            (*g)[0] = listGauss[i].getAmpl();
            (*g)[1] = listGauss[i].getCenter();
            (*g)[2] = listGauss[i].getFWHM();
            ++count;
        }
    }
    estimate_.resize();
    listGauss.evaluate(estimate_,x_);
    return true;
}

std::vector<float> SDFitter::getEstimate() const
{
    if (estimate_.nelements() == 0)
        throw (AipsError("No estimate set."));
    std::vector<float> stlout;
    estimate_.tovector(stlout);
    return stlout;
}


bool SDFitter::setExpression(const std::string& expr, int ncomp)
{
    clear();
    if (expr == "gauss") {
        if (ncomp < 1) throw (AipsError("Need at least one gaussian to fit."));
        funcs_.resize(ncomp);
        for (Int k=0; k<ncomp; ++k) {
            funcs_[k] = new Gaussian1D<Float>();
        }
    } else if (expr == "poly") {
        funcs_.resize(1);
        funcs_[0] = new Polynomial<Float>(ncomp);
    } else {
        //cerr << " compiled functions not yet implemented" << endl;
        //funcs_.resize(1);
        //funcs_[0] = new CompiledFunction<Float>();
        //funcs_[0]->setFunction(String(expr));
        return false;
    };
    return true;
}

bool SDFitter::setData(std::vector<float> absc, std::vector<float> spec,
                       std::vector<bool> mask)
{
    x_.resize();
    y_.resize();
    m_.resize();
    // convert std::vector to casa Vector
    Vector<Float> tmpx(absc);
    Vector<Float> tmpy(spec);
    Vector<Bool> tmpm(mask);
    AlwaysAssert(tmpx.nelements() == tmpy.nelements(), AipsError);
    x_ = tmpx;
    y_ = tmpy;
    m_ = tmpm;
    return true;
}

std::vector<float> SDFitter::getResidual() const
{
    if (residual_.nelements() == 0)
        throw (AipsError("Function not yet fitted."));
    std::vector<float> stlout;
    residual_.tovector(stlout);
    return stlout;
}

std::vector<float> SDFitter::getFit() const
{
    Vector<Float> out = thefit_;
    std::vector<float> stlout;
    out.tovector(stlout);
    return stlout;

}

std::vector<float> SDFitter::getErrors() const
{
    Vector<Float> out = error_;
    std::vector<float> stlout;
    out.tovector(stlout);
    return stlout;
}

bool SDFitter::setParameters(std::vector<float> params)
{
    Vector<Float> tmppar(params);
    if (funcs_.nelements() == 0)
        throw (AipsError("Function not yet set."));
    if (parameters_.nelements() > 0 && tmppar.nelements() != parameters_.nelements())
        throw (AipsError("Number of parameters inconsistent with function."));
    if (parameters_.nelements() == 0)
        parameters_.resize(tmppar.nelements());
        fixedpar_.resize(tmppar.nelements());
        fixedpar_ = False;
    if (dynamic_cast<Gaussian1D<Float>* >(funcs_[0]) != 0) {
        uInt count = 0;
        for (uInt j=0; j < funcs_.nelements(); ++j) {
            for (uInt i=0; i < funcs_[j]->nparameters(); ++i) {
                (funcs_[j]->parameters())[i] = tmppar[count];
                parameters_[count] = tmppar[count];
                ++count;
            }
        }
    } else if (dynamic_cast<Polynomial<Float>* >(funcs_[0]) != 0) {
        for (uInt i=0; i < funcs_[0]->nparameters(); ++i) {
            parameters_[i] = tmppar[i];
            (funcs_[0]->parameters())[i] =  tmppar[i];
        }
    }
    return true;
}

bool SDFitter::setFixedParameters(std::vector<bool> fixed)
{
    Vector<Bool> tmp(fixed);
    if (funcs_.nelements() == 0)
        throw (AipsError("Function not yet set."));
    if (fixedpar_.nelements() > 0 && tmp.nelements() != fixedpar_.nelements())
        throw (AipsError("Number of mask elements inconsistent with function."));
    if (dynamic_cast<Gaussian1D<Float>* >(funcs_[0]) != 0) {
        uInt count = 0;
        for (uInt j=0; j < funcs_.nelements(); ++j) {
            for (uInt i=0; i < funcs_[j]->nparameters(); ++i) {
                funcs_[j]->mask(i) = !tmp[count];
                fixedpar_[count] = !tmp[count];
                ++count;
            }
        }
    } else if (dynamic_cast<Polynomial<Float>* >(funcs_[0]) != 0) {
        for (uInt i=0; i < funcs_[0]->nparameters(); ++i) {
            fixedpar_[i] = tmp[i];
            funcs_[0]->mask(i) =  tmp[i];
        }
    }
    //fixedpar_ = !tmpmsk;
    return true;
}

std::vector<float> SDFitter::getParameters() const {
    Vector<Float> out = parameters_;
    std::vector<float> stlout;
    out.tovector(stlout);
    return stlout;
}

std::vector<bool> SDFitter::getFixedParameters() const {
  Vector<Bool> out(parameters_.nelements());
  if (fixedpar_.nelements() == 0) {
    out = False;
    //throw (AipsError("No parameter mask set."));
  } else {
    out = fixedpar_;
  }
  std::vector<bool> stlout;
  out.tovector(stlout);
  return stlout;
}

float SDFitter::getChisquared() const {
    return chisquared_;
}

bool SDFitter::fit() {
    NonLinearFitLM<Float> fitter;
    //CompiledFunction<AutoDiff<Float> > comp;
    //Polynomial<AutoDiff<Float> > poly;
    CompoundFunction<AutoDiff<Float> > func;
    if (dynamic_cast<Gaussian1D<Float>* >(funcs_[0]) != 0) {
        //computeEstimates();
        for (uInt i=0; i<funcs_.nelements(); i++) {
            Gaussian1D<AutoDiff<Float> > gauss;//(*funcs_[i]);
	    
	    for (uInt j=0; j<funcs_[i]->nparameters(); j++) {
		gauss[j] = AutoDiff<Float>((*funcs_[i])[j], 
					   gauss.nparameters(), j);
		gauss.mask(j) = funcs_[i]->mask(j);
            }
	    
            func.addFunction(gauss);
        }
    } else if (dynamic_cast<Polynomial<Float>* >(funcs_[0]) != 0) {
	Polynomial<AutoDiff<Float> > poly(funcs_[0]->nparameters()-1);
	//Polynomial<AutoDiff<Float> > poly(*funcs_[0]);
	for (uInt j=0; j<funcs_[0]->nparameters(); j++) {
	    poly[j] = AutoDiff<Float>(0, poly.nparameters(), j);
	    poly.mask(j) = funcs_[0]->mask(j);
	}
	func.addFunction(poly);
    } else if (dynamic_cast<CompiledFunction<Float>* >(funcs_[0]) != 0) {
	
	//         CompiledFunction<AutoDiff<Float> > comp;
	//         for (uInt j=0; j<funcs_[0]->nparameters(); j++) {
//             comp[j] = AutoDiff<Float>(0, comp.nparameters(), j);
//             comp.mask(j) = funcs_[0]->mask(j);
//         }
//         func.addFunction(comp);
	
        cout << "NYI." << endl;
    } else {
        throw(AipsError("Fitter not set up correctly."));
    }
    fitter.setFunction(func);
    fitter.setMaxIter(50+funcs_.nelements()*10);
    // Convergence criterium
    fitter.setCriteria(0.001);

    // Fit
    Vector<Float> sigma(x_.nelements());
    sigma = 1.0;

    parameters_.resize();
    parameters_ = fitter.fit(x_, y_, sigma, &m_);

    error_.resize();
    error_ = fitter.errors();

    chisquared_ = fitter.getChi2();

    residual_.resize();
    residual_ =  y_;
    fitter.residual(residual_,x_);

    // use fitter.residual(model=True) to get the model
    thefit_.resize(x_.nelements());
    fitter.residual(thefit_,x_,True);
    return true;
}