//  ************************************************************************************************
//
//  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/Math/Numeric.h"
#include "Base/Spin/SpinMatrix.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/Beam.h"
#include "Device/Beam/IFootprint.h"
#include "Device/Pol/PolFilter.h"
#include "Param/Distrib/Distributions.h"
#include "Param/Distrib/ParameterSample.h"
#include "Resample/Slice/KzComputation.h"

BeamScan::BeamScan(Scale* axis)
    : m_axis(axis)
{
}

BeamScan::~BeamScan() = default;

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->BlochVector(), m_pol_analyzer->transmission());
}

std::vector<const INode*> BeamScan::nodeChildren() const
{
    std::vector<const INode*> result;
    if (m_pol_analyzer)
        result << m_pol_analyzer.get();
    return result;
}

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

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

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

void BeamScan::setIntensity(double intensity)
{
    for (auto& b : m_beams)
        b->setIntensity(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();
}

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

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

const IFootprint* BeamScan::commonFootprint() const
{
    if (!isCommonFootprint())
        throw std::runtime_error("Footprint function changes during scan. "
                                 "Use 'footprintAt(i)' instead.");
    return m_beams.front()->footprint();
}

const IFootprint* BeamScan::footprintAt(size_t i) const
{
    return m_beams[i]->footprint();
}

void BeamScan::setFootprint(const IFootprint* footprint)
{
    for (auto& b : m_beams)
        b->setFootprint(footprint);
}

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

Frame BeamScan::scanCoordSystem() const
{
    return {coordinateAxis()->clone()};
}

std::vector<complex_t> BeamScan::produceKz(const SliceStack& slices, const R3& k) const
{
    return Compute::Kz::computeKzFromRefIndices(slices, k);
}

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

bool BeamScan::isCommonFootprint() const
{
    const auto* ref_fp = m_beams.front()->footprint();
    if (!ref_fp) {
        // footprint is not set
        for (const auto& b : m_beams)
            if (b->footprint())
                return false;
    } else {
        // footprint is set
        double ref = ref_fp->widthRatio();
        for (const auto& b : m_beams) {
            if (!b->footprint())
                return false;
            // TODO check that footprint type is the same? Isn't it too expensive and redundant?
            if (!Numeric::almostEqual(b->footprint()->widthRatio(), ref, 1))
                return false;
        }
    }
    return true;
}

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