//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Particle/IReParticle.cpp
//! @brief     Implements class interface IReParticle.
//!
//! @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 "Resample/Particle/IReParticle.h"
#include "Base/Util/Assert.h"
#include "Resample/Element/DiffuseElement.h"
#include "Resample/Flux/MatrixFlux.h"
#include "Resample/Flux/ScalarFlux.h"
#include "Sample/Material/Admixtures.h"

IReParticle::IReParticle() = default;

IReParticle::IReParticle(const std::optional<size_t>& i_layer)
    : m_i_layer(i_layer)
{
}

IReParticle::~IReParticle() = default;

OneAdmixture IReParticle::admixed() const
{
    return {m_admixed_fraction, *m_admixed_material};
}

void IReParticle::setAdmixedFraction(double fraction)
{
    m_admixed_fraction = fraction;
}

void IReParticle::setAdmixedMaterial(const Material& material)
{
    m_admixed_material = std::make_unique<Material>(material);
}

bool IReParticle::consideredEqualTo(const IReParticle& ire) const
{
    bool same_admixed_material =
        m_admixed_material ? (*m_admixed_material == ire.admixed().material) : true;

    return m_i_layer == ire.i_layer() && m_admixed_fraction == ire.admixedFraction()
           && same_admixed_material;
}

R3 IReParticle::posDiff(const R3* a, const R3* b)
{
    R3 pos_a = a ? *a : R3();
    R3 pos_b = b ? *b : R3();
    return pos_a - pos_b;
}

complex_t IReParticle::phaseFactor(const WavevectorInfo& wavevectors, const R3* position)
{
    if (!position)
        return 1;
    return exp_I(position->dot(wavevectors.getQ()));
}

complex_t IReParticle::coherentFF(const std::vector<ScalarTerm>& components, const R3& shift) const
{
    complex_t result;
    for (const auto& term : components)
        result += term.dwba_term * phaseFactor(term.wavevectors, &shift);
    return result;
}

std::vector<ScalarTerm> IReParticle::calcCoherentComponents(const DiffuseElement& ele) const
{
    const WavevectorInfo& wavevectors = ele.wavevectorInfo();

    if (!i_layer().has_value()) {
        // no slicing, pure Born approximation
        return {{wavevectors, theFF(wavevectors)}};
    }

    ASSERT(i_layer().has_value());
    const size_t iLayer = i_layer().value();

    const auto* inFlux = dynamic_cast<const ScalarFlux*>(ele.fluxIn(iLayer));
    const auto* outFlux = dynamic_cast<const ScalarFlux*>(ele.fluxOut(iLayer));

    // Retrieve the two different incoming wavevectors in the layer
    const C3& ki = wavevectors.getKi();
    const complex_t kiz = inFlux->getScalarKz();
    const C3 k_i_T{ki.x(), ki.y(), -kiz};
    const C3 k_i_R{ki.x(), ki.y(), +kiz};

    // Retrieve the two different outgoing wavevector bins in the layer
    const C3& kf = wavevectors.getKf();
    const complex_t kfz = outFlux->getScalarKz();
    const C3 k_f_T{kf.x(), kf.y(), +kfz};
    const C3 k_f_R{kf.x(), kf.y(), -kfz};

    // Construct the four different scattering contributions wavevector infos
    const double wavelength = wavevectors.vacuumLambda();
    const WavevectorInfo q_TT(k_i_T, k_f_T, wavelength);
    const WavevectorInfo q_RT(k_i_R, k_f_T, wavelength);
    const WavevectorInfo q_TR(k_i_T, k_f_R, wavelength);
    const WavevectorInfo q_RR(k_i_R, k_f_R, wavelength);

    // Get the four R,T coefficients
    const complex_t T_in = inFlux->getScalarT();
    const complex_t R_in = inFlux->getScalarR();
    const complex_t T_out = outFlux->getScalarT();
    const complex_t R_out = outFlux->getScalarR();

    // The four different scattering contributions;
    // S stands for scattering off the particle,
    // R for reflection off the layer interface.

    // Note that the order of multiplication matters:
    // If a prefactor is 0, then theFF() won't be called.

    const complex_t term_S = T_in * T_out * theFF(q_TT);
    const complex_t term_RS = R_in * T_out * theFF(q_RT);
    const complex_t term_SR = T_in * R_out * theFF(q_TR);
    const complex_t term_RSR = R_in * R_out * theFF(q_RR);

    return {
        {q_TT, term_S},
        {q_RT, term_RS},
        {q_TR, term_SR},
        {q_RR, term_RSR},
    };
}

