//  ************************************************************************************************
//
//  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/Distrib/Distributions.h"
#include "Base/Distrib/ParameterSample.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/Beam.h"
#include "Device/Pol/PolFilter.h"
#include "Sample/Element/ScanElement.h"
#include <algorithm> // is_sorted
#include <numbers>

using std::numbers::pi;

QzScan::QzScan(const 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(defaultGrazingAngle) / m_axis->binCenter(i);
        auto beam = std::make_unique<Beam>(defaultIntensity, lambda, defaultGrazingAngle);
        // Since the grazing geometry is not clear in q-space, the footprint should not be used
        beam->setFootprint(nullptr);
        m_beams.emplace_back(std::move(beam));
    }
}

QzScan::QzScan(std::vector<double> qs_nm)
    : QzScan(*std::unique_ptr<Scale>(ListScan("q_z (1/nm)", std::move(qs_nm))))
{
}

QzScan::QzScan(int nbins, double qz_min, double qz_max)
    : QzScan(*std::unique_ptr<Scale>(EquiScan("q_z (1/nm)", nbins, qz_min, qz_max)))
{
}

QzScan::~QzScan() = default;

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

void QzScan::setAbsoluteQResolution(const Distribution1D& distr, double std_dev)
{
    m_qz_distrib.reset(distr.clone());
    m_relative_resolution = false;
    m_resol_width_factor = {std_dev};
}

void QzScan::setVectorResolution(const Distribution1D& distr, const std::vector<double>& std_devs)
{
    m_qz_distrib.reset(distr.clone());
    m_relative_resolution = false;
    ASSERT(std_devs.size() > 1);
    m_resol_width_factor = std_devs;
}
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_factor = m_resol_width_factor;
        result->m_relative_resolution = m_relative_resolution;
    }
    result->setOffset(m_offset);
    return result;
}

std::vector<ScanElement> QzScan::generateElements() const
{
    std::vector<ScanElement> 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_factor.empty());
                if (m_relative_resolution)
                    qz += q0 * m_resol_width_factor[0] * j.value;
                else if (m_resol_width_factor.size() > 1)
                    qz += m_resol_width_factor[i] * j.value;
                else
                    qz += m_resol_width_factor[0] * j.value;
                result.emplace_back(i, qz >= 0, j.weight, intensityAt(i), 1., polarizerMatrixAt(i),
                                    analyzerMatrix(), 0, 0, 0, R3(0, 0, -(qz + m_offset) / 2));
            }
        } else {
            result.emplace_back(i, q0 >= 0, 1., intensityAt(i), 1., polarizerMatrixAt(i),
                                analyzerMatrix(), 0, 0, 0, R3(0, 0, -(q0 + m_offset) / 2));
        }
    }
    return result;
}

size_t QzScan::nDistributionSamples() const
{
    return m_qz_distrib ? m_qz_distrib->nSamples() : 1;
}
