/*============================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center (DKFZ)
All rights reserved.

Use of this source code is governed by a 3-clause BSD license that can be
found in the LICENSE file.

============================================================================*/

#ifndef mitkModelFitFunctorBase_h
#define mitkModelFitFunctorBase_h

#include <itkObject.h>

#include <mitkVector.h>

#include <mitkModelBase.h>
#include <mitkSVModelFitCostFunction.h>

#include <MitkModelFitExports.h>

#include <mutex>

namespace mitk
{

  /**
   * \class ModelFitFunctorBase
   * \brief Abstract base class for functors that perform model fitting.
   *
   * A model fit functor takes an observed signal, a model instance, and initial parameters,
   * then performs the fitting optimization. Subclasses implement specific fitting algorithms
   * (e.g. Levenberg-Marquardt). The functor also supports user-defined evaluation cost
   * functions and optional debug parameter maps.
   *
   * \sa LevenbergMarquardtModelFitFunctor, DummyModelFitFunctor
   */
  class MITKMODELFIT_EXPORT ModelFitFunctorBase: public ::itk::Object
  {
  public:
    typedef ModelFitFunctorBase Self;
    typedef itk::Object Superclass;
    typedef itk::SmartPointer< Self >                            Pointer;
    typedef itk::SmartPointer< const Self >                      ConstPointer;

    itkTypeMacro(ModelFitFunctorBase, itk::Object);

    typedef ScalarType ParameterImagePixelType;
    typedef std::vector<ParameterImagePixelType> InputPixelArrayType;
    typedef std::vector<ParameterImagePixelType> OutputPixelArrayType;

    /**
     * \brief Performs model fitting and returns the results.
     *
     * The values in the returned vector are ordered in the following sequence:
     * - model parameters (see also GetParameterNames())
     * - derived model parameters (see also GetDerivedParameterNames())
     * - criterion(s) (see also GetCriterionNames())
     * - evaluation parameters (see also GetEvaluationParameterNames())
     *
     * \param value Signal the model should be fitted onto.
     * \param model Pointer to the preconfigured model instance.
     * \param initialParameters Starting point parameters for the fitting process.
     * \return Vector of fitted parameter values, derived parameters, criteria, and evaluation parameters.
     * \pre model must point to a valid instance.
     * \pre Size of initialParameters must be equal to model->GetNumberOfParameters().
     */
    OutputPixelArrayType Compute(const InputPixelArrayType& value, const ModelBase* model,
                                 const ModelBase::ParametersType& initialParameters) const;

    /**
     * \brief Returns the number of outputs the fit functor will produce.
     *
     * The number depends in part on the passed model.
     * \param model The model for which the output count is determined.
     * \return The total number of output values.
     * \throw Exception if no valid model is passed.
     */
    unsigned int GetNumberOfOutputs(const ModelBase* model) const;

    typedef ModelBase::ParameterNamesType ParameterNamesType;

    /** \brief Returns names of all evaluation parameters defined by the user. */
    ParameterNamesType GetEvaluationParameterNames() const;

    /** \brief Removes all registered evaluation parameters. */
    void ResetEvaluationParameters();

    /**
     * \brief Registers an evaluation parameter with its cost function.
     * \param parameterName The name of the evaluation parameter.
     * \param evaluationCostFunction The cost function used to compute this evaluation parameter.
     */
    void RegisterEvaluationParameter(const std::string& parameterName,
                                     SVModelFitCostFunction* evaluationCostFunction);

    /**
     * \brief Returns the cost function associated with the given evaluation parameter.
     * \param parameterName The name of the evaluation parameter.
     * \return Pointer to the associated cost function.
     */
    const SVModelFitCostFunction* GetEvaluationParameterCostFunction(const std::string& parameterName)
    const;

    /** \brief Returns names of the criteria used to fit the model. */
    virtual ParameterNamesType GetCriterionNames() const = 0 ;

    /**
     * \brief Returns names of the debug parameters generated by the functor.
     *
     * Returns an empty list if debug is deactivated.
     */
    ParameterNamesType GetDebugParameterNames() const;

    itkBooleanMacro(DebugParameterMaps);
    itkSetMacro(DebugParameterMaps, bool);
    itkGetConstMacro(DebugParameterMaps, bool);

  protected:

    typedef ModelBase::ParametersType ParametersType;
    typedef ModelFitCostFunctionInterface::SignalType SignalType;

    ModelFitFunctorBase();

    ~ModelFitFunctorBase() override;

    /**
     * \brief Internal method called by Compute() to get the final criterion values that drove the fit.
     *
     * Must be implemented by concrete functor classes.
     */
    virtual OutputPixelArrayType GetCriteria(const ModelBase* model, const ParametersType& parameters,
        const SignalType& sample) const = 0;

    /** \brief Internal method called by Compute() to get derived parameters with the final fit parameters. */
    OutputPixelArrayType GetDerivedParameters(const ModelBase* model,
        const ParametersType& parameters) const;

    /** \brief Internal method called by Compute() to get evaluation parameters from user-defined cost functions. */
    OutputPixelArrayType GetEvaluationParameters(const ModelBase* model,
        const ParametersType& parameters, const SignalType& sample) const;

    typedef std::map<std::string, ParameterImagePixelType> DebugParameterMapType;

    /**
     * \brief Internal method that performs the actual model fitting.
     *
     * Must be implemented by derived classes. Returns the found parameters and
     * optionally populates the debug parameter map.
     *
     * \param value Signal the model should be fitted against.
     * \param model Pointer to the model that should be fitted.
     * \param initialParameters Initial model parameters for the fit.
     * \param[out] debugParameters Map containing all debug parameters (only valid if m_DebugParameterMaps is true).
     * \return The optimized parameter values.
     * \post If m_DebugParameterMaps is true, all debug parameters defined by
     * DefineDebugParameterNames() must be returned via debugParameters.
     */
    virtual ParametersType DoModelFit(const SignalType& value, const ModelBase* model,
                                      const ModelBase::ParametersType& initialParameters,
                                      DebugParameterMapType& debugParameters) const = 0;

    /**
     * \brief Returns names of the debug parameters generated by the functor.
     *
     * Will be called by GetDebugParameterNames() if debug is activated.
     */
    virtual ParameterNamesType DefineDebugParameterNames()const = 0;

  private:

    typedef std::map<std::string, SVModelFitCostFunction::Pointer> CostFunctionMapType;
    CostFunctionMapType m_CostFunctionMap;
    bool m_DebugParameterMaps;
    mutable std::mutex m_Mutex;
  };

}


#endif
