//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Scan/QzScan.cpp
//! @brief     Implements QzScan class.
//!
//! @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/Scan/QzScan.h"
#include "Base/Axis/FixedBinAxis.h"
#include "Base/Axis/PointwiseAxis.h"
#include "Device/Coord/CoordSystem1D.h"
#include "Device/Pol/PolFilter.h"
#include "Param/Distrib/ParameterSample.h"
#include "Param/Distrib/RangedDistributions.h"
#include "Resample/Element/SpecularElement.h"
#include "Sim/Scan/ScanResolution.h"
#include <algorithm>

QzScan::QzScan(IAxis* qs_nm)
    : m_qs(std::move(qs_nm))
    , m_resolution(scanEmptyResolution())
{
    std::vector<double> axis_values = m_qs->binCenters();
    if (!std::is_sorted(axis_values.begin(), axis_values.end()))
        throw std::runtime_error("Error in QzScan::checkInitialization: q-vector values shall "
                                 "be sorted in ascending order.");

    if (axis_values.front() < 0)
        throw std::runtime_error("Error in QzScan::checkInitialization: q-vector values are out "
                                 "of acceptable range.");
}

QzScan::QzScan(std::vector<double> qs_nm)
    : QzScan(new PointwiseAxis("qs", std::move(qs_nm)))
{
}

QzScan::QzScan(const IAxis& qs_nm)
    : QzScan(qs_nm.clone())
{
}

QzScan::QzScan(int nbins, double qz_min, double qz_max)
    : QzScan(new FixedBinAxis("qs", nbins, qz_min, qz_max))
{
}

QzScan::~QzScan() = default;

QzScan* QzScan::clone() const
{
    auto* result = new QzScan(*m_qs);
    result->setQResolution(*m_resolution);
    result->setOffset(m_offset);
    // TODO merge with same code in AlphaScan
    if (m_beamPolarization)
        result->m_beamPolarization.reset(new R3(*m_beamPolarization));
    if (m_polAnalyzer)
        result->m_polAnalyzer.reset(new PolFilter(*m_polAnalyzer));
    return result;
}

//! Generates simulation elements for specular simulations
std::vector<SpecularElement> QzScan::generateElements() const
{
    const auto samples = m_resolution->generateResolutionSamples(m_qs->binCenters());

    std::vector<double> qz;
    qz.reserve(numberOfElements());
    for (size_t i = 0, size_out = samples.size(); i < size_out; ++i)
        for (size_t j = 0, size_in = samples[i].size(); j < size_in; ++j)
            qz.push_back(samples[i][j].value);

    std::vector<SpecularElement> result;
    result.reserve(qz.size());
    for (size_t i = 0, size = qz.size(); i < size; ++i)
        result.emplace_back(
            SpecularElement::FromQzScan(-(qz[i] + m_offset) / 2.0, polMatrices(), qz[i] >= 0));

    return result;
}

std::vector<double> QzScan::footprint(size_t i, size_t n_elements) const
{
    if (i + n_elements > numberOfElements())
        throw std::runtime_error("Error in QzScan::footprint: given index exceeds the "
                                 "number of simulation elements");
    return std::vector<double>(n_elements, 1.0);
}

//! Returns the number of simulation elements
size_t QzScan::numberOfElements() const
{
    return m_qs->size() * m_resolution->nSamples();
}

CoordSystem1D* QzScan::scanCoordSystem() const
{
    return new WavenumberReflectometryCoords(coordinateAxis()->clone());
}

std::vector<double> QzScan::createIntensities(const std::vector<SpecularElement>& eles) const
{
    const size_t axis_size = m_qs->size();
    std::vector<double> result(axis_size, 0.0);

    const auto samples = m_resolution->generateResolutionSamples(m_qs->binCenters());

    size_t elem_pos = 0;
    for (size_t i = 0; i < axis_size; ++i) {
        double& current = result[i];
        for (size_t j = 0, size = samples[i].size(); j < size; ++j) {
            current += eles[elem_pos].intensity() * samples[i][j].weight;
            ++elem_pos;
        }
    }
    return result;
}

void QzScan::setQResolution(const IScanResolution& resolution)
{
    m_resolution.reset(resolution.clone());
}

void QzScan::setRelativeQResolution(const IRangedDistribution& distr, double rel_dev)
{
    std::unique_ptr<IScanResolution> resolution(scanRelativeResolution(distr, rel_dev));
    setQResolution(*resolution);
}

void QzScan::setRelativeQVectorResolution(const IRangedDistribution& distr,
                                          const std::vector<double>& rel_dev)
{
    std::unique_ptr<IScanResolution> resolution(scanRelativeVectorResolution(distr, rel_dev));
    setQResolution(*resolution);
}

void QzScan::setAbsoluteQResolution(const IRangedDistribution& distr, double std_dev)
{
    std::unique_ptr<IScanResolution> resolution(scanAbsoluteResolution(distr, std_dev));
    setQResolution(*resolution);
}

void QzScan::setAbsoluteQVectorResolution(const IRangedDistribution& distr,
                                          const std::vector<double>& std_dev)
{
    std::unique_ptr<IScanResolution> resolution(scanAbsoluteVectorResolution(distr, std_dev));
    setQResolution(*resolution);
}