//  ************************************************************************************************
//
//  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/MakeScale.h"
#include "Base/Axis/Scale.h"
#include "Base/Util/Assert.h"
#include "Device/Coord/CoordSystem1D.h"
#include "Device/Pol/PolFilter.h"
#include "Param/Distrib/Distributions.h"
#include "Param/Distrib/ParameterSample.h"
#include "Resample/Element/SpecularElement.h"
#include <algorithm>

QzScan::QzScan(Scale* qs_nm)
    : IBeamScan(qs_nm, std::nan(""))
{
    std::vector<double> axis_values = m_axis->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(newListScan("qs", std::move(qs_nm)))
{
}

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

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

QzScan::~QzScan() = default;

QzScan* QzScan::clone() const
{
    auto* result = new QzScan(*m_axis);
    result->setIntensity(intensity());
    if (m_qz_distrib) {
        result->m_qz_distrib.reset(m_qz_distrib->clone());
        result->m_resol_width = m_resol_width;
        result->m_relative_resolution = m_relative_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;
}

std::vector<const INode*> QzScan::nodeChildren() const
{
    std::vector<const INode*> result;
    for (const INode* n : IBeamScan::nodeChildren())
        result << n;
    if (m_qz_distrib)
        result << m_qz_distrib.get();
    return result;
}

//! Generates simulation elements for specular simulations
std::vector<SpecularElement> QzScan::generateElements() const
{
    std::vector<SpecularElement> result;
    result.reserve(nDistributionSamples());
    for (size_t i = 0; i < m_axis->size(); ++i) {
        const double q0 = m_axis->binCenters()[i];
        if (m_qz_distrib) {
            const std::vector<ParameterSample> qzDistrib = m_qz_distrib->distributionSamples();
            for (size_t j = 0; j < qzDistrib.size(); ++j) {
                double qz = q0;
                ASSERT(m_resol_width.size() > 0);
                if (m_relative_resolution)
                    qz += q0 * m_resol_width[0] * qzDistrib[j].value;
                else if (m_resol_width.size() > 1)
                    qz += m_resol_width[i] * qzDistrib[j].value;
                else
                    qz += m_resol_width[0] * qzDistrib[j].value;
                result.emplace_back(
                    SpecularElement::FromQzScan(i, qzDistrib[j].weight, -(qz + m_offset) / 2,
                                                polarizerMatrix(), analyzerMatrix(), qz >= 0));
            }
        } else {
            result.emplace_back(SpecularElement::FromQzScan(
                i, 1., -(q0 + m_offset) / 2, polarizerMatrix(), analyzerMatrix(), q0 >= 0));
        }
    }
    return result;
}

//! Returns the number of simulation elements
size_t QzScan::nDistributionSamples() const
{
    return m_qz_distrib ? m_qz_distrib->nSamples() : 1;
}

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

void QzScan::setRelativeQResolution(const IDistribution1D& distr, double rel_dev)
{
    m_qz_distrib.reset(distr.clone());
    m_relative_resolution = true;
    m_resol_width = {rel_dev};
}

void QzScan::setAbsoluteQResolution(const IDistribution1D& distr, double std_dev)
{
    m_qz_distrib.reset(distr.clone());
    m_resol_width = {std_dev};
}

void QzScan::setVectorResolution(const IDistribution1D& distr, const std::vector<double>& std_devs)
{
    m_qz_distrib.reset(distr.clone());
    ASSERT(std_devs.size() > 1);
    m_resol_width = std_devs;
}
