feos 0.9.5

use crate::saftvrqmie::parameters::SaftVRQMiePars;
use feos_core::FeosResult;
use feos_dft::{FunctionalContribution, WeightFunction, WeightFunctionInfo, WeightFunctionShape};
use nalgebra::DVector;
use ndarray::*;
use num_dual::DualNum;
use std::f64::consts::PI;

pub const N0_CUTOFF: f64 = 1e-9;

#[derive(Clone)]
pub struct NonAddHardSphereFunctional<'a> {
    parameters: &'a SaftVRQMiePars,
}

impl<'a> NonAddHardSphereFunctional<'a> {
    pub fn new(parameters: &'a SaftVRQMiePars) -> Self {
        Self { parameters }
    }
}

impl<'a> FunctionalContribution for NonAddHardSphereFunctional<'a> {
    fn name(&self) -> &'static str {
        "Non-additive hard-sphere functional"
    }

    fn weight_functions<N: DualNum<f64> + Copy>(&self, temperature: N) -> WeightFunctionInfo<N> {
        let p = &self.parameters;
        let r = p.hs_diameter(temperature) * N::from(0.5);
        WeightFunctionInfo::new(DVector::from_fn(r.len(), |i, _| i), false)
            .add(
                WeightFunction::new_scaled(r.clone(), WeightFunctionShape::Delta),
                false,
            )
            .add(
                WeightFunction {
                    prefactor: p.m.map(N::from),
                    kernel_radius: r.clone(),
                    shape: WeightFunctionShape::DeltaVec,
                },
                false,
            )
            .add(
                WeightFunction {
                    prefactor: p.m.map(N::from),
                    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>> {
        let p = &self.parameters;
        // number of components
        let n = p.m.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 d = p.hs_diameter(temperature);
        let mut n2i = Array::zeros(n0i.raw_dim());
        for i in 0..n {
            n2i.index_axis_mut(Axis(0), i)
                .assign(&(&n0i.index_axis(Axis(0), i) * (d[i].powi(2) * (p.m[i] * PI))));
        }
        let mut rho0: Array2<N> = (n2vi
            .iter()
            .map(|n2vi| n2vi * n2vi)
            .fold(Array::zeros(n0i.raw_dim()), |acc, x| acc + x)
            / -(&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()
            .map(|n2v| n2v * n2v)
            .fold(Array::zeros(n3.raw_dim()), |acc, x| acc + x)
            / -(&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());

        // temperature dependent segment radius // calc & store this in struct
        let s_eff_ij =
            Array2::from_shape_fn((n, n), |(i, j)| p.calc_sigma_eff_ij(i, j, temperature));

        // temperature dependent segment radius // calc & store this in struct
        let d_hs_ij = Array2::from_shape_fn((n, n), |(i, j)| {
            p.hs_diameter_ij(i, j, temperature, s_eff_ij[[i, j]])
        });

        // Additive hard-sphere diameter
        let d_hs_add_ij =
            Array2::from_shape_fn((n, n), |(i, j)| (d_hs_ij[[i, i]] + d_hs_ij[[j, j]]) * 0.5);

        Ok(rho0
            .view()
            .into_shape_with_order([n, rho0.len() / n])
            .unwrap()
            .axis_iter(Axis(1))
            .zip(n2.iter())
            .zip(n3i.iter())
            .zip(xi.iter())
            .map(|(((rho0, &n2), &n3i), &xi)| {
                non_additive_hs_energy_density(p, &d_hs_ij, &d_hs_add_ij, &rho0, n2, n3i, xi)
            })
            .collect::<Array1<N>>()
            .into_shape_with_order(n2.raw_dim())
            .unwrap())
    }
}

pub fn non_additive_hs_energy_density<S, N: DualNum<f64> + Copy>(
    parameters: &SaftVRQMiePars,
    d_hs_ij: &Array2<N>,
    d_hs_add_ij: &Array2<N>,
    rho0: &ArrayBase<S, Ix1>,
    n2: N,
    n3i: N,
    xi: N,
) -> N
where
    S: Data<Elem = N>,
{
    // auxiliary variables
    let n = rho0.len();
    let p = parameters;
    let d = Array1::from_shape_fn(n, |i| d_hs_ij[[i, i]]);
    let g_hs_ij = Array2::from_shape_fn((n, n), |(i, j)| {
        let mu = d[i] * d[j] / (d[i] + d[j]);
        n3i + mu * n2 * xi * n3i.powi(2) / 2.0 + (mu * n2 * xi).powi(2) * n3i.powi(3) / 18.0
    });

    // segment densities
    let rho0_s = Array1::from_shape_fn(n, |i| -> N { rho0[i] * p.m[i] });

    Array2::from_shape_fn((n, n), |(i, j)| {
        -rho0_s[i]
            * rho0_s[j]
            * d_hs_add_ij[[i, j]].powi(2)
            * g_hs_ij[[i, j]]
            * (d_hs_add_ij[[i, j]] - d_hs_ij[[i, j]])
            * 2.0
            * PI
    })
    .sum()
}