//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Fitting/ObjectiveMetric.cpp
//! @brief     Implements class ObjectiveMetrices.
//!
//! @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/ObjectiveMetric.h"
#include "Base/Util/Assert.h"
#include "Device/Data/Datafield.h"
#include "Sim/Fitting/ObjectiveMetricUtil.h"
#include "Sim/Fitting/SimDataPair.h"
#include <cmath>
#include <limits>
#include <stdexcept>

namespace {

const double double_max = std::numeric_limits<double>::max();
const double double_min = std::numeric_limits<double>::min();
const double ln10 = std::log(10.0);

template <typename T> T* copyMetric(const T& metric)
{
    auto* result = new T;
    result->setNorm(metric.norm());
    return result;
}

void checkIntegrity(const std::vector<double>& sim_data, const std::vector<double>& exp_data)
{
    const size_t sim_size = sim_data.size();
    if (sim_size != exp_data.size())
        throw std::runtime_error("Error in ObjectiveMetric: input arrays have different sizes");

    for (size_t i = 0; i < sim_size; ++i)
        if (sim_data[i] < 0.0)
            throw std::runtime_error(
                "Error in ObjectiveMetric: simulation data array contains negative values");
}

void checkIntegrity(const std::vector<double>& sim_data, const std::vector<double>& exp_data,
                    const std::vector<double>& exp_stdv)
{
    if (sim_data.size() != exp_stdv.size())
        throw std::runtime_error("Error in ObjectiveMetric: input arrays have different sizes");

    checkIntegrity(sim_data, exp_data);
}

} // namespace

ObjectiveMetric::ObjectiveMetric(const std::function<double(double)>& norm)
    : m_norm(norm)
{
}

double ObjectiveMetric::computeMetric(const SimDataPair& data_pair, bool use_weights) const
{
    if (use_weights && !data_pair.containsUncertainties())
        throw std::runtime_error("Error in ObjectiveMetric::compute: the metric is weighted, but "
                                 "the simulation-data pair does not contain uncertainties");

    if (use_weights)
        return computeFromArrays(data_pair.simulation_array(), data_pair.experimental_array(),
                                 data_pair.uncertainties_array());
    return computeFromArrays(data_pair.simulation_array(), data_pair.experimental_array());
}

void ObjectiveMetric::setNorm(std::function<double(double)> norm)
{
    m_norm = std::move(norm);
}

// ----------------------- Chi2 metric ---------------------------

Chi2Metric::Chi2Metric()
    : ObjectiveMetric(ObjectiveMetricUtil::l2Norm())
{
}

Chi2Metric* Chi2Metric::clone() const
{
    return copyMetric(*this);
}

double Chi2Metric::computeFromArrays(std::vector<double> sim_data, std::vector<double> exp_data,
                                     std::vector<double> exp_stdv) const
{
    checkIntegrity(sim_data, exp_data, exp_stdv);

    double result = 0.0;
    auto norm_fun = norm();
    for (size_t i = 0, sim_size = sim_data.size(); i < sim_size; ++i)
        if (exp_data[i] >= 0.0 && exp_stdv[i] > 0.0)
            result += norm_fun((exp_data[i] - sim_data[i]) / exp_stdv[i]);

    return std::isfinite(result) ? result : double_max;
}

double Chi2Metric::computeFromArrays(std::vector<double> sim_data,
                                     std::vector<double> exp_data) const
{
    checkIntegrity(sim_data, exp_data);

    auto norm_fun = norm();
    double result = 0.0;
    for (size_t i = 0, sim_size = sim_data.size(); i < sim_size; ++i)
        if (exp_data[i] >= 0.0)
            result += norm_fun(exp_data[i] - sim_data[i]);

    return std::isfinite(result) ? result : double_max;
}

// ----------------------- Poisson-like metric ---------------------------

PoissonLikeMetric::PoissonLikeMetric() = default;

PoissonLikeMetric* PoissonLikeMetric::clone() const
{
    return copyMetric(*this);
}

