//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Device/Detector/IDetector.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/IDetector.h"
#include "Base/Axis/Scale.h"
#include "Base/Util/Assert.h"
#include "Device/Mask/InfinitePlane.h"
#include "Device/Mask/MaskStack.h"
#include "Device/Resolution/ConvolutionDetectorResolution.h"
#include <iostream>

//... 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 IDetector

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

IDetector::IDetector(const IDetector& other)
    : INode(other)
    , m_explicitROI(other.m_explicitROI)
    , m_frame(other.m_frame->clone())
    , m_pol_analyzer(other.m_pol_analyzer)
    , m_resolution(other.m_resolution ? other.m_resolution->clone() : nullptr)
    , m_mask(other.m_mask->clone())
{
}

IDetector::~IDetector() = default;

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

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

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

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

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

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

Frame IDetector::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)))};
}

void IDetector::setAnalyzer(const R3& Bloch_vector, double mean_transmission)
{
    m_pol_analyzer = PolFilter(Bloch_vector, mean_transmission);
}

void IDetector::setDetectorResolution(const IDetectorResolution& detector_resolution)
{
    m_resolution.reset(detector_resolution.clone());
}

// TODO: pass dimension-independent argument to this function
void IDetector::setResolutionFunction(const IResolutionFunction2D& resFunc)
{
    ConvolutionDetectorResolution convFunc(resFunc);
    setDetectorResolution(convFunc);
}

void IDetector::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.);
}

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

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

std::vector<const INode*> IDetector::nodeChildren() const
{
    return std::vector<const INode*>() << &m_pol_analyzer << m_resolution.get();
}

size_t IDetector::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;
}

void IDetector::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);
}

std::vector<size_t> IDetector::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;
}

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

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

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

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