//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Simulation/ScatteringSimulation.cpp
//! @brief     Implements interface Simulation.
//!
//! @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/Distrib/DistributionHandler.h"
#include "Base/Progress/ProgressHandler.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/Beam.h"
#include "Device/Beam/Footprint.h"
#include "Device/Data/Datafield.h"
#include "Device/Detector/Detector.h"
#include "Sample/Aggregate/Dense2D.h"
#include "Sample/Element/DiffuseElement.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Option/SimulationOptions.h"
#include "Sample/Sample/Sample.h"
#include "Sim/Background/Background.h"
#include "Sim/Evaluate/EvaluateScattering.h"
#include "Sim/Evaluate/EvaluateSpecularGisas.h"

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

ScatteringSimulation::~ScatteringSimulation() = default;

void ScatteringSimulation::addParameterDistribution(
    ParameterDistribution::WhichParameter whichParameter, const Distribution1D& 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::BeamGrazingAngle:
            distributionHandler().defineCallbackForDistribution(
                &distribution, [&](double d) { m_beam->setGrazingAngle(d); });
            break;
        case ParameterDistribution::BeamWavelength:
            distributionHandler().defineCallbackForDistribution(
                &distribution, [&](double d) { m_beam->setWavelength(d); });
            break;
        default:
            ASSERT_NEVER;
        }
    }
}

void ScatteringSimulation::prepareSimulation()
{
    // Dense2D's PY packing limit is a scattering-domain restriction. Scene export must still be
    // able to visualize overpacked Dense2D samples, so this check stays out of Sample::validate().
    const auto validate_dense_structs = [](const auto& structs) {
        for (const auto& s : structs)
            if (const auto* dense = dynamic_cast<const Dense2D*>(s.get()))
                dense->validate();
    };
    for (const auto& layer : sample()->unwrapped()) {
        validate_dense_structs(layer->structsAtBottom());
        validate_dense_structs(layer->structsAtTop());
    }

    sample()->checkRefractiveIndex(m_beam->wavelength());
    m_active_indices = m_detector->activeIndices();
    m_pixels.reserve(m_active_indices.size());
    for (size_t i : m_active_indices)
        m_pixels.emplace_back(m_detector->createPixel(m_detector->roiToFullIndex(i)));
}

//! Computes reflection and scattering towards m_pixel[i], and stores it in m_cache[i].

void ScatteringSimulation::runComputation(const ReSample& re_sample, size_t i, double weight)
{
    ASSERT(m_cache.size() == nElements());

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

    DiffuseElement ele(m_beam->wavelength(), m_beam->alpha_i(), m_beam->phi_i(), m_pixels[i].get(),
                       m_beam->polMatrix(), m_detector->analyzerMatrix());

    double intensity = Evaluate::scatteredContribution(re_sample, ele, options(), doItPolarized());

    if (isSpecular && options().includeSpecular())
        intensity += Evaluate::gisasSpecularContribution(re_sample, ele, doItPolarized());

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

    double sin_alpha_i = std::abs(std::sin(m_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::forcePolarized() const
{
    return m_detector->analyzer() != nullptr;
}

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()->perturbed(result[i]);

    return result;
}