SpinMatrix IReParticle::coherentPolFF(const std::vector<PolarizedTerm>& components,
                                      const R3& shift) const
{
    SpinMatrix result;
    for (const auto& term : components)
        result += term.dwba_term * phaseFactor(term.wavevectors, &shift);
    return result;
}

std::vector<PolarizedTerm> IReParticle::calcCoherentPolComponents(const DiffuseElement& ele) const
{
    const WavevectorInfo& wavevectors = ele.wavevectorInfo();

    if (!i_layer().has_value()) {
        // no slicing, pure Born approximation
        SpinMatrix o = thePolFF(wavevectors);
        return {{wavevectors, {-o.c, +o.a, -o.d, +o.b}}};
    }
    const size_t iLayer = i_layer().value();

    const auto* inFlux = dynamic_cast<const MatrixFlux*>(ele.fluxIn(iLayer));
    const auto* outFlux = dynamic_cast<const MatrixFlux*>(ele.fluxOut(iLayer));

    // the required wavevectors inside the layer for
    // different eigenmodes and in- and outgoing wavevector;
    const complex_t kix = wavevectors.getKi().x();
    const complex_t kiy = wavevectors.getKi().y();
    const Spinor& kiz = inFlux->getKz();
    const C3 ki_1R{kix, kiy, +kiz.u};
    const C3 ki_1T{kix, kiy, -kiz.u};
    const C3 ki_2R{kix, kiy, +kiz.v};
    const C3 ki_2T{kix, kiy, -kiz.v};

    const complex_t kfx = wavevectors.getKf().x();
    const complex_t kfy = wavevectors.getKf().y();
    const Spinor& kfz = outFlux->getKz();
    const C3 kf_1R{kfx, kfy, -kfz.u};
    const C3 kf_1T{kfx, kfy, +kfz.u};
    const C3 kf_2R{kfx, kfy, -kfz.v};
    const C3 kf_2T{kfx, kfy, +kfz.v};

    // now each of the 16 matrix terms of the polarized DWBA is calculated:
    // NOTE: when the underlying reflection/transmission coefficients are
    // scalar, the eigenmodes have identical eigenvalues and spin polarization
    // is used as a basis; in this case however the matrices get mixed:
    //     real m_M11 = calculated m_M12
    //     real m_M12 = calculated m_M11
    //     real m_M21 = calculated m_M22
    //     real m_M22 = calculated m_M21
    // since both eigenvalues are identical, this does not influence the result.
    SpinMatrix ff_BA; // will be overwritten

    double wavelength = wavevectors.vacuumLambda();

    // The following matrices each contain the four polarization conditions
    // (p->p, p->m, m->p, m->m)
    // The first two indices indicate a scattering from the 1/2 eigenstate into
    // the 1/2 eigenstate, while the capital indices indicate a reflection
    // before and/or after the scattering event (first index is in-state,
    // second is out-state; this also applies to the internal matrix indices)

    // eigenmode 1 -> eigenmode 1: direct scattering
    const WavevectorInfo q11_TT(ki_1T, kf_1T, wavelength);
    ff_BA = thePolFF(q11_TT);
    SpinMatrix M11_S(-DotProduct(outFlux->T1min(), ff_BA * inFlux->T1plus()),
                     +DotProduct(outFlux->T1plus(), ff_BA * inFlux->T1plus()),
                     -DotProduct(outFlux->T1min(), ff_BA * inFlux->T1min()),
                     +DotProduct(outFlux->T1plus(), ff_BA * inFlux->T1min()));
    // eigenmode 1 -> eigenmode 1: reflection and then scattering
    const WavevectorInfo q11_RT(ki_1R, kf_1T, wavelength);
    ff_BA = thePolFF(q11_RT);
    SpinMatrix M11_RS(-DotProduct(outFlux->T1min(), ff_BA * inFlux->R1plus()),
                      +DotProduct(outFlux->T1plus(), ff_BA * inFlux->R1plus()),
                      -DotProduct(outFlux->T1min(), ff_BA * inFlux->R1min()),
                      +DotProduct(outFlux->T1plus(), ff_BA * inFlux->R1min()));
    // eigenmode 1 -> eigenmode 1: scattering and then reflection
    const WavevectorInfo q11_TR(ki_1T, kf_1R, wavelength);
    ff_BA = thePolFF(q11_TR);
    SpinMatrix M11_SR(-DotProduct(outFlux->R1min(), ff_BA * inFlux->T1plus()),
                      +DotProduct(outFlux->R1plus(), ff_BA * inFlux->T1plus()),
                      -DotProduct(outFlux->R1min(), ff_BA * inFlux->T1min()),
                      +DotProduct(outFlux->R1plus(), ff_BA * inFlux->T1min()));
    // eigenmode 1 -> eigenmode 1: reflection, scattering and again reflection
    const WavevectorInfo q11_RR(ki_1R, kf_1R, wavelength);
    ff_BA = thePolFF(q11_RR);
    SpinMatrix M11_RSR(-DotProduct(outFlux->R1min(), ff_BA * inFlux->R1plus()),
                       +DotProduct(outFlux->R1plus(), ff_BA * inFlux->R1plus()),
                       -DotProduct(outFlux->R1min(), ff_BA * inFlux->R1min()),
                       +DotProduct(outFlux->R1plus(), ff_BA * inFlux->R1min()));

    // eigenmode 1 -> eigenmode 2: direct scattering
    const WavevectorInfo q12_TT(ki_1T, kf_2T, wavelength);
    ff_BA = thePolFF(q12_TT);
    SpinMatrix M12_S(-DotProduct(outFlux->T2min(), ff_BA * inFlux->T1plus()),
                     +DotProduct(outFlux->T2plus(), ff_BA * inFlux->T1plus()),
                     -DotProduct(outFlux->T2min(), ff_BA * inFlux->T1min()),
                     +DotProduct(outFlux->T2plus(), ff_BA * inFlux->T1min()));
    // eigenmode 1 -> eigenmode 2: reflection and then scattering
    const WavevectorInfo q12_RT(ki_1R, kf_2T, wavelength);
    ff_BA = thePolFF(q12_RT);
    SpinMatrix M12_RS(-DotProduct(outFlux->T2min(), ff_BA * inFlux->R1plus()),
                      +DotProduct(outFlux->T2plus(), ff_BA * inFlux->R1plus()),
                      -DotProduct(outFlux->T2min(), ff_BA * inFlux->R1min()),
                      +DotProduct(outFlux->T2plus(), ff_BA * inFlux->R1min()));
    // eigenmode 1 -> eigenmode 2: scattering and then reflection
    const WavevectorInfo q12_TR(ki_1T, kf_2R, wavelength);
    ff_BA = thePolFF(q12_TR);
    SpinMatrix M12_SR(-DotProduct(outFlux->R2min(), ff_BA * inFlux->T1plus()),
                      +DotProduct(outFlux->R2plus(), ff_BA * inFlux->T1plus()),
                      -DotProduct(outFlux->R2min(), ff_BA * inFlux->T1min()),
                      +DotProduct(outFlux->R2plus(), ff_BA * inFlux->T1min()));
    // eigenmode 1 -> eigenmode 2: reflection, scattering and again reflection
    const WavevectorInfo q12_RR(ki_1R, kf_2R, wavelength);
    ff_BA = thePolFF(q12_RR);
    SpinMatrix M12_RSR(-DotProduct(outFlux->R2min(), ff_BA * inFlux->R1plus()),
                       +DotProduct(outFlux->R2plus(), ff_BA * inFlux->R1plus()),
                       -DotProduct(outFlux->R2min(), ff_BA * inFlux->R1min()),
                       +DotProduct(outFlux->R2plus(), ff_BA * inFlux->R1min()));

    // eigenmode 2 -> eigenmode 1: direct scattering
    const WavevectorInfo q21_TT(ki_2T, kf_1T, wavelength);
    ff_BA = thePolFF(q21_TT);
    SpinMatrix M21_S(-DotProduct(outFlux->T1min(), ff_BA * inFlux->T2plus()),
                     +DotProduct(outFlux->T1plus(), ff_BA * inFlux->T2plus()),
                     -DotProduct(outFlux->T1min(), ff_BA * inFlux->T2min()),
                     +DotProduct(outFlux->T1plus(), ff_BA * inFlux->T2min()));
    // eigenmode 2 -> eigenmode 1: reflection and then scattering
    const WavevectorInfo q21_RT(ki_2R, kf_1T, wavelength);
    ff_BA = thePolFF(q21_RT);
    SpinMatrix M21_RS(-DotProduct(outFlux->T1min(), ff_BA * inFlux->R2plus()),
                      +DotProduct(outFlux->T1plus(), ff_BA * inFlux->R2plus()),
                      -DotProduct(outFlux->T1min(), ff_BA * inFlux->R2min()),
                      +DotProduct(outFlux->T1plus(), ff_BA * inFlux->R2min()));
    // eigenmode 2 -> eigenmode 1: scattering and then reflection
    const WavevectorInfo q21_TR(ki_2T, kf_1R, wavelength);
    ff_BA = thePolFF(q21_TR);
    SpinMatrix M21_SR(-DotProduct(outFlux->R1min(), ff_BA * inFlux->T2plus()),
                      +DotProduct(outFlux->R1plus(), ff_BA * inFlux->T2plus()),
                      -DotProduct(outFlux->R1min(), ff_BA * inFlux->T2min()),
                      +DotProduct(outFlux->R1plus(), ff_BA * inFlux->T2min()));
    // eigenmode 2 -> eigenmode 1: reflection, scattering and again reflection
    const WavevectorInfo q21_RR(ki_2R, kf_1R, wavelength);
    ff_BA = thePolFF(q21_RR);
    SpinMatrix M21_RSR(-DotProduct(outFlux->R1min(), ff_BA * inFlux->R2plus()),
                       +DotProduct(outFlux->R1plus(), ff_BA * inFlux->R2plus()),
                       -DotProduct(outFlux->R1min(), ff_BA * inFlux->R2min()),
                       +DotProduct(outFlux->R1plus(), ff_BA * inFlux->R2min()));

    // eigenmode 2 -> eigenmode 2: direct scattering
    const WavevectorInfo q22_TT(ki_2T, kf_2T, wavelength);
    ff_BA = thePolFF(q22_TT);
    SpinMatrix M22_S(-DotProduct(outFlux->T2min(), ff_BA * inFlux->T2plus()),
                     +DotProduct(outFlux->T2plus(), ff_BA * inFlux->T2plus()),
                     -DotProduct(outFlux->T2min(), ff_BA * inFlux->T2min()),
                     +DotProduct(outFlux->T2plus(), ff_BA * inFlux->T2min()));
    // eigenmode 2 -> eigenmode 2: reflection and then scattering
    const WavevectorInfo q22_RT(ki_2R, kf_2T, wavelength);
    ff_BA = thePolFF(q22_RT);
    SpinMatrix M22_RS(-DotProduct(outFlux->T2min(), ff_BA * inFlux->R2plus()),
                      +DotProduct(outFlux->T2plus(), ff_BA * inFlux->R2plus()),
                      -DotProduct(outFlux->T2min(), ff_BA * inFlux->R2min()),
                      +DotProduct(outFlux->T2plus(), ff_BA * inFlux->R2min()));
    // eigenmode 2 -> eigenmode 2: scattering and then reflection
    const WavevectorInfo q22_TR(ki_2T, kf_2R, wavelength);
    ff_BA = thePolFF(q22_TR);
    SpinMatrix M22_SR(-DotProduct(outFlux->R2min(), ff_BA * inFlux->T2plus()),
                      +DotProduct(outFlux->R2plus(), ff_BA * inFlux->T2plus()),
                      -DotProduct(outFlux->R2min(), ff_BA * inFlux->T2min()),
                      +DotProduct(outFlux->R2plus(), ff_BA * inFlux->T2min()));
    // eigenmode 2 -> eigenmode 2: reflection, scattering and again reflection
    const WavevectorInfo q22_RR(ki_2R, kf_2R, wavelength);
    ff_BA = thePolFF(q22_RR);
    SpinMatrix M22_RSR(-DotProduct(outFlux->R2min(), ff_BA * inFlux->R2plus()),
                       +DotProduct(outFlux->R2plus(), ff_BA * inFlux->R2plus()),
                       -DotProduct(outFlux->R2min(), ff_BA * inFlux->R2min()),
                       +DotProduct(outFlux->R2plus(), ff_BA * inFlux->R2min()));

    return {{q11_TT, M11_S},   {q12_TT, M12_S},   {q21_TT, M21_S},   {q22_TT, M22_S},
            {q11_RT, M11_RS},  {q12_RT, M12_RS},  {q21_RT, M21_RS},  {q22_RT, M22_RS},
            {q11_TR, M11_SR},  {q12_TR, M12_SR},  {q21_TR, M21_SR},  {q22_TR, M22_SR},
            {q11_RR, M11_RSR}, {q12_RR, M12_RSR}, {q21_RR, M21_RSR}, {q22_RR, M22_RSR}};
}
