//  ************************************************************************************************
//
//  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/Scale.h"
#include "Base/Progress/ProgressHandler.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/IFootprint.h"
#include "Device/Coord/CoordSystem1D.h"
#include "Device/Data/Datafield.h"
#include "Device/Histo/SimulationResult.h"
#include "Param/Distrib/Distributions.h"
#include "Param/Distrib/ParameterSample.h"
#include "Resample/Element/SpecularElement.h"
#include "Resample/Processed/ReSample.h"
#include "Resample/Specular/ComputeFluxMagnetic.h"
#include "Resample/Specular/ComputeFluxScalar.h"
#include "Sim/Background/IBackground.h"
#include "Sim/Computation/SpecularComputation.h"
#include "Sim/Scan/AlphaScan.h"

namespace {

//! Performs the actual simulation, and sets intensity, for given range of SpecularElement%s.
double reflectedIntensity(const ReSample& re_sample, SpecularElement& ele)
{
    if (!ele.isCalculated())
        return 0;

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

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

        return Compute::magneticR(R, polMatrix, anaMatrix);

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

} // namespace


SpecularSimulation::SpecularSimulation(const IBeamScan& scan, const MultiLayer& 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;

const ICoordSystem* SpecularSimulation::simCoordSystem() const
{
    return m_scan->scanCoordSystem();
}

//... Overridden executors:

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

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

    progress().incrementDone(1);
}

//... Overridden getters:

bool SpecularSimulation::force_polarized() const
{
    return m_scan->polarized();
}

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

SimulationResult SpecularSimulation::packResult()
{
    std::vector<double> vec(m_scan->nScan(), 0.0);
    for (size_t i = 0; i < nElements(); i++) {
        const SpecularElement& ele = m_eles.at(i);
        vec.at(ele.i_out()) += m_scan->intensity() * m_cache.at(i) * ele.weight();
    }
    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++)
        vec[i] /= m_scan->intensity();

    Datafield data({m_scan->coordinateAxis()->clone()}, vec);
    return {data, simCoordSystem()};
}