double PoissonLikeMetric::computeFromArrays(std::vector<double> sim_data,
                                            std::vector<double> exp_data) const
{
    checkIntegrity(sim_data, exp_data);

    double result = 0.0;
    auto norm_fun = norm();
    for (size_t i = 0, sim_size = sim_data.size(); i < sim_size; ++i) {
        if (exp_data[i] < 0.0)
            continue;
        const double variance = std::max(1.0, sim_data[i]);
        const double value = (sim_data[i] - exp_data[i]) / std::sqrt(variance);
        result += norm_fun(value);
    }

    return std::isfinite(result) ? result : double_max;
}

// ----------------------- Log metric ---------------------------

LogMetric::LogMetric()
    : ObjectiveMetric(ObjectiveMetricUtil::l2Norm())
{
}

LogMetric* LogMetric::clone() const
{
    return copyMetric(*this);
}

double LogMetric::computeFromArrays(std::vector<double> sim_data, std::vector<double> exp_data,
                                    std::vector<double> exp_stdv) const
{
    checkIntegrity(sim_data, exp_data, exp_stdv);

    double result = 0.0;
    auto norm_fun = norm();
    for (size_t i = 0, sim_size = sim_data.size(); i < sim_size; ++i) {
        if (exp_data[i] < 0.0 || exp_stdv[i] <= 0.0)
            continue;
        const double sim_val = std::max(double_min, sim_data[i]);
        const double exp_val = std::max(double_min, exp_data[i]);
        double value = std::log10(sim_val) - std::log10(exp_val);
        value *= exp_val * ln10 / exp_stdv[i];
        result += norm_fun(value);
    }

    return std::isfinite(result) ? result : double_max;
}

double LogMetric::computeFromArrays(std::vector<double> sim_data,
                                    std::vector<double> exp_data) const
{
    checkIntegrity(sim_data, exp_data);

    double result = 0.0;
    auto norm_fun = norm();
    for (size_t i = 0, sim_size = sim_data.size(); i < sim_size; ++i) {
        if (exp_data[i] < 0.0)
            continue;
        const double sim_val = std::max(double_min, sim_data[i]);
        const double exp_val = std::max(double_min, exp_data[i]);
        result += norm_fun(std::log10(sim_val) - std::log10(exp_val));
    }

    return std::isfinite(result) ? result : double_max;
}

// ----------------------- Relative difference ---------------------------

meanRelativeDifferenceMetric::meanRelativeDifferenceMetric() = default;

meanRelativeDifferenceMetric* meanRelativeDifferenceMetric::clone() const
{
    return copyMetric(*this);
}

double meanRelativeDifferenceMetric::computeFromArrays(std::vector<double> sim_data,
                                                       std::vector<double> exp_data) const
{
    checkIntegrity(sim_data, exp_data);

    double result = 0.0;
    auto norm_fun = norm();
    for (size_t i = 0, sim_size = sim_data.size(); i < sim_size; ++i) {
        if (exp_data[i] < 0.0)
            continue;
        const double sim_val = std::max(double_min, sim_data[i]);
        const double exp_val = std::max(double_min, exp_data[i]);
        result += norm_fun((exp_val - sim_val) / (exp_val + sim_val));
    }

    return std::isfinite(result) ? result : double_max;
}

// ----------------------- RQ4 metric ---------------------------

RQ4Metric::RQ4Metric() = default;

RQ4Metric* RQ4Metric::clone() const
{
    return copyMetric(*this);
}

double RQ4Metric::computeMetric(const SimDataPair& data_pair, bool use_weights) const
{
    if (use_weights)
        return Chi2Metric::computeMetric(data_pair, use_weights);

    // fetching data in RQ4 form
    // TODO weight data points correctly
    // https://jugit.fz-juelich.de/mlz/bornagain/-/issues/568
    throw std::runtime_error("\"rq4\" metric is temporary disabled.\n"
                             "Please choose another one.");

    const std::vector<double> sim_vec = data_pair.simulationResult().flatVector();
    const std::vector<double> exp_vec = data_pair.experimentalData().flatVector();

    return computeFromArrays(sim_vec, exp_vec);
}