//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Simulation/SpecularSimulation.cpp
//! @brief     Implements class OffspecSimulation.
//!
//! @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/Simulation/SpecularSimulation.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/Scale.h"
#include "Base/Progress/ProgressHandler.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/IFootprint.h"
#include "Device/Data/Datafield.h"
#include "Device/Pol/PolFilter.h"
#include "Param/Distrib/Distributions.h"
#include "Param/Distrib/ParameterSample.h"
#include "Resample/Element/ScanElement.h"
#include "Resample/Processed/ReSample.h"
#include "Resample/Specular/ComputeFluxMagnetic.h"
#include "Resample/Specular/ComputeFluxScalar.h"
#include "Sim/Background/ConstantBackground.h"
#include "Sim/Computation/SpecularComputation.h"
#include "Sim/Scan/BeamScan.h"

SpecularSimulation::SpecularSimulation(const BeamScan& scan, const Sample& sample)
    : ISimulation(sample)
    , m_scan(scan.clone())
{
    // TODO: move inside AlphaScan when pointwise resolution is implemented
    if (scan.coordinateAxis()->min() < 0.0)
        throw std::runtime_error("Invalid scan: minimum value on coordinate axis is negative");
}

SpecularSimulation::~SpecularSimulation() = default;

//... Overridden executors:

void SpecularSimulation::initScanElementVector()
{
    m_eles = m_scan->generateElements();
}

void SpecularSimulation::runComputation(const ReSample& re_sample, size_t iElement, double weight)
{
    ScanElement& ele = *(m_eles.begin() + static_cast<long>(iElement));

    double refl = 0;
    if (ele.isCalculated()) {

        const SliceStack& slices = re_sample.averageSlices();
        std::vector<complex_t> kz_stack = m_scan->produceKz(slices, ele.k());

        if (re_sample.polarizing()) {
            const SpinMatrix R = Compute::polarizedReflectivity(slices, kz_stack, true);
            const SpinMatrix& polMatrix = ele.polarizer();
            const SpinMatrix& anaMatrix = ele.analyzer();

            refl = Compute::magneticR(R, polMatrix, anaMatrix);

        } else {
            const complex_t R = Compute::scalarReflectivity(slices, kz_stack);
            refl = Compute::scalarR(R);
        }
    }

    m_cache[iElement] += refl * ele.footprint() * weight * ele.beamIntensity() * ele.weight();

    progress().incrementDone(1);
}

//... Overridden getters:

bool SpecularSimulation::force_polarized() const
{
    return m_scan->analyzer() ? m_scan->analyzer()->BlochVector() != R3{} : false;
}

size_t SpecularSimulation::nElements() const
{
    return m_scan->nElements();
}

Datafield SpecularSimulation::packResult() const
{
    std::vector<double> vec(m_scan->nScan(), 0.0);
    for (size_t i = 0; i < nElements(); i++) {
        const ScanElement& ele = m_eles.at(i);
        vec.at(ele.i_out()) += m_cache.at(i);
    }

    double common_intensity = m_scan->commonIntensity();
    if (const auto b = dynamic_cast<const ConstantBackground*>(background()))
        common_intensity += b->backgroundValue();

    if (background())
        for (size_t i = 0; i < m_scan->nScan(); i++)
            vec[i] = background()->addBackground(vec[i]);

    for (size_t i = 0; i < m_scan->nScan(); i++)
        if (common_intensity > 0)
            vec[i] /= common_intensity;
        else
            vec[i] = 0;

    return {std::vector<const Scale*>{m_scan->coordinateAxis()->clone()}, vec};
}
