//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Simulation/OffspecSimulation.cpp
//! @brief     Implements class OffspecSimulation.
//!
//! @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/OffspecSimulation.h"
#include "Base/Axis/Scale.h"
#include "Base/Pixel/IPixel.h"
#include "Base/Progress/ProgressHandler.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/IFootprint.h"
#include "Device/Coord/CoordSystem2D.h"
#include "Device/Data/Datafield.h"
#include "Device/Detector/OffspecDetector.h"
#include "Device/Histo/SimulationResult.h"
#include "Param/Distrib/DistributionHandler.h"
#include "Param/Distrib/Distributions.h"
#include "Resample/Element/DiffuseElement.h"
#include "Sim/Background/IBackground.h"
#include "Sim/Computation/DWBAComputation.h"
#include "Sim/Scan/AlphaScan.h"

OffspecSimulation::OffspecSimulation(const IBeamScan& scan, const MultiLayer& sample,
                                     const OffspecDetector& detector)
    : ISimulation(sample)
    , m_scan(scan.clone())
    , m_detector(detector.clone())
{
}

OffspecSimulation::~OffspecSimulation() = default;

std::vector<const INode*> OffspecSimulation::nodeChildren() const
{
    std::vector<const INode*> result = ISimulation::nodeChildren();
    result.push_back(m_scan.get());
    if (m_detector)
        result.push_back(m_detector.get());
    return result;
}

void OffspecSimulation::prepareSimulation()
{
    m_pixels.reserve(m_detector->totalSize());
    for (size_t i = 0; i < m_detector->totalSize(); ++i)
        m_pixels.emplace_back(m_detector->createPixel(i));
}

const ICoordSystem* OffspecSimulation::simCoordSystem() const
{
    return new OffspecCoords(
        std::vector<const Scale*>{m_scan->coordinateAxis()->clone(), m_detector->axis(1).clone()});
}

//... Overridden executors:

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

        switch (distribution.whichParameter()) {
        case ParameterDistribution::BeamWavelength:
            distributionHandler().defineCallbackForDistribution(
                &distribution, [&](double d) { m_scan->setWavelength(d); });
            break;
        default:
            ASSERT(false);
        }
    }
}

void OffspecSimulation::runComputation(const ReSample& re_sample, size_t i, double weight)
{
    if (auto alpha_scan = dynamic_cast<AlphaScan*>(m_scan.get()))
        if (alpha_scan->wavelengthDistribution() || alpha_scan->angleDistribution())
            throw std::runtime_error("Offspecular simulation supports neither alpha nor lambda "
                                     "distributions (issue #752).");

    if (m_cache.empty())
        m_cache.resize(nElements(), 0.0);

    const size_t Na = m_detector->totalSize();
    size_t j = i / Na;
    size_t k = i % Na;

    const double alpha = m_scan->coordinateAxis()->binCenter(j);
    const double phi = 0;
    const bool isSpecular = k == m_detector->indexOfSpecular(alpha, phi);

    DiffuseElement ele(m_scan->wavelength(), alpha, phi, m_pixels[k], m_scan->polarizerMatrix(),
                       m_scan->analyzerMatrix(), isSpecular);

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

    if (const auto* footprint = m_scan->footprint())
        intensity *= footprint->calculate(alpha);

    double sin_alpha_i = std::abs(std::sin(alpha));
    if (sin_alpha_i == 0.0) {
        intensity = 0;
    } else {
        const double solid_angle = ele.solidAngle();
        intensity *= m_scan->intensity() * solid_angle / sin_alpha_i;
    }

    m_cache[i] += intensity * weight;

    progress().incrementDone(1);
}

//... Overridden getters:

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

size_t OffspecSimulation::nElements() const
{
    return m_detector->totalSize() * m_scan->coordinateAxis()->size();
}

SimulationResult OffspecSimulation::packResult()
{

    // update intensity map
    Datafield intensity_map({m_scan->coordinateAxis()->clone(), m_detector->axis(1).clone()});
    intensity_map.setAllTo(0.);

    size_t ny = m_detector->axis(1).size();

    // Apply detector resolution and transfer detector image
    for (size_t j = 0; j < m_scan->coordinateAxis()->size(); ++j) {
        Datafield detector_image({m_detector->axis(0).clone(), m_detector->axis(1).clone()});
        size_t N = detector_image.size();
        for (size_t i = 0; i < N; ++i)
            detector_image[i] = m_cache[j * N + i];
        // TODO restore resolution m_detector->applyDetectorResolution(&detector_image);
        for (size_t i = 0; i < N; ++i)
            intensity_map[j * ny + i % ny] += detector_image[i];

        if (background())
            for (size_t i = 0; i < N; ++i)
                intensity_map[j * ny + i % ny] =
                    background()->addBackground(intensity_map[j * ny + i % ny]);
    }

    return {intensity_map, simCoordSystem()};
}
