//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Fitting/FitObjective.cpp
//! @brief     Implements class FitObjective.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2018
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Sim/Fitting/FitObjective.h"
#include "Base/Util/Assert.h"
#include "Device/Data/DataUtil.h"
#include "Device/Data/Datafield.h"
#include "Fit/Minimizer/MinimizerResult.h"
#include "Sim/Fitting/FitStatus.h"
#include "Sim/Fitting/ObjectiveMetric.h"
#include "Sim/Fitting/ObjectiveMetricUtil.h"
#include "Sim/Fitting/PyFittingCallbacks.h"
#include "Sim/Fitting/SimDataPair.h"
#include "Sim/Residual/ChiSquaredModule.h"
#include "Sim/Simulation/ISimulation.h"
#include <algorithm>
#include <iostream>
#include <stdexcept>

namespace {

simulation_builder_t simulationBuilder(const PyBuilderCallback& callback)
{
    return [&callback](const mumufit::Parameters& params) {
        return std::unique_ptr<ISimulation>(callback.build_simulation(params));
    };
}

} // namespace

//  ************************************************************************************************
//  definition of auxiliary class hierarchy IMetricWrapper
//  ************************************************************************************************

class IMetricWrapper {
public:
    virtual ~IMetricWrapper() = default;
    virtual double compute(const std::vector<SimDataPair>& fit_objects, size_t n_pars) const = 0;
};

//! Metric wrapper for back-compaptibility with old scripts
class ChiModuleWrapper : public IMetricWrapper {
public:
    explicit ChiModuleWrapper(std::unique_ptr<IChiSquaredModule> module);
    double compute(const std::vector<SimDataPair>& fit_objects, size_t n_pars) const override;

private:
    std::unique_ptr<IChiSquaredModule> m_module;
};

class ObjectiveMetricWrapper : public IMetricWrapper {
public:
    explicit ObjectiveMetricWrapper(std::unique_ptr<ObjectiveMetric> module);
    double compute(const std::vector<SimDataPair>& fit_objects, size_t n_pars) const override;

private:
    std::unique_ptr<ObjectiveMetric> m_module;
};

//  ************************************************************************************************
//  implementation of auxiliary class hierarchy IMetricWrapper
//  ************************************************************************************************

ChiModuleWrapper::ChiModuleWrapper(std::unique_ptr<IChiSquaredModule> module)
    : m_module(std::move(module))
{
    ASSERT(m_module);
}

double ChiModuleWrapper::compute(const std::vector<SimDataPair>& fit_objects, size_t n_pars) const
{
    size_t n_points = 0;
    double result = 0;
    for (const auto& obj : fit_objects) {
        const auto sim_array = obj.simulation_array();
        const auto exp_array = obj.experimental_array();
        const size_t n_elements = sim_array.size();
        double contrib = 0;
        for (size_t i = 0; i < n_elements; ++i) {
            double value = m_module->residual(sim_array[i], exp_array[i]);
            contrib += std::pow(value, 2);
        }
        result += obj.weight() * contrib;
        n_points += n_elements;
    }

    int fnorm = static_cast<int>(n_points) - static_cast<int>(n_pars);
    if (fnorm <= 0)
        throw std::runtime_error("Error in ChiModuleWrapper: Normalization shall be positive");

    return result / fnorm;
}

ObjectiveMetricWrapper::ObjectiveMetricWrapper(std::unique_ptr<ObjectiveMetric> module)
    : m_module(std::move(module))
{
    ASSERT(m_module);
}

double ObjectiveMetricWrapper::compute(const std::vector<SimDataPair>& fit_objects, size_t) const
{
    // deciding whether to use uncertainties in metrics computation.
    bool use_uncertainties = true;
    for (const auto& obj : fit_objects)
        use_uncertainties = use_uncertainties && obj.containsUncertainties();

    double result = 0.0;
    for (const auto& obj : fit_objects)
        result += obj.weight() * m_module->computeMetric(obj, use_uncertainties);
    return result;
}

//  ************************************************************************************************
//  implementation of class FitObjective
//  ************************************************************************************************

FitObjective::FitObjective()
    : m_metric_module(
        std::make_unique<ObjectiveMetricWrapper>(std::make_unique<PoissonLikeMetric>()))
    , m_fit_status(std::make_unique<FitStatus>(this))
{
}

FitObjective::~FitObjective() = default;

//! Constructs simulation/data pair for later fit.
//! @param builder: simulation builder capable of producing simulations
//! @param data: experimental data array
//! @param stdv: data uncertainties array
//! @param weight: weight of dataset in metric calculations
//!
//! Should be private, but is used in several tests.
void FitObjective::execAddSimulationAndData(const simulation_builder_t& builder,
                                            const Datafield& data, double weight)
{
    m_fit_objects.emplace_back(builder, data, weight);
}

void FitObjective::addFitPair(const PyBuilderCallback& callback, Datafield expData, double weight)
{
    execAddSimulationAndData(simulationBuilder(callback), expData, weight);
}

double FitObjective::evaluate(const mumufit::Parameters& params)
{
    execSimulations(params);
    const double metric_value = m_metric_module->compute(m_fit_objects, params.size());
    m_fit_status->update(params, metric_value);
    return metric_value;
}

std::vector<double> FitObjective::evaluate_residuals(const mumufit::Parameters& params)
{
    evaluate(params);

    std::vector<double> result = flatExpData(); // init result with experimental data values
    const std::vector<double> sim_values = flatSimData();
    std::transform(result.begin(), result.end(), sim_values.begin(), result.begin(),
                   [](double lhs, double rhs) { return lhs - rhs; });
    return result;
}

//! Returns simulation result in the form of Datafield.
Datafield FitObjective::simulationResult(size_t i_item) const
{
    return dataPair(i_item).simulationResult();
}

//! Returns experimental data in the form of Datafield.
Datafield FitObjective::experimentalData(size_t i_item) const
{
    return dataPair(i_item).experimentalData();
}

//! Returns relative difference between simulation and experimental data
//! in the form of Datafield, for use in plotting.
Datafield FitObjective::relativeDifference(size_t i_item) const
{
    return dataPair(i_item).relativeDifference();
}

//! Returns absolute value of difference between simulation and experimental data
//! in the form of Datafield, for use in plotting.
Datafield FitObjective::absoluteDifference(size_t i_item) const
{
    return dataPair(i_item).absoluteDifference();
}

//! Returns one-dimensional array representing merged experimental data.
//! The area outside of the region of interest is not included, masked data is nullified.
std::vector<double> FitObjective::flatExpData() const
{
    return composeArray(&SimDataPair::experimental_array);
}

//! Returns one-dimensional array representing merged simulated intensities data.
//! The area outside of the region of interest is not included, masked data is nullified.
std::vector<double> FitObjective::flatSimData() const
{
    return composeArray(&SimDataPair::simulation_array);
}

const SimDataPair& FitObjective::dataPair(size_t i_item) const
{
    return m_fit_objects.at(i_item);
}

void FitObjective::initPrint(int every_nth)
{
    m_fit_status->initPrint(every_nth);
}

void FitObjective::initPlot(int every_nth, fit_observer_t&& observer)
{
    m_fit_status->addObserver(every_nth, std::move(observer));
}

void FitObjective::initPlot(int every_nth, PyObserverCallback& callback)
{
    fit_observer_t observer = [&](const FitObjective& objective) { callback.update(objective); };
    m_fit_status->addObserver(every_nth, std::move(observer));
}

bool FitObjective::isCompleted() const
{
    return m_fit_status->isCompleted();
}

IterationInfo FitObjective::iterationInfo() const
{
    return m_fit_status->iterationInfo();
}

mumufit::MinimizerResult FitObjective::minimizerResult() const
{
    return m_fit_status->minimizerResult();
}

void FitObjective::finalize(const mumufit::MinimizerResult& result)
{
    m_fit_status->finalize(result);
}

unsigned FitObjective::fitObjectCount() const
{
    return static_cast<unsigned>(m_fit_objects.size());
}

void FitObjective::interruptFitting()
{
    m_fit_status->setInterrupted();
}

bool FitObjective::isInterrupted() const
{
    return m_fit_status->isInterrupted();
}

bool FitObjective::isFirstIteration() const
{
    return iterationInfo().iterationCount() == 1;
}

void FitObjective::execSimulations(const mumufit::Parameters& params)
{
    if (m_fit_status->isInterrupted())
        throw std::runtime_error("Fitting was interrupted by the user.");

    if (m_fit_objects.empty())
        throw std::runtime_error("FitObjective::execSimulations -> Error. "
                                 "No simulation/data defined.");

    for (auto& obj : m_fit_objects)
        obj.execSimulation(params);
}

void FitObjective::setChiSquaredModule(const IChiSquaredModule& module)
{
    std::cout << "Warning in FitObjective::setChiSquaredModule: setChiSquaredModule is deprecated "
                 "and will be removed in future versions. Please use "
                 "FitObjective::setObjectiveMetric instead."
              << std::endl;

    std::unique_ptr<IChiSquaredModule> chi_module(module.clone());
    m_metric_module = std::make_unique<ChiModuleWrapper>(std::move(chi_module));
}

void FitObjective::setObjectiveMetric(std::unique_ptr<ObjectiveMetric> metric)
{
    m_metric_module = std::make_unique<ObjectiveMetricWrapper>(std::move(metric));
}

void FitObjective::setObjectiveMetric(const std::string& metric)
{
    m_metric_module = std::make_unique<ObjectiveMetricWrapper>(
        ObjectiveMetricUtil::createMetric(metric, ObjectiveMetricUtil::defaultNormName()));
}

void FitObjective::setObjectiveMetric(const std::string& metric, const std::string& norm)
{
    m_metric_module =
        std::make_unique<ObjectiveMetricWrapper>(ObjectiveMetricUtil::createMetric(metric, norm));
}

//! Returns true if the specified DataPair element contains uncertainties
bool FitObjective::containsUncertainties(size_t i_item) const
{
    return dataPair(i_item).containsUncertainties();
}

//! Returns true if all the data pairs in FitObjective instance contain uncertainties
bool FitObjective::allPairsHaveUncertainties() const
{
    bool result = true;
    for (size_t i = 0, size = fitObjectCount(); i < size; ++i)
        result = result && dataPair(i).containsUncertainties();
    return result;
}

std::vector<double> FitObjective::composeArray(DataPairAccessor getter) const
{
    const size_t n_objs = m_fit_objects.size();
    if (n_objs == 0)
        return {};
    if (n_objs == 1)
        return (m_fit_objects[0].*getter)();

    std::vector<double> result;
    for (const auto& pair : m_fit_objects) {
        std::vector<double> array = (pair.*getter)();
        std::move(array.begin(), array.end(), std::back_inserter(result));
    }
    return result;
}
