//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Fitting/SimDataPair.cpp
//! @brief     Defines class SimDataPair.
//!
//! @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/SimDataPair.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/Scale.h"
#include "Base/Math/Numeric.h"
#include "Base/Util/Assert.h"
#include "Device/Data/Datafield.h"
#include "Device/Detector/IDetector.h"
#include "Sim/Simulation/ScatteringSimulation.h"
#include <utility>

namespace {

bool haveSameSizes(const IDetector& detector, const Datafield& data)
{
    if (data.rank() != 2)
        return false;

    for (size_t i = 0; i < 2; ++i)
        if (data.axis(i).size() != detector.axis(i).size())
            return false;

    return true;
}

//! Convert user data to Datafield object for later drawing in various axes units.
//! User data will be cropped to the ROI defined in the simulation, amplitudes in areas
//! corresponding to the masked areas of the detector will be set to zero.

Datafield repositionData(const ScatteringSimulation& simulation, const Datafield& data)
{
    auto frame = std::make_unique<Frame>(simulation.detector().clippedFrame());

    std::vector<double> values(frame->size(), 0.);
    std::vector<double> errors;
    if (data.hasErrorSigmas())
        errors = std::vector<double>(frame->size(), 0.);

    const IDetector& det = simulation.detector();
    std::vector<size_t> ai = det.activeIndices();
    if (frame->hasSameSizes(data.frame())) {
        for (unsigned long i : ai) {
            values[i] = data[i];
            if (data.hasErrorSigmas())
                errors[i] = data.errorSigmas()[i];
        }
    } else if (haveSameSizes(simulation.detector(), data)) {
        // experimental data has same shape as the detector, we have to copy the original
        // data to a smaller roi map
        for (unsigned long i : ai) {
            values[i] = data[det.roiToFullIndex(i)];
            if (data.hasErrorSigmas())
                errors[i] = data.errorSigmas()[det.roiToFullIndex(i)];
        }
    } else
        throw std::runtime_error(
            "FitObject::init_dataset: Detector and experimental data have different shape");

    return {*frame, values, errors};
}

} // namespace

//  ************************************************************************************************
//  class implementation
//  ************************************************************************************************

SimDataPair::SimDataPair(const SimulationWrapper& sim, const Datafield& raw_data,
                         const double weight)
    : m_simulation_builder(sim)
    , m_raw_data(raw_data.clone())
    , m_weight(weight)
{
    validate();
}

SimDataPair::SimDataPair(SimDataPair&& other) noexcept = default;

SimDataPair::~SimDataPair() = default;

void SimDataPair::execSimulation(const mumufit::Parameters& params)
{
    m_sim_data = std::make_unique<Datafield>(m_simulation_builder.simulate(params));

    ASSERT(!m_sim_data->empty());

    if (m_exp_data) {
        if (!m_exp_data->empty()) {
            // discard the simulation artifacts (if any)
            m_simulation_builder.discard();
            // TODO: why this early return?
            return;
        }
    }

    // TODO: why experimental data is replaced in this case?
    auto* const sim2d = dynamic_cast<ScatteringSimulation*>(m_simulation_builder.simulation.get());

    if (sim2d)
        m_exp_data = std::make_unique<Datafield>(repositionData(*sim2d, *m_raw_data));
    else
        m_exp_data = std::make_unique<Datafield>(*m_raw_data);
}

bool SimDataPair::containsUncertainties() const
{
    return m_raw_data->hasErrorSigmas();
}

Datafield SimDataPair::simulationResult() const
{
    ASSERT(m_sim_data);
    ASSERT(!m_sim_data->empty());
    return *m_sim_data;
}

Datafield SimDataPair::experimentalData() const
{
    ASSERT(m_exp_data);
    ASSERT(!m_exp_data->empty());
    return *m_exp_data;
}

std::vector<double> SimDataPair::simulation_array() const
{
    return simulationResult().flatVector();
}

std::vector<double> SimDataPair::experimental_array() const
{
    return experimentalData().flatVector();
}

std::vector<double> SimDataPair::uncertainties_array() const
{
    ASSERT(m_exp_data);
    return m_exp_data->errorSigmas();
}

//! Returns relative difference between simulation and experimental data.

Datafield SimDataPair::relativeDifference() const
{
    size_t N = m_sim_data->size();
    if (N == 0)
        throw std::runtime_error("Empty simulation data => won't compute relative difference");
    if (!m_exp_data || m_exp_data->size() != N)
        throw std::runtime_error("Different data shapes => won't compute relative difference");

    std::vector<double> data(N, 0.);
    for (size_t i = 0; i < N; ++i)
        data[i] = Numeric::relativeDifference((*m_sim_data)[i], (*m_exp_data)[i]);

    return {Datafield(m_sim_data->frame(), data)};
}

Datafield SimDataPair::absoluteDifference() const
{
    size_t N = m_sim_data->size();
    if (N == 0)
        throw std::runtime_error("Empty simulation data => won't compute absolute difference");
    if (!m_exp_data || m_exp_data->size() != N)
        throw std::runtime_error("Different data shapes => won't compute absolute difference");

    std::vector<double> data(N, 0.);
    for (size_t i = 0; i < N; ++i)
        data[i] = std::abs((*m_sim_data)[i] - (*m_exp_data)[i]);

    return {Datafield(m_sim_data->frame(), data)};
}

void SimDataPair::validate() const
{
    m_simulation_builder.check();

    if (!m_raw_data)
        throw std::runtime_error("Error in SimDataPair: passed experimental data array is empty");
}
