//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Particle/DwbaChannels.cpp
//! @brief     Implements DwbaChannels helpers.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2026
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Sample/Particle/DwbaChannels.h"
#include "Base/Util/Assert.h"
#include "Sample/Element/DiffuseElement.h"
#include "Sample/Flux/MatrixFlux.h"
#include "Sample/Flux/ScalarFlux.h"

DwbaChannels extractDwbaChannels(const DiffuseElement& ele, size_t i_slice)
{
    const WavevectorInfo& wv = ele.wavevectorInfo();
    const auto* in_flux = dynamic_cast<const ScalarFlux*>(ele.fluxIn(i_slice));
    const auto* out_flux = dynamic_cast<const ScalarFlux*>(ele.fluxOut(i_slice));
    ASSERT(in_flux && out_flux);

    if (in_flux->only_outgoing && out_flux->only_outgoing)
        // Born approximation: single channel with unit weight
        return {{{wv, wv, wv, wv}}, {1.0, 0.0, 0.0, 0.0}, 1};

    const C3& ki = wv.getKi();
    const complex_t kiz = in_flux->getScalarKz();
    const C3 k_i_T{ki.x(), ki.y(), -kiz};
    const C3 k_i_R{ki.x(), ki.y(), +kiz};

    const C3& kf = wv.getKf();
    const complex_t kfz = out_flux->getScalarKz();
    const C3 k_f_T{kf.x(), kf.y(), +kfz};
    const C3 k_f_R{kf.x(), kf.y(), -kfz};

    const double lambda = wv.vacuumLambda();
    const WavevectorInfo q_TT(k_i_T, k_f_T, lambda);
    const WavevectorInfo q_RT(k_i_R, k_f_T, lambda);
    const WavevectorInfo q_TR(k_i_T, k_f_R, lambda);
    const WavevectorInfo q_RR(k_i_R, k_f_R, lambda);

    const complex_t T_in = in_flux->getScalarT();
    const complex_t R_in = in_flux->getScalarR();
    const complex_t T_out = out_flux->getScalarT();
    const complex_t R_out = out_flux->getScalarR();

    return {{q_TT, q_RT, q_TR, q_RR}, {T_in * T_out, R_in * T_out, T_in * R_out, R_in * R_out}};
}

SpinMatrix applyWeight(const SpinMatrix& F, const PolChannelFlux& fl)
{
    const Spinor f_ip = F * fl.in_plus;
    const Spinor f_im = F * fl.in_minus;
    return {-Spinor::dotProduct(fl.out_minus, f_ip), +Spinor::dotProduct(fl.out_plus, f_ip),
            -Spinor::dotProduct(fl.out_minus, f_im), +Spinor::dotProduct(fl.out_plus, f_im)};
}

