//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Scan/BeamScan.cpp
//! @brief     Implements interface BeamScan.
//!
//! @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/BeamScan.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/Scale.h"
#include "Base/Distrib/Distributions.h"
#include "Base/Distrib/ParameterSample.h"
#include "Base/Math/Numeric.h"
#include "Base/Spin/SpinMatrix.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/Beam.h"
#include "Device/Pol/PolFilter.h"
#include "Resample/Slice/ComputeKz.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/Slice.h"

BeamScan::BeamScan(const Scale& axis)
    : m_axis(axis.clone())
{
}

BeamScan::~BeamScan() = default;

void BeamScan::setIntensity(double intensity)
{
    for (auto& b : m_beams)
        b->setIntensity(intensity);
}

void BeamScan::setPolarization(const R3& bloch_vector)
{
    for (auto& b : m_beams)
        b->setPolarization(bloch_vector);
}

void BeamScan::setAnalyzer(const R3& direction, double transmission_selected,
                           double transmission_suppressed)
{
    m_pol_analyzer =
        std::make_unique<PolFilter>(direction, transmission_selected, transmission_suppressed);
}

double BeamScan::commonIntensity() const
{
    if (!isCommonIntensity())
        throw std::runtime_error("Intensity changes during scan. "
                                 "Use 'intensityAt(i)' instead.");
    return m_beams.front()->intensity();
}

double BeamScan::intensityAt(size_t i) const
{
    return m_beams[i]->intensity();
}

const R3& BeamScan::commonPolarization() const
{
    if (!isCommonPolarization())
        throw std::runtime_error("Polarization changes during scan. "
                                 "Use 'polarizationAt(i)' instead.");
    return m_beams.front()->polVector();
}

R3 BeamScan::polarizationAt(size_t i) const
{
    return m_beams[i]->polVector();
}

SpinMatrix BeamScan::polarizerMatrixAt(size_t i) const
{
    return m_beams[i]->polMatrix();
}

SpinMatrix BeamScan::analyzerMatrix() const
{
    return m_pol_analyzer ? m_pol_analyzer->matrix() : SpinMatrix::One();
}

size_t BeamScan::nScan() const
{
    return m_axis->size();
}

std::vector<complex_t> BeamScan::produceKz(const std::vector<std::unique_ptr<const Slice>>& slices,
                                           const R3& k) const
{
    return Compute::kzFromRefractivity(slices, k);
}

void BeamScan::copyBeamScan(BeamScan* dest) const
{
    dest->m_beams = m_beams; // TODO cleanup: here we are overwriting what was set in the c'tor
    if (m_pol_analyzer)
        dest->setAnalyzer(m_pol_analyzer->direction(), m_pol_analyzer->transmissionParallel(),
                          m_pol_analyzer->transmissionAntiparallel());
}

std::vector<ParameterSample> BeamScan::drawDistribution(Distribution1D* distrib, double point) const
{
    if (!distrib)
        return {{point, 1}};
    distrib->setMean(point);
    return distrib->distributionSamples();
}

bool BeamScan::isCommonIntensity() const
{
    const auto ref = m_beams.front()->intensity();
    for (const auto& b : m_beams)
        if (!Numeric::almostEqual(b->intensity(), ref, 1))
            return false;
    return true;
}

bool BeamScan::isCommonPolarization() const
{
    const auto ref = m_beams.front()->polVector();
    for (const auto& b : m_beams)
        if (!Numeric::almostEqual(b->polVector(), ref, 1))
            return false;
    return true;
}
