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

#include "Sample/Particle/Mixture.h"
#include "Base/Spin/SpinMatrix.h"
#include "Base/Util/Assert.h"
#include "Sample/Element/DiffuseElement.h"
#include "Sample/Material/Admixtures.h"
#include "Sample/Particle/DwbaChannels.h"
#include "Sample/Particle/Rotations.h"
#include <stdexcept>

Mixture::Mixture() = default;

Mixture::~Mixture() = default;

void Mixture::addParticle(const IParticle& particle, double abundance)
{
    addParticle(std::unique_ptr<IParticle>(particle.clone()), abundance);
}

void Mixture::addParticle(std::unique_ptr<IParticle> particle, double abundance)
{
    if (abundance <= 0)
        throw std::runtime_error("Mixture::addParticle: abundance must be positive");
    m_particles.emplace_back(std::move(particle));
    m_abundances.push_back(abundance);
}

Mixture* Mixture::clone() const
{
    auto* result = new Mixture();
    for (size_t i = 0; i < m_particles.size(); ++i)
        result->addParticle(*m_particles[i], m_abundances[i]);
    if (m_validated)
        result->validate();
    instructClone(result);
    return result;
}

void Mixture::validate() const
{
    if (m_validated)
        return;

    for (const auto& particle : m_particles)
        particle->validate();

    // Compute cached normalized m_weights.
    double total = 0;
    for (const double a : m_abundances)
        total += a;
    if (total == 0)
        throw std::runtime_error("Empty Mixture");
    ASSERT(total > 0);
    for (const double a : m_abundances)
        m_weights.push_back(a / total);

    m_validated = true;
}

std::vector<const IParticle*> Mixture::particles() const
{
    std::vector<const IParticle*> result;
    result.reserve(m_particles.size());
    for (const auto& p : m_particles)
        result.emplace_back(p.get());
    return result;
}

std::vector<IParticle*> Mixture::mutableParticles()
{
    std::vector<IParticle*> result;
    result.reserve(m_particles.size());
    for (const auto& p : m_particles)
        result.emplace_back(p.get());
    return result;
}

void Mixture::setSliceNumber(size_t n)
{
    IParticle::setSliceNumber(n);
    for (auto& particle : m_particles)
        particle->setSliceNumber(n);
}

double Mixture::bottomZ() const
{
    if (m_particles.empty())
        return particlePosition().z();
    double result = m_particles[0]->bottomZ();
    for (size_t i = 1; i < m_particles.size(); ++i)
        result = std::min(result, m_particles[i]->bottomZ());
    return result + particlePosition().z();
}

double Mixture::topZ() const
{
    if (m_particles.empty())
        return particlePosition().z();
    double result = m_particles[0]->topZ();
    for (size_t i = 1; i < m_particles.size(); ++i)
        result = std::max(result, m_particles[i]->topZ());
    return result + particlePosition().z();
}

double Mixture::volumeToAdmix() const
{
    ASSERT(m_validated);
    double weighted_volume = 0.0;
    for (size_t i = 0; i < m_particles.size(); ++i)
        weighted_volume += m_weights[i] * m_particles[i]->volumeToAdmix();
    return weighted_volume;
}

bool Mixture::hasRadius() const
{
    if (!m_validated)
        return false;
    for (const auto& particle : m_particles)
        if (!particle->hasRadius())
            return false;
    return true;
}

double Mixture::radius() const
{
    ASSERT(m_validated);
    double weighted_radius = 0;
    for (size_t i = 0; i < m_particles.size(); ++i)
        weighted_radius += m_weights[i] * m_particles[i]->radius();
    return weighted_radius;
}

double Mixture::radialExtension() const
{
    ASSERT(m_validated);
    double weighted_extension = 0;
    for (size_t i = 0; i < m_particles.size(); ++i)
        weighted_extension += m_weights[i] * m_particles[i]->radialExtension();
    return weighted_extension;
}

bool Mixture::containsSoftParticle() const
{
    for (const auto& particle : m_particles)
        if (particle->containsSoftParticle())
            return true;
    return false;
}

std::unique_ptr<const Material> Mixture::materialToAdmix() const
{
    ASSERT(m_validated);
    std::vector<OneAdmixture> admixtures;
    for (size_t i = 0; i < m_particles.size(); ++i)
        admixtures.emplace_back(m_weights[i] * m_particles[i]->volumeToAdmix(),
                                *m_particles[i]->materialToAdmix());
    return MixUtil::averagedMaterial("Mixture_avg", admixtures);
}

complex_t Mixture::theFF(const WavevectorInfo& wavevectors) const
{
    ASSERT(m_validated);
    complex_t result = 0;
    for (size_t i = 0; i < m_particles.size(); ++i)
        result += m_weights[i] * m_particles[i]->theFF(wavevectors);
    return result;
}

SpinMatrix Mixture::thePolFF(const WavevectorInfo& wavevectors) const
{
    ASSERT(m_validated);
    SpinMatrix result;
    for (size_t i = 0; i < m_particles.size(); ++i)
        result += m_weights[i] * m_particles[i]->thePolFF(wavevectors);
    return result;
}

ScatteringMoments Mixture::scatteringMoments(const DiffuseElement& ele) const
{
    // Welford's algorithm for numerically stable variance computation,
    // adapted for mixture of distributions where each component has intrinsic variance.
    ASSERT(m_validated);
    double W = 0;
    complex_t mean = 0;
    double S = 0;
    for (size_t i = 0; i < m_particles.size(); ++i) {
        const double w = m_weights[i];
        const auto [f, v] = m_particles[i]->scatteringMoments(ele);
        const complex_t old_mean = mean;
        W += w;
        mean += (w / W) * (f - mean);
        S += w * v + w * std::real((f - old_mean) * std::conj(f - mean));
    }
    return {mean, S};
}