PolDwbaChannels extractPolDwbaChannels(const DiffuseElement& ele, size_t i_slice)
{
    const WavevectorInfo& wv = ele.wavevectorInfo();
    const auto* in_flux = dynamic_cast<const MatrixFlux*>(ele.fluxIn(i_slice));
    const auto* out_flux = dynamic_cast<const MatrixFlux*>(ele.fluxOut(i_slice));
    ASSERT(in_flux && out_flux);

    const SpinMatrix polarizer = ele.polarizer();
    const Spinor zero = Spinor::Zero();

    if (in_flux->only_outgoing && out_flux->only_outgoing) {
        // Born approximation: single active channel with identity-like flux.
        // W_0(F) = {-F.c, +F.a, -F.d, +F.b}, remaining channels zero.
        const PolChannelFlux born{Spinor(1, 0), Spinor(0, 1), Spinor(1, 0), Spinor(0, 1)};
        const PolChannelFlux nil{zero, zero, zero, zero};
        return {{wv, wv, wv, wv, wv, wv, wv, wv, wv, wv, wv, wv, wv, wv, wv, wv},
                {born, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil},
                polarizer,
                1};
    }

    // Incoming wavevectors for the two eigenmodes, transmitted and reflected
    const complex_t kix = wv.getKi().x();
    const complex_t kiy = wv.getKi().y();
    const Spinor& kiz = in_flux->getKz();
    const C3 ki_1T{kix, kiy, -kiz.u};
    const C3 ki_1R{kix, kiy, +kiz.u};
    const C3 ki_2T{kix, kiy, -kiz.v};
    const C3 ki_2R{kix, kiy, +kiz.v};

    // Outgoing wavevectors for the two eigenmodes, transmitted and reflected
    const complex_t kfx = wv.getKf().x();
    const complex_t kfy = wv.getKf().y();
    const Spinor& kfz = out_flux->getKz();
    const C3 kf_1T{kfx, kfy, +kfz.u};
    const C3 kf_1R{kfx, kfy, -kfz.u};
    const C3 kf_2T{kfx, kfy, +kfz.v};
    const C3 kf_2R{kfx, kfy, -kfz.v};

    const double lambda = wv.vacuumLambda();

    // Scattering wavevectors for 16 channels (eigenmode pair × T/R combination)

    // eigenmode 1 → eigenmode 1
    const WavevectorInfo q11_TT(ki_1T, kf_1T, lambda);
    const WavevectorInfo q11_RT(ki_1R, kf_1T, lambda);
    const WavevectorInfo q11_TR(ki_1T, kf_1R, lambda);
    const WavevectorInfo q11_RR(ki_1R, kf_1R, lambda);

    // eigenmode 1 → eigenmode 2
    const WavevectorInfo q12_TT(ki_1T, kf_2T, lambda);
    const WavevectorInfo q12_RT(ki_1R, kf_2T, lambda);
    const WavevectorInfo q12_TR(ki_1T, kf_2R, lambda);
    const WavevectorInfo q12_RR(ki_1R, kf_2R, lambda);

    // eigenmode 2 → eigenmode 1
    const WavevectorInfo q21_TT(ki_2T, kf_1T, lambda);
    const WavevectorInfo q21_RT(ki_2R, kf_1T, lambda);
    const WavevectorInfo q21_TR(ki_2T, kf_1R, lambda);
    const WavevectorInfo q21_RR(ki_2R, kf_1R, lambda);

    // eigenmode 2 → eigenmode 2
    const WavevectorInfo q22_TT(ki_2T, kf_2T, lambda);
    const WavevectorInfo q22_RT(ki_2R, kf_2T, lambda);
    const WavevectorInfo q22_TR(ki_2T, kf_2R, lambda);
    const WavevectorInfo q22_RR(ki_2R, kf_2R, lambda);

    // Flux spinors for 16 channels: {in_plus, in_minus, out_plus, out_minus}

    // eigenmode 1 → eigenmode 1
    const PolChannelFlux fl_11_TT{in_flux->T1plus(), in_flux->T1min(), out_flux->T1plus(),
                                  out_flux->T1min()};
    const PolChannelFlux fl_11_RT{in_flux->R1plus(), in_flux->R1min(), out_flux->T1plus(),
                                  out_flux->T1min()};
    const PolChannelFlux fl_11_TR{in_flux->T1plus(), in_flux->T1min(), out_flux->R1plus(),
                                  out_flux->R1min()};
    const PolChannelFlux fl_11_RR{in_flux->R1plus(), in_flux->R1min(), out_flux->R1plus(),
                                  out_flux->R1min()};

    // eigenmode 1 → eigenmode 2
    const PolChannelFlux fl_12_TT{in_flux->T1plus(), in_flux->T1min(), out_flux->T2plus(),
                                  out_flux->T2min()};
    const PolChannelFlux fl_12_RT{in_flux->R1plus(), in_flux->R1min(), out_flux->T2plus(),
                                  out_flux->T2min()};
    const PolChannelFlux fl_12_TR{in_flux->T1plus(), in_flux->T1min(), out_flux->R2plus(),
                                  out_flux->R2min()};
    const PolChannelFlux fl_12_RR{in_flux->R1plus(), in_flux->R1min(), out_flux->R2plus(),
                                  out_flux->R2min()};

    // eigenmode 2 → eigenmode 1
    const PolChannelFlux fl_21_TT{in_flux->T2plus(), in_flux->T2min(), out_flux->T1plus(),
                                  out_flux->T1min()};
    const PolChannelFlux fl_21_RT{in_flux->R2plus(), in_flux->R2min(), out_flux->T1plus(),
                                  out_flux->T1min()};
    const PolChannelFlux fl_21_TR{in_flux->T2plus(), in_flux->T2min(), out_flux->R1plus(),
                                  out_flux->R1min()};
    const PolChannelFlux fl_21_RR{in_flux->R2plus(), in_flux->R2min(), out_flux->R1plus(),
                                  out_flux->R1min()};

    // eigenmode 2 → eigenmode 2
    const PolChannelFlux fl_22_TT{in_flux->T2plus(), in_flux->T2min(), out_flux->T2plus(),
                                  out_flux->T2min()};
    const PolChannelFlux fl_22_RT{in_flux->R2plus(), in_flux->R2min(), out_flux->T2plus(),
                                  out_flux->T2min()};
    const PolChannelFlux fl_22_TR{in_flux->T2plus(), in_flux->T2min(), out_flux->R2plus(),
                                  out_flux->R2min()};
    const PolChannelFlux fl_22_RR{in_flux->R2plus(), in_flux->R2min(), out_flux->R2plus(),
                                  out_flux->R2min()};

    return {{q11_TT, q11_RT, q11_TR, q11_RR,          // eigenmode 1 → 1
             q12_TT, q12_RT, q12_TR, q12_RR,          // eigenmode 1 → 2
             q21_TT, q21_RT, q21_TR, q21_RR,          // eigenmode 2 → 1
             q22_TT, q22_RT, q22_TR, q22_RR},         // eigenmode 2 → 2
            {fl_11_TT, fl_11_RT, fl_11_TR, fl_11_RR,  // eigenmode 1 → 1
             fl_12_TT, fl_12_RT, fl_12_TR, fl_12_RR,  // eigenmode 1 → 2
             fl_21_TT, fl_21_RT, fl_21_TR, fl_21_RR,  // eigenmode 2 → 1
             fl_22_TT, fl_22_RT, fl_22_TR, fl_22_RR}, // eigenmode 2 → 2
            polarizer};
}
