//  ************************************************************************************************
//
//  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/SimDataPair.h"
#include "Sim/Fitting/SimulationWrapper.h"
#include "Sim/Residual/ChiSquaredModule.h"
#include "Sim/Simulation/ISimulation.h"
#include <algorithm>
#include <iostream>
#include <stdexcept>

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

class IMetricWrapper {
public:
    virtual ~IMetricWrapper() = default;
    virtual double compute(const std::vector<SimDataPair>& sim_data_pairs, 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>& sim_data_pairs, 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>& sim_data_pairs, 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>& sim_data_pairs,
                                 size_t n_pars) const
{
    size_t n_points = 0;
    double result = 0;
    for (const auto& obj : sim_data_pairs) {
        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>& sim_data_pairs, size_t) const
{
    // deciding whether to use uncertainties in metrics computation.
    bool use_uncertainties = true;
    for (const auto& obj : sim_data_pairs)
        use_uncertainties = use_uncertainties && obj.containsUncertainties();

    double result = 0.0;
    for (const auto& obj : sim_data_pairs)
        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

void FitObjective::addFitPair(const simulation_builder_t& sim_fn, const Datafield& expData,
                              const double weight)
{
    // add C++ simulation building function
    SimulationWrapper sim;
    sim.cSimulationFn = sim_fn;
    m_sim_data_pairs.emplace_back(sim, expData, weight);
}

void FitObjective::addFitPair(void* pSimulationCallable, PySimulate_t pSimulationCaller,
                              PyFree_t pFree, const Datafield& expData, const double weight)
{
    // add Python simulation building function
    SimulationWrapper sim;
    sim.pySimulationFn = pSimulationCallable;
    sim.pySimulate = pSimulationCaller;
    sim.pyFree = pFree;
    m_sim_data_pairs.emplace_back(sim, expData, weight);
}

double FitObjective::evaluate(const mumufit::Parameters& params)
{
    execSimulations(params);
    const double metric_value = m_metric_module->compute(m_sim_data_pairs, 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_sim_data_pairs.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, void* pCallable, PyCaller_t pCall)
{
    fit_observer_t observer = [pCallable, pCall](const FitObjective& objective) {
        pCall(pCallable, 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);
}

size_t FitObjective::nPairs() const
{
    return m_sim_data_pairs.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_sim_data_pairs.empty())
        throw std::runtime_error("Cannot start fit as no simulation/data pairs are defined.");

    for (auto& obj : m_sim_data_pairs)
        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; i < m_sim_data_pairs.size(); ++i)
        result = result && dataPair(i).containsUncertainties();
    return result;
}

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

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