//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Material/MaterialUtil.cpp
//! @brief     Implements functions in namespace MaterialUtil.
//!
//! @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 "Sample/Material/MaterialUtil.h"
#include "Base/Const/PhysicalConstants.h"
#include "Base/Spin/SpinMatrix.h"
#include "Base/Util/Assert.h"
#include "Sample/Material/IMaterialImpl.h"
#include "Sample/Material/MaterialFactoryFuncs.h"

using PhysConsts::g_factor_n;
using PhysConsts::h_bar;
using PhysConsts::m_n;
using PhysConsts::mu_N;

// The factor 1e-18 is here to have unit: 1/T*nm^-2
constexpr double magnetic_prefactor = (m_n * g_factor_n * mu_N / h_bar / h_bar) * 1e-18;

namespace {

// Pauli matrices
const SpinMatrix Pauli_X(0, 1, 1, 0);
const SpinMatrix Pauli_Y(0, -I, I, 0);
const SpinMatrix Pauli_Z(1, 0, 0, -1);

} // namespace


template <typename T>
SpinMatrix MaterialUtil::MagnetizationCorrection(complex_t unit_factor, double magnetic_factor,
                                                 const Vec3<T>& polarization)
{
    SpinMatrix result = unit_factor * SpinMatrix::One()
                        + magnetic_factor
                              * (Pauli_X * polarization.x() + Pauli_Y * polarization.y()
                                 + Pauli_Z * polarization.z());
    return result;
}

// Prompt compilation for real and complex vectors:
template SpinMatrix MaterialUtil::MagnetizationCorrection(complex_t unit_factor,
                                                          double magnetic_factor,
                                                          const R3& polarization);
template SpinMatrix MaterialUtil::MagnetizationCorrection(complex_t unit_factor,
                                                          double magnetic_factor,
                                                          const C3& polarization);


complex_t MaterialUtil::ScalarReducedPotential(complex_t n, const R3& k, double n_ref)
{
    return n * n - n_ref * n_ref * R3Util::sin2Theta(k);
}

SpinMatrix MaterialUtil::PolarizedReducedPotential(complex_t n, const R3& b_field, const R3& k,
                                                   double n_ref)
{
    double factor = magnetic_prefactor / k.mag2();
    complex_t unit_factor = ScalarReducedPotential(n, k, n_ref);
    return MagnetizationCorrection(unit_factor, factor, b_field);
}

MATERIAL_TYPES MaterialUtil::checkMaterialTypes(const std::vector<const Material*>& materials)
{
    MATERIAL_TYPES result = MATERIAL_TYPES::RefractiveMaterial;
    bool isDefault = true;
    for (const Material* mat : materials) {
        if (isDefault) {
            result = mat->typeID();
            isDefault = mat->isDefaultMaterial();
            continue;
        }
        if (mat->typeID() != result && !mat->isDefaultMaterial())
            return MATERIAL_TYPES::InvalidMaterialType;
    }
    return result;
}

Material MaterialUtil::averagedMaterial(const std::string& name, const std::vector<double>& weights,
                                        const std::vector<Material>& materials)
{
    const size_t N = materials.size();
    ASSERT(weights.size() == N);
    ASSERT(N > 0);

    double totalWeight = 0;
    for (double w : weights) {
        ASSERT(w >= 0);
        totalWeight += w;
    }
    ASSERT(totalWeight > 0);

    R3 avgeMagn;
    complex_t avgeData{0., 0.};
    const auto type = materials[0].typeID();

    for (size_t i = 0; i < N; ++i) {
        double w = weights[i] / totalWeight;
        const Material& mat = materials[i];
        if (mat.typeID() != type)
            throw std::runtime_error(
                "Invalid mixture of different material definitions (refractive index vs SLD)");
        avgeMagn += w * mat.magnetization();
        if (type == MATERIAL_TYPES::RefractiveMaterial) {
            const complex_t mdc = std::conj(mat.refractiveIndex_or_SLD());
            avgeData += w * (mdc * mdc - 2.0 * mdc);
        } else if (type == MATERIAL_TYPES::MaterialBySLD) {
            avgeData += w * mat.refractiveIndex_or_SLD();
        } else
            ASSERT_NEVER;
    }

    if (type == MATERIAL_TYPES::RefractiveMaterial) {
        avgeData = std::conj(1.0 - std::sqrt(1.0 + avgeData));
        return RefractiveMaterial(name, avgeData.real(), avgeData.imag(), avgeMagn);
    }
    if (type == MATERIAL_TYPES::MaterialBySLD)
        return MaterialBySLD(name, avgeData.real(), avgeData.imag(), avgeMagn);
    ASSERT_NEVER;
}

// Tested by Tests/Unit/Sim/Sample/MaterialTest.cpp

Material MaterialUtil::averagedMaterial(const Material& base_mat, const Admixtures& admixtures)
{
    double totalAdmixedFraction = 0;
    for (const OneAdmixture& admix : admixtures) {
        ASSERT(admix.fraction >= 0);
        if (admix.fraction > 1)
            throw std::runtime_error("Volume fraction of one admixture component is more than 1");
        totalAdmixedFraction += admix.fraction;
    }
    if (totalAdmixedFraction > 1)
        throw std::runtime_error("Volume fractions of sample components add to more than 1");

    std::vector<double> weights;
    std::vector<Material> materials;

    weights.push_back(1 - totalAdmixedFraction);
    materials.push_back(base_mat);

    for (const OneAdmixture& admix : admixtures) {
        weights.push_back(admix.fraction);
        materials.push_back(admix.material);
    }

    const std::string avge_name = base_mat.materialName() + "_avg";

    return averagedMaterial(avge_name, weights, materials);
}
