feos 0.9.5

use super::*;
use feos_core::FeosResult;
use feos_core::parameter::GenericParameters;
use feos_dft::{FunctionalContribution, WeightFunction, WeightFunctionInfo, WeightFunctionShape};
use ndarray::{Array1, Array2, ArrayBase, ArrayView2, Axis, Data, Ix1, Slice};
use num_dual::DualNum;
use std::f64::consts::PI;
use std::ops::MulAssign;

pub const N0_CUTOFF: f64 = 1e-9;

impl Association {
    /// Association strength for functional of Yu and Wu with xi parameter for inhomogeneity.
    ///
    /// Uses the contact value of hard-sphere pair correlation function and model-specific
    /// implementations for the bonding volume.
    #[expect(clippy::too_many_arguments)]
    fn yu_wu_association_strength<A: AssociationStrength, D: DualNum<f64> + Copy>(
        &self,
        parameters: &AssociationParameters<A::Record>,
        model: &A,
        temperature: D,
        diameter: &DVector<D>,
        n2: D,
        n3i: D,
        xi: D,
    ) -> [DMatrix<D>; 2] {
        let p = parameters;

        let mut delta_ab = DMatrix::zeros(p.sites_a.len(), p.sites_b.len());
        for b in &p.binary_ab {
            let [i, j] = [b.id1, b.id2];
            let di = diameter[p.sites_a[i].assoc_comp];
            let dj = diameter[p.sites_b[j].assoc_comp];
            let k = di * dj / (di + dj) * (n2 * n3i);
            delta_ab[(i, j)] = n3i
                * (k * xi * (k / 18.0 + 0.5) + 1.0)
                * model.association_strength_ij(
                    temperature,
                    p.sites_a[i].assoc_comp,
                    p.sites_b[j].assoc_comp,
                    &b.model_record,
                )
        }
        let mut delta_cc = DMatrix::zeros(p.sites_c.len(), p.sites_c.len());
        for b in &p.binary_cc {
            let [i, j] = [b.id1, b.id2];
            let di = diameter[p.sites_c[i].assoc_comp];
            let dj = diameter[p.sites_c[j].assoc_comp];
            let k = di * dj / (di + dj) * (n2 * n3i);
            delta_cc[(i, j)] = n3i
                * (k * xi * (k / 18.0 + 0.5) + 1.0)
                * model.association_strength_ij(
                    temperature,
                    p.sites_c[i].assoc_comp,
                    p.sites_c[j].assoc_comp,
                    &b.model_record,
                )
        }
        [delta_ab, delta_cc]
    }
}

/// Implementation of the association Helmholtz energy functional of Yu and Wu.
///
/// [Yang-Xin Yu and Jianzhong Wu (2002)](https://aip.scitation.org/doi/abs/10.1063/1.1463435)
///
/// # Note
///  
/// Uses the contact value of the hard-sphere pair correlation function and model-specific
/// implementations for the bonding volume.
pub struct YuWuAssociationFunctional<'a, A: AssociationStrength> {
    model: &'a A,
    association_parameters: &'a AssociationParameters<A::Record>,
    association: Association,
}

impl<'a, A: AssociationStrength> YuWuAssociationFunctional<'a, A> {
    pub fn new<P, B, Bo, C, Data>(
        model: &'a A,
        parameters: &'a GenericParameters<P, B, A::Record, Bo, C, Data>,
        association: Option<Association>,
    ) -> Option<Self> {
        association.map(|a| Self {
            model,
            association_parameters: &parameters.association,
            association: a,
        })
    }
}

impl<'a, A: AssociationStrength + Sync + Send> FunctionalContribution
    for YuWuAssociationFunctional<'a, A>