PolScatteringMoments Mixture::polScatteringMoments(const DiffuseElement& ele) const
{
    // Welford's algorithm for numerically stable variance computation,
    // adapted for polarized case with SpinMatrix values.
    ASSERT(m_validated);
    const SpinMatrix& polarizer = ele.polarizer();
    double W = 0;
    SpinMatrix mean;
    SpinMatrix S;
    for (size_t i = 0; i < m_particles.size(); ++i) {
        const double w = m_weights[i];
        const auto [f, v] = m_particles[i]->polScatteringMoments(ele);
        const SpinMatrix old_mean = mean;
        W += w;
        mean += (w / W) * (f - mean);
        S += w * v + w * (f - old_mean) * polarizer * (f - mean).adjoint();
    }
    return {mean, S};
}

ChannelMoments Mixture::channelMoments(const DwbaChannels& channels) const
{
    // Law of total covariance over mixture components alpha with weights p_alpha:
    //   mu[m]   = sum_alpha p_alpha * mu^alpha[m]
    //   C[m][n] = sum_alpha p_alpha * C^alpha[m][n]
    //           + sum_alpha p_alpha * (mu^alpha[m] - mu[m]) * conj(mu^alpha[n] - mu[n])
    ASSERT(m_validated);
    const size_t N = m_particles.size();
    // Pass 1: collect per-component moments and compute mixture mean
    std::vector<ChannelMoments> child(N);
    ChannelMoments total;
    for (size_t i = 0; i < N; ++i) {
        child[i] = m_particles[i]->channelMoments(channels);
        for (size_t m = 0; m < channels.n_active; ++m)
            total.mu[m] += m_weights[i] * child[i].mu[m];
    }
    // Pass 2: accumulate within-component and between-component covariance
    for (size_t i = 0; i < N; ++i) {
        const double w = m_weights[i];
        std::array<complex_t, N_DWBA> dev{};
        for (size_t m = 0; m < channels.n_active; ++m)
            dev[m] = child[i].mu[m] - total.mu[m];
        for (size_t m = 0; m < channels.n_active; ++m)
            for (size_t n = 0; n < channels.n_active; ++n)
                total.C[m][n] += w * child[i].C[m][n] + w * dev[m] * std::conj(dev[n]);
    }
    return total;
}

PolChannelMoments Mixture::polChannelMoments(const PolDwbaChannels& channels) const
{
    // Law of total covariance for the W-contracted polarized covariance:
    //   mu[m]   = sum_alpha p_alpha * mu^alpha[m]
    //   C[m][n] = sum_alpha p_alpha * C^alpha[m][n]
    //           + sum_alpha p_alpha * W_m(dev_m) * P * W_n(dev_n)†
    // where dev_m = mu^alpha[m] - mu[m].
    ASSERT(m_validated);
    const size_t N = m_particles.size();
    const SpinMatrix& P = channels.polarizer;

    // Pass 1: collect per-component moments and compute mixture mean
    std::vector<PolChannelMoments> child(N);
    PolChannelMoments total;
    for (size_t i = 0; i < N; ++i) {
        child[i] = m_particles[i]->polChannelMoments(channels);
        for (size_t m = 0; m < channels.n_active; ++m)
            total.mu[m] += m_weights[i] * child[i].mu[m];
    }

    // Pass 2: accumulate within-component and between-component covariance
    for (size_t i = 0; i < N; ++i) {
        const double w = m_weights[i];
        // Compute W_m(dev_m) and cache the channel-independent factors used in C[m][n].
        std::array<SpinMatrix, N_POL_DWBA> w_dev;
        std::array<SpinMatrix, N_POL_DWBA> P_w_dev_adj;
        for (size_t m = 0; m < channels.n_active; ++m) {
            w_dev[m] = applyWeight(child[i].mu[m] - total.mu[m], channels.flux[m]);
            P_w_dev_adj[m] = P * w_dev[m].adjoint();
        }
        for (size_t m = 0; m < channels.n_active; ++m)
            for (size_t n = 0; n < channels.n_active; ++n)
                total.C[m][n] += w * child[i].C[m][n] + w * w_dev[m] * P_w_dev_adj[n];
    }
    return total;
}

SscaMoments Mixture::sscaMoments(const DiffuseElement& ele, double kappa_qp,
                                 double mean_extension) const
{
    ASSERT(m_validated);
    double intensity = 0;
    complex_t plus = 0;
    complex_t minus = 0;
    complex_t coupling = 0;
    for (size_t i = 0; i < m_particles.size(); ++i) {
        const double w = m_weights[i];
        const auto ssca = m_particles[i]->sscaMoments(ele, kappa_qp, mean_extension);
        intensity += w * ssca.intensity;
        plus += w * ssca.shifted_plus;
        minus += w * ssca.shifted_minus;
        coupling += w * ssca.phase_coupling;
    }
    return {intensity, plus, minus, coupling};
}

PolSscaMoments Mixture::polSscaMoments(const DiffuseElement& ele, double kappa_qp,
                                       double mean_extension) const
{
    ASSERT(m_validated);
    SpinMatrix intensity;
    SpinMatrix plus;
    SpinMatrix minus;
    complex_t coupling = 0;
    for (size_t i = 0; i < m_particles.size(); ++i) {
        const double w = m_weights[i];
        const auto ssca = m_particles[i]->polSscaMoments(ele, kappa_qp, mean_extension);
        intensity += w * ssca.intensity;
        plus += w * ssca.shifted_plus;
        minus += w * ssca.shifted_minus;
        coupling += w * ssca.phase_coupling;
    }
    return {intensity, plus, minus, coupling};
}
