//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Scan/QzScan.cpp
//! @brief     Implements class QzScan.
//!
//! @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/Pol/PolFilter.h"
#include "Param/Distrib/Distributions.h"
#include "Param/Distrib/ParameterSample.h"
#include "Resample/Element/SpecularElement.h"
#include <algorithm> // is_sorted
#include <numbers>

using std::numbers::pi;

QzScan::QzScan(Scale* qs_nm)
    : BeamScan(qs_nm)
{
    std::vector<double> axis_values = m_axis->binCenters();
    if (!std::is_sorted(axis_values.begin(), axis_values.end()))
        throw std::runtime_error("QzScan: q-vector values are not "
                                 "sorted in ascending order.");
    if (axis_values.front() < 0)
        throw std::runtime_error("QzScan: negative q-values.");

    m_beams.clear();
    for (size_t i = 0; i < nScan(); i++) {
        // Qz scan is internally understood as wavelength scan
        double lambda = 4 * pi * std::sin(defaultInclination) / m_axis->binCenter(i);
        auto* beam = new Beam(defaultIntensity, lambda, defaultInclination);
        // Since the grazing geometry is not clear in q-space, the footprint should not be used
        beam->setFootprint(nullptr);
        m_beams.push_back(beam);
    }
}

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

QzScan::QzScan(std::vector<double> qs_nm)
    : QzScan(newListScan("q (1/nm)", std::move(qs_nm)))
{
}

QzScan::QzScan(int nbins, double qz_min, double qz_max)
    : QzScan(newEquiScan("q (1/nm)", nbins, qz_min, qz_max))
{
}

QzScan::~QzScan() = default;

QzScan* QzScan::clone() const
{
    auto* result = new QzScan(*m_axis);
    copyBeamScan(result);
    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);
    return result;
}

std::vector<const INode*> QzScan::nodeChildren() const
{
    std::vector<const INode*> result;
    for (const INode* n : BeamScan::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->binCenter(i);
        if (m_qz_distrib) {
            const std::vector<ParameterSample> qzDistrib = m_qz_distrib->distributionSamples();
            for (auto j : qzDistrib) {
                double qz = q0;
                ASSERT(!m_resol_width.empty());
                if (m_relative_resolution)
                    qz += q0 * m_resol_width[0] * j.value;
                else if (m_resol_width.size() > 1)
                    qz += m_resol_width[i] * j.value;
                else
                    qz += m_resol_width[0] * j.value;
                result.emplace_back(i, qz >= 0, j.weight, intensityAt(i), 1., polarizerMatrixAt(i),
                                    analyzerMatrix(), R3(0, 0, -(qz + m_offset) / 2));
            }
        } else {
            result.emplace_back(i, q0 >= 0, 1., intensityAt(i), 1., polarizerMatrixAt(i),
                                analyzerMatrix(), R3(0, 0, -(q0 + m_offset) / 2));
        }
    }
    return result;
}

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

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;
}
