//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Device/Detector/Detector.cpp
//! @brief     Implements common detector interface.
//!
//! @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 "Device/Detector/Detector.h"
#include "Base/Axis/Bin.h"
#include "Base/Axis/Scale.h"
#include "Base/Util/Assert.h"
#include "Device/Data/Datafield.h"
#include "Device/Mask/InfinitePlane.h"
#include "Device/Mask/MaskStack.h"
#include "Device/Pol/PolFilter.h"
#include "Device/Resolution/ConvolutionDetectorResolution.h"

//... struct RoiOfAxis

//! Keeps RegionOfInterest (ROI) data of one axis.
struct RoiOfAxis {
    double lower;
    double upper;

    // Precomputed values, to speed up repeated calculations.
    size_t lowerIndex;   //!< index corresponding to 'lower'
    size_t upperIndex;   //!< index corresponding to 'upper'
    size_t roiSize;      //!< number of bins on axis of ROI
    size_t detectorSize; //!< number of bins on axis of detector

    RoiOfAxis(const Scale& axis, double lower, double upper);

    std::pair<double, double> bounds() const;
};

RoiOfAxis::RoiOfAxis(const Scale& axis, double _lower, double _upper)
    : lower(_lower)
    , upper(_upper)
{
    ASSERT(lower < upper);
    detectorSize = axis.size();
    lowerIndex = axis.closestIndex(lower);
    upperIndex = axis.closestIndex(upper);
    // suppress tiny bins that are most likely due to floating-point inaccuracy
    if (axis.bin(lowerIndex).binSize() < 1e-12 * axis.span() / axis.size()) {
        ASSERT(lowerIndex < axis.size() - 1);
        ++lowerIndex;
    }
    if (axis.bin(upperIndex).binSize() < 1e-12 * axis.span() / axis.size()) {
        ASSERT(upperIndex > 0);
        --upperIndex;
    }
    roiSize = upperIndex - lowerIndex + 1;
}

std::pair<double, double> RoiOfAxis::bounds() const
{
    return {lower, upper};
}


//... class Detector

Detector::Detector(Frame* frame)
    : m_frame(frame)
    , m_mask(new MaskStack)
{
    ASSERT(frame);
    ASSERT(frame->rank() == 2);
}

Detector::~Detector() = default;

void Detector::setAnalyzer(const R3& direction, double transmission_selected,
                           double transmission_suppressed)
{
    m_pol_analyzer =
        std::make_unique<PolFilter>(direction, transmission_selected, transmission_suppressed);
}

void Detector::setResolutionFunction(const ResolutionFunction2D& resFunc)
{
    ConvolutionDetectorResolution convFunc(resFunc);
    setDetectorResolution(convFunc);
}

const Frame& Detector::frame() const
{
    ASSERT(m_frame);
    return *m_frame;
}

const MaskStack* Detector::detectorMask() const
{
    return m_mask.get();
}

void Detector::addMask(const Shape2D& shape, bool mask_value)
{
    m_mask->pushMask(shape, mask_value);
}

void Detector::maskAll()
{
    addMask(InfinitePlane(), true);
}

const Scale& Detector::axis(size_t i) const
{
    ASSERT(i < 2);
    return m_frame->axis(i);
}

void Detector::setRegionOfInterest(double xlow, double ylow, double xup, double yup)
{
    m_explicitROI.clear();
    m_explicitROI.emplace_back(axis(0), xlow, xup);
    m_explicitROI.emplace_back(axis(1), ylow, yup);
}

Frame Detector::clippedFrame() const
{
    ASSERT(m_frame);
    ASSERT(m_frame->rank() == 2);
    return {new Scale(m_frame->axis(0).clipped(regionOfInterestBounds(0))),
            new Scale(m_frame->axis(1).clipped(regionOfInterestBounds(1)))};
}

size_t Detector::roiToFullIndex(const size_t i) const
{
    if (m_explicitROI.size() != 2)
        return i;

    const auto& x = m_explicitROI[0];
    const auto& y = m_explicitROI[1];

    const size_t globalIndex0 = y.lowerIndex * x.detectorSize + x.lowerIndex;
    const size_t xcoord = i % x.roiSize;
    const size_t ycoord = i / x.roiSize;
    return globalIndex0 + xcoord + ycoord * x.detectorSize;
}

std::vector<size_t> Detector::activeIndices() const
{
    std::vector<size_t> result;
    for (size_t i = 0; i < sizeOfRegionOfInterest(); ++i)
        if (!detectorMask() || !detectorMask()->isMasked(roiToFullIndex(i), frame()))
            result.push_back(i);

    return result;
}

size_t Detector::totalSize() const
{
    return m_frame->axis(0).size() * m_frame->axis(1).size();
}

void Detector::setDetectorResolution(const DetectorResolution& detector_resolution)
{
    m_resolution.reset(detector_resolution.clone());
}

void Detector::applyDetectorResolution(Datafield* df) const
{
    ASSERT(df);

    if (!m_resolution)
        return;
    m_resolution->execDetectorResolution(df);

    if (detectorMask())
        // set amplitude in masked areas to zero
        for (size_t i = 0; i < totalSize(); ++i)
            if (detectorMask()->isMasked(i, frame()))
                df->setAt(i, 0.);
}

bool Detector::hasExplicitRegionOfInterest() const
{
    return m_explicitROI.size() == 2;
}

size_t Detector::sizeOfRegionOfInterest() const
{
    if (hasExplicitRegionOfInterest())
        return m_explicitROI[0].roiSize * m_explicitROI[1].roiSize;
    return totalSize();
}

std::pair<double, double> Detector::regionOfInterestBounds(size_t iAxis) const
{
    ASSERT(iAxis < 2);
    if (iAxis < m_explicitROI.size())
        return m_explicitROI[iAxis].bounds();
    return m_frame->axis(iAxis).bounds();
}

Datafield Detector::createDetectorMap() const
{
    std::vector<const Scale*> axes;
    axes.reserve(2);
    for (size_t i = 0; i < 2; ++i)
        axes.push_back(new Scale(axis(i).clipped(regionOfInterestBounds(i))));
    return {Frame(axes)};
}

SpinMatrix Detector::analyzerMatrix() const
{
    return m_pol_analyzer ? m_pol_analyzer->matrix() : SpinMatrix::One();
}

void Detector::copySettings(Detector* other) const
{
    other->m_explicitROI = m_explicitROI;
    other->m_mask.reset(m_mask->clone());
    if (m_resolution)
        other->m_resolution.reset(m_resolution->clone());
    if (m_pol_analyzer)
        other->setAnalyzer(m_pol_analyzer->direction(), m_pol_analyzer->transmissionParallel(),
                           m_pol_analyzer->transmissionAntiparallel());
}

size_t Detector::globalIndex(size_t x, size_t y) const
{
    return y * axis(0).size() + x;
}

size_t Detector::axisBinIndex(size_t i, size_t k_axis) const
{
    return m_frame->projectedIndex(i, k_axis);
}
