//  ************************************************************************************************
//
//  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 "Resample/Element/ScanElement.h"
#include "Sim/Background/IBackground.h"
#include "Sim/Computation/DWBAComputation.h"
#include "Sim/Scan/PhysicalScan.h"

OffspecSimulation::OffspecSimulation(const PhysicalScan& scan, const Sample& 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:

void OffspecSimulation::initScanElementVector()
{
    m_eles = m_scan->generateElements();
}

void OffspecSimulation::runComputation(const ReSample& re_sample, size_t i, double weight)
{
    const size_t Na = m_detector->totalSize();
    size_t j = i / Na; // index in unwrapped scan (including distribution sampling)
    size_t k = i % Na; // index in detector

    ASSERT(m_cache.size() / Na == m_eles.size());

    ScanElement& scan_ele = *(m_eles.begin() + static_cast<long>(j));
    const double alpha_i = scan_ele.alpha();
    const double lambda_i = scan_ele.lambda();
    const double phi_i = scan_ele.phi();
    const bool isSpecular = k == m_detector->indexOfSpecular(alpha_i, phi_i);

    DiffuseElement diff_ele(lambda_i, alpha_i, phi_i, m_pixels[k],
                            m_scan->polarizerMatrixAt(scan_ele.i_out()),
                            m_detector->analyzer().matrix(), isSpecular);

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

    double sin_alpha_i = std::abs(std::sin(alpha_i));
    if (sin_alpha_i == 0.0) {
        intensity = 0;
    } else {
        const double solid_angle = diff_ele.solidAngle();
        intensity *= scan_ele.beamIntensity() * scan_ele.footprint() * solid_angle / sin_alpha_i;
    }

    m_cache[i] += intensity * weight * scan_ele.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->nElements();
}

Datafield OffspecSimulation::packResult() const
{
    const size_t ns = m_scan->nScan();
    const size_t nslong = m_scan->nElements();

    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.);

    for (size_t j = 0; j < nslong; ++j) {
        for (size_t ia = 0; ia < nalp; ++ia) {
            double val = 0;

            // sum over detector phi to the distribution point
            for (size_t ip = 0; ip < nphi; ++ip)
                val += m_cache[j * Ndet + ia * nphi + ip];

            // accumulate distribution points to the scan point
            const ScanElement& ele = m_eles.at(j);
            const size_t j_real = ele.i_out();
            out[ia * ns + j_real] += val;
        }
    }

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

    if (background())
        for (double& o : out)
            o = background()->addBackground(o);

    std::vector<const Scale*> axes{m_scan->coordinateAxis()->clone(), m_detector->axis(1).clone()};
    return {axes, out};
}
