//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Simulation/DepthprobeSimulation.cpp
//! @brief     Implements class DepthprobeSimulation.
//!
//! @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/DepthprobeSimulation.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/Scale.h"
#include "Base/Progress/ProgressHandler.h"
#include "Base/Util/Assert.h"
#include "Base/Vector/GisasDirection.h"
#include "Device/Beam/Footprint.h"
#include "Device/Data/Datafield.h"
#include "Resample/Sample/ReSample.h"
#include "Sample/Element/Element.h"
#include "Sample/Element/ScanElement.h"
#include "Sample/Flux/ScalarFlux.h"
#include "Sample/Multilayer/Layer.h"
#include "Sim/Scan/AlphaScan.h"
#include <valarray>

const int ZDirection_None = 0;
const int ZDirection_Reflected = 1;
const int ZDirection_Transmitted = 2;
const int WaveProperty_Intensity = 0;
const int WaveProperty_Modulus = 4;
const int WaveProperty_Phase = 8;

DepthprobeSimulation::DepthprobeSimulation(const BeamScan& scan, const Sample& sample,
                                           const Scale& zaxis, int flags)
    : Simulation(sample)
    , m_scan(scan.clone())
    , m_z_axis(zaxis.clone())
    , m_flags(flags)
{
}

DepthprobeSimulation::~DepthprobeSimulation() = default;

//... Public non-virtual methods:

Datafield DepthprobeSimulation::transmitted() const
{
    const ReSample re_sample = ReSample::processedSample(*sample(), options());
    std::vector<ScanElement> eles = m_scan->generateElements();

    const size_t n_layers = re_sample.numberOfSlices();
    std::vector<double> out(m_scan->nScan(), 0.0);

    for (ScanElement& ele : eles) {
        if (!ele.isCalculated() || ele.k().z() >= 0)
            continue;
        const Fluxes fluxes = re_sample.fluxesIn(ele.k(), doItPolarized());
        const auto* f_vac = dynamic_cast<const ScalarFlux*>(fluxes[0].get());
        const auto* f_sub = dynamic_cast<const ScalarFlux*>(fluxes[n_layers - 1].get());
        ASSERT(f_vac && f_sub);
        const double kz_vac = f_vac->getScalarKz().real();
        const double kz_sub = f_sub->getScalarKz().real();
        const double T = kz_vac > 0 ? kz_sub / kz_vac * std::norm(f_sub->getScalarT()) : 0.0;
        out[ele.i_out()] += T * ele.beamIntensity() * ele.weight();
    }

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

Datafield DepthprobeSimulation::reflected() const
{
    const ReSample re_sample = ReSample::processedSample(*sample(), options());
    std::vector<ScanElement> eles = m_scan->generateElements();

    std::vector<double> out(m_scan->nScan(), 0.0);

    for (ScanElement& ele : eles) {
        if (!ele.isCalculated() || ele.k().z() >= 0)
            continue;
        const Fluxes fluxes = re_sample.fluxesIn(ele.k(), doItPolarized());
        const auto* f_vac = dynamic_cast<const ScalarFlux*>(fluxes[0].get());
        ASSERT(f_vac);
        const double R = std::norm(f_vac->getScalarR());
        out[ele.i_out()] += R * ele.beamIntensity() * ele.weight();
    }

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

//... Overridden executors:

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

void DepthprobeSimulation::runComputation(const ReSample& re_sample, size_t i, double weight)
{
    const size_t Nz = m_z_axis->size();
    ASSERT(m_cache.size() / Nz == m_eles.size());

    std::valarray<double> z_intensities; //!< simulated intensity for given z positions
    z_intensities.resize(Nz, 0.0);

    ScanElement& scan_ele = *(m_eles.begin() + static_cast<long>(i));

    const R3 ki = scan_ele.k();

    if (scan_ele.isCalculated() && ki.z() < 0) {
        const size_t n_layers = re_sample.numberOfSlices();
        size_t start_z_ind = Nz;

        const R3 ki = scan_ele.k();
        const Fluxes fluxes = re_sample.fluxesIn(ki, doItPolarized());

        double z_layer_bottom(0.0);
        double z_layer_top(0.0);
        for (size_t i_layer = 0; i_layer < n_layers && start_z_ind != 0; ++i_layer) {
            z_layer_bottom = re_sample.avgeSlice(i_layer).low();
            z_layer_top = i_layer ? re_sample.avgeSlice(i_layer).hig() : z_layer_bottom;

            // get R & T coefficients for current layer
            const auto* flux = dynamic_cast<const ScalarFlux*>(fluxes[i_layer].get());
            ASSERT(flux);
            const complex_t R = flux->getScalarR();
            const complex_t T = flux->getScalarT();
            const complex_t kz_out = flux->getScalarKz();
            const complex_t kz_in = -kz_out;

            // Compute intensity for z's of the layer
            size_t ip1_z = start_z_ind;
            for (; ip1_z > 0; --ip1_z) {
                const size_t i_z = ip1_z - 1;
                if (i_layer + 1 != n_layers && m_z_axis->binCenter(i_z) <= z_layer_bottom)
                    break;
                const double z = m_z_axis->binCenter(i_z) - z_layer_top;
                complex_t psi;
                if ((m_flags & 3) == ZDirection_None)
                    psi = R * exp_I(kz_out * z) + T * exp_I(kz_in * z);
                else if (m_flags & ZDirection_Reflected)
                    psi = R * exp_I(kz_out * z);
                else if (m_flags & ZDirection_Transmitted)
                    psi = T * exp_I(kz_in * z);
                else
                    throw std::runtime_error("Invalid combination of ZDirection flags");
                if ((m_flags & 12) == WaveProperty_Intensity)
                    z_intensities[i_z] = std::norm(psi);
                else if (m_flags & WaveProperty_Modulus)
                    z_intensities[i_z] = std::abs(psi);
                else if (m_flags & WaveProperty_Phase)
                    z_intensities[i_z] = std::arg(psi);
                else
                    throw std::runtime_error("Invalid combination of WaveProperty flags");
            }
            start_z_ind = ip1_z;
        }
    }

    for (double& v : z_intensities)
        v *= scan_ele.beamIntensity();

    for (size_t j = 0; j < Nz; ++j)
        m_cache[j * m_scan->nElements() + i] += z_intensities[j] * weight * scan_ele.weight();

    progress().incrementDone(1);
}

//... Overridden getters:

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

size_t DepthprobeSimulation::nOutChannels() const
{
    return nElements() * m_z_axis->size();
}

Datafield DepthprobeSimulation::packResult() const
{
    const size_t Nz = m_z_axis->size();
    const size_t Ns = m_scan->nScan();
    const size_t Nslong = m_scan->nElements();
    std::vector<double> out(Ns * Nz, 0.0);

    for (size_t j = 0; j < Nz; ++j)
        for (size_t i = 0; i < Nslong; ++i) {
            const ScanElement& ele = m_eles.at(i);
            const size_t i_real = ele.i_out();
            out[j * Ns + i_real] += m_cache[j * Nslong + i];
        }

    if (background())
        throw std::runtime_error("nonzero background is not supported by DepthprobeSimulation");

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