//  ************************************************************************************************
//
//  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/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/IFootprint.h"
#include "Device/Data/Datafield.h"
#include "Device/Detector/OffspecDetector.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/PhysicalScan.h"

OffspecSimulation::OffspecSimulation(const PhysicalScan& 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.push_back(m_detector->createPixel(i));
}

//... 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_NEVER;
        }
    }
}

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

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

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

    const double alpha_i = m_scan->inclinationAt(j);
    const double phi_i = 0;
    const bool isSpecular = k == m_detector->indexOfSpecular(alpha_i, phi_i);

    DiffuseElement ele(m_scan->wavelengthAt(j), alpha_i, phi_i, m_pixels[k],
                       m_scan->polarizerMatrixAt(j), m_detector->analyzer().matrix(), isSpecular);

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

    if (const auto* footprint = m_scan->footprintAt(j))
        intensity *= footprint->calculate(alpha_i);

    double sin_alpha_i = std::abs(std::sin(alpha_i));
    if (sin_alpha_i == 0.0) {
        intensity = 0;
    } else {
        const double solid_angle = ele.solidAngle();
        intensity *= m_scan->intensityAt(j) * 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->nScan();
}

Datafield OffspecSimulation::packResult()
{
    const size_t ns = m_scan->nScan();
    const size_t nphi = m_detector->axis(0).size();
    const size_t nalp = m_detector->axis(1).size();
    const size_t Ndet = nalp * nphi;
    std::vector<double> out(ns * nalp, 0.);

    // Apply detector resolution and transfer detector image
    for (size_t j = 0; j < ns; ++j) {

        // TODO restore resolution m_detector->applyDetectorResolution(&detector_image);

        for (size_t ia = 0; ia < nalp; ++ia) {
            double val = 0;
            for (size_t ip = 0; ip < nphi; ++ip)
                val += m_cache[j * Ndet + ia * nphi + ip];
            if (background())
                val = background()->addBackground(val);
            out[ia * ns + j] = val;
        }
    }

    return {{m_scan->coordinateAxis()->clone(), m_detector->axis(1).clone()}, out};
}
