//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Simulation/ScatteringSimulation.cpp
//! @brief     Implements interface ISimulation.
//!
//! @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/Simulation/ScatteringSimulation.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/Pixel.h"
#include "Base/Axis/Scale.h"
#include "Base/Progress/ProgressHandler.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/Beam.h"
#include "Device/Beam/IFootprint.h"
#include "Device/Data/Datafield.h"
#include "Device/Detector/IDetector.h"
#include "Param/Distrib/DistributionHandler.h"
#include "Resample/Element/DiffuseElement.h"
#include "Sample/Multilayer/Sample.h"
#include "Sim/Background/IBackground.h"
#include "Sim/Computation/DWBAComputation.h"

ScatteringSimulation::ScatteringSimulation(const Beam& beam, const Sample& sample,
                                           const IDetector& detector)
    : ISimulation(sample)
    , m_beam(beam.clone())
    , m_detector(detector.clone())
{
}

ScatteringSimulation::~ScatteringSimulation() = default;

void ScatteringSimulation::addParameterDistribution(
    ParameterDistribution::WhichParameter whichParameter, const IDistribution1D& distribution)
{
    ParameterDistribution par_distr(whichParameter, distribution);
    distributionHandler().addDistribution(par_distr);
}

//... Overridden executors:

//! init callbacks for setting the parameter values
void ScatteringSimulation::initDistributionHandler()
{
    for (const auto& distribution : distributionHandler().paramDistributions()) {

        switch (distribution.whichParameter()) {
        case ParameterDistribution::BeamAzimuthalAngle:
            distributionHandler().defineCallbackForDistribution(
                &distribution, [&](double d) { m_beam->setAzimuthalAngle(d); });
            break;
        case ParameterDistribution::BeamInclinationAngle:
            distributionHandler().defineCallbackForDistribution(
                &distribution, [&](double d) { m_beam->setInclination(d); });
            break;
        case ParameterDistribution::BeamWavelength:
            distributionHandler().defineCallbackForDistribution(
                &distribution, [&](double d) { m_beam->setWavelength(d); });
            break;
        default:
            ASSERT_NEVER;
        }
    }
}

void ScatteringSimulation::prepareSimulation()
{
    sample()->checkMaterials(beam().wavelength());
    m_active_indices = m_detector->activeIndices();
    m_pixels.reserve(m_active_indices.size());
    for (size_t i : m_active_indices)
        m_pixels.push_back(m_detector->createPixel(m_detector->roiToFullIndex(i)));
}

void ScatteringSimulation::runComputation(const ReSample& re_sample, size_t i, double weight)
{
    if (m_cache.empty())
        m_cache.resize(nElements(), 0.0);

    const bool isSpecular = m_active_indices[i] == m_detector->indexOfSpecular(beam());

    DiffuseElement ele(beam().wavelength(), beam().alpha_i(), beam().phi_i(), m_pixels[i],
                       beam().polMatrix(), m_detector->analyzer().matrix(), isSpecular);

    double intensity = Compute::scattered_and_reflected(re_sample, options(), ele);

    if (const auto* footprint = beam().footprint())
        intensity *= footprint->calculate(beam().alpha_i());

    double sin_alpha_i = std::abs(std::sin(beam().alpha_i()));
    if (sin_alpha_i == 0.0) {
        intensity = 0;
    } else {
        const double solid_angle = m_pixels[i]->solidAngle();
        intensity *= m_beam->intensity() * solid_angle / sin_alpha_i;
    }

    m_cache[i] += intensity * weight;

    progress().incrementDone(1);
}

//... Overridden getters:

bool ScatteringSimulation::force_polarized() const
{
    return m_detector->analyzer().BlochVector() != R3{};
}

size_t ScatteringSimulation::nElements() const
{
    return m_active_indices.size();
}

Datafield ScatteringSimulation::packResult() const
{
    Datafield result(m_detector->createDetectorMap());
    for (size_t i = 0; i < m_active_indices.size(); ++i)
        result.setAt(m_active_indices[i], m_cache[i]);

    m_detector->applyDetectorResolution(&result);

    if (background())
        for (size_t i : m_active_indices)
            result[i] = background()->addBackground(result[i]);

    return result;
}