where
    A::Record: Sync + Send,
{
    fn name(&self) -> &'static str {
        "Association"
    }

    fn weight_functions<N: DualNum<f64> + Copy>(&self, temperature: N) -> WeightFunctionInfo<N> {
        let p = self.model;
        let r = p.hs_diameter(temperature) * N::from(0.5);
        let [_, _, _, c3] = p.geometry_coefficients(temperature);
        WeightFunctionInfo::new(p.component_index().into_owned().into(), false)
            .add(
                WeightFunction::new_scaled(r.clone(), WeightFunctionShape::Delta),
                false,
            )
            .add(
                WeightFunction {
                    prefactor: c3.clone(),
                    kernel_radius: r.clone(),
                    shape: WeightFunctionShape::DeltaVec,
                },
                false,
            )
            .add(
                WeightFunction {
                    prefactor: c3,
                    kernel_radius: r,
                    shape: WeightFunctionShape::Theta,
                },
                true,
            )
    }

    fn helmholtz_energy_density<N: DualNum<f64> + Copy>(
        &self,
        temperature: N,
        weighted_densities: ArrayView2<N>,
    ) -> FeosResult<Array1<N>> {
        // number of segments
        let n = self.association_parameters.component_index.len();

        // number of dimensions
        let dim = (weighted_densities.shape()[0] - 1) / n - 1;

        // weighted densities
        let n0i = weighted_densities.slice_axis(Axis(0), Slice::new(0, Some(n as isize), 1));
        let n2vi: Vec<_> = (0..dim)
            .map(|i| {
                weighted_densities.slice_axis(
                    Axis(0),
                    Slice::new((n * (i + 1)) as isize, Some((n * (i + 2)) as isize), 1),
                )
            })
            .collect();
        let n3 = weighted_densities.index_axis(Axis(0), n * (dim + 1));

        // calculate rho0
        let [_, _, c2, _] = self.model.geometry_coefficients(temperature);
        let diameter = self.model.hs_diameter(temperature);
        let mut n2i = n0i.to_owned();
        for (i, mut n2i) in n2i.outer_iter_mut().enumerate() {
            n2i.mul_assign(diameter[i].powi(2) * c2[i] * PI);
        }
        let mut rho0: Array2<N> = (n2vi
            .iter()
            .fold(Array2::zeros(n0i.raw_dim()), |acc, n2vi| acc + n2vi * n2vi)
            / -(&n2i * &n2i)
            + 1.0)
            * n0i;
        rho0.iter_mut().zip(&n0i).for_each(|(rho0, &n0i)| {
            if n0i.re() < N0_CUTOFF {
                *rho0 = n0i;
            }
        });

        // calculate xi
        let n2v: Vec<_> = n2vi.iter().map(|n2vi| n2vi.sum_axis(Axis(0))).collect();
        let n2 = n2i.sum_axis(Axis(0));
        let mut xi = n2v
            .iter()
            .fold(Array1::zeros(n2.raw_dim()), |acc, n2v| acc + n2v * n2v)
            / -(&n2 * &n2)
            + 1.0;
        xi.iter_mut()
            .zip(&n0i.sum_axis(Axis(0)))
            .for_each(|(xi, &n0i)| {
                if n0i.re() < N0_CUTOFF {
                    *xi = N::one();
                }
            });

        // auxiliary variables
        let n3i = n3.mapv(|n3| (-n3 + 1.0).recip());

        self._helmholtz_energy_density(temperature, &rho0, &n2, &n3i, &xi)
    }
}

impl<'a, A: AssociationStrength> YuWuAssociationFunctional<'a, A> {
    pub fn _helmholtz_energy_density<N: DualNum<f64> + Copy, S: Data<Elem = N>>(
        &self,
        temperature: N,
        rho0: &Array2<N>,
        n2: &ArrayBase<S, Ix1>,
        n3i: &Array1<N>,
        xi: &Array1<N>,
    ) -> FeosResult<Array1<N>> {
        let a = &self.association_parameters;
        let t = temperature;

        let d = self.model.hs_diameter(t);

        match (
            a.sites_a.len() * a.sites_b.len(),
            a.sites_c.len(),
            self.association.force_cross_association,
        ) {
            (0, 0, _) => Ok(Array1::zeros(n3i.len())),
            (1, 0, false) => {
                let params = a.binary_ab.first().map(|r| &r.model_record);
                Ok(self.helmholtz_energy_density_ab_analytic(params, t, rho0, &d, n2, n3i, xi))
            }
            (0, 1, false) => {
                let params = a.binary_cc.first().map(|r| &r.model_record);
                Ok(self.helmholtz_energy_density_cc_analytic(params, t, rho0, &d, n2, n3i, xi))
            }
            (1, 1, false) => {
                let params_ab = a.binary_ab.first().map(|r| &r.model_record);
                let params_cc = a.binary_cc.first().map(|r| &r.model_record);
                Ok(
                    self.helmholtz_energy_density_ab_analytic(params_ab, t, rho0, &d, n2, n3i, xi)
                        + self.helmholtz_energy_density_cc_analytic(
                            params_cc, t, rho0, &d, n2, n3i, xi,
                        ),
                )
            }
            _ => {
                let mut x = DVector::from_element(a.sites_a.len() + a.sites_b.len(), 0.2);
                let (assoc_comp_ab, n_ab): (Vec<_>, Vec<_>) = a
                    .sites_a
                    .iter()
                    .chain(a.sites_b.iter())
                    .map(|s| (s.assoc_comp, s.n))
                    .unzip();
                let rhoab = Array2::from_shape_fn((x.len(), n3i.len()), |(i, j)| {
                    rho0[(assoc_comp_ab[i], j)] * n_ab[i]
                });
                rhoab
                    .axis_iter(Axis(1))
                    .zip(n2.iter())
                    .zip(n3i.iter())
                    .zip(xi.iter())
                    .map(|(((rho, &n2), &n3i), &xi)| {
                        let [delta_ab, delta_cc] = self.association.yu_wu_association_strength(
                            self.association_parameters,
                            self.model,
                            t,
                            &d,
                            n2,
                            n3i,
                            xi,
                        );
                        let rho = DVector::from(rho.to_vec());
                        self.association.helmholtz_energy_density_cross_association(
                            &rho,
                            &delta_ab,
                            &delta_cc,
                            Some(&mut x),
                        )
                    })
                    .collect()
            }
        }
    }

    #[expect(clippy::too_many_arguments)]
    fn helmholtz_energy_density_ab_analytic<N: DualNum<f64> + Copy, S: Data<Elem = N>>(
        &self,
        parameters: Option<&A::Record>,
        temperature: N,
        rho0: &Array2<N>,
        diameter: &DVector<N>,
        n2: &ArrayBase<S, Ix1>,
        n3i: &Array1<N>,
        xi: &Array1<N>,
    ) -> Array1<N> {
        let a = &self.association_parameters;
        let Some(par) = parameters else {
            return Array1::zeros(xi.len());
        };

        // site densities
        let i = a.sites_a[0].assoc_comp;
        let j = a.sites_b[0].assoc_comp;
        let rhoa = &rho0.index_axis(Axis(0), i) * a.sites_a[0].n;
        let rhob = &rho0.index_axis(Axis(0), j) * a.sites_b[0].n;

        // association strength
        let di = diameter[i];
        let dj = diameter[j];
        let k = n2 * n3i * (di * dj / (di + dj));
        let delta = (((&k / 18.0 + 0.5) * &k * xi + 1.0) * n3i)
            * self.model.association_strength_ij(temperature, 0, 0, par);

        // no cross association, two association sites
        let aux = &delta * (&rhob - &rhoa) + 1.0;
        let xa = ((&aux * &aux + &delta * &rhoa * 4.0).map(N::sqrt) + &aux).map(N::recip) * 2.0;
        let aux = -aux + 2.0;
        let xb = ((&aux * &aux + delta * &rhob * 4.0).map(N::sqrt) + aux).map(N::recip) * 2.0;

        let f = |x: N| x.ln() - x * 0.5 + 0.5;
        rhoa * xa.mapv(f) + rhob * xb.mapv(f)
    }

    #[expect(clippy::too_many_arguments)]
    fn helmholtz_energy_density_cc_analytic<N: DualNum<f64> + Copy, S: Data<Elem = N>>(
        &self,
        parameters: Option<&A::Record>,
        temperature: N,
        rho0: &Array2<N>,
        diameter: &DVector<N>,
        n2: &ArrayBase<S, Ix1>,
        n3i: &Array1<N>,
        xi: &Array1<N>,
    ) -> Array1<N> {
        let a = &self.association_parameters;
        let Some(par) = parameters else {
            return Array1::zeros(xi.len());
        };

        // site densities
        let i = a.sites_c[0].assoc_comp;
        let rhoc = &rho0.index_axis(Axis(0), i) * a.sites_c[0].n;

        // association strength
        let di = diameter[i];
        let k = n2 * n3i * (di * 0.5);
        let delta = (((&k / 18.0 + 0.5) * &k * xi + 1.0) * n3i)
            * self.model.association_strength_ij(temperature, 0, 0, par);

        // no cross association, two association sites
        let xc = ((delta * 4.0 * &rhoc + 1.0).map(N::sqrt) + 1.0).map(N::recip) * 2.0;

        let f = |x: N| x.ln() - x * 0.5 + 0.5;
        rhoc * xc.mapv(f)
    }
}