//! Functionalities for the calculation of pair correlation functions.

use crate::functional::HelmholtzEnergyFunctional;

use crate::profile::MAX_POTENTIAL;

use crate::solver::DFTSolver;

use crate::{Axis, DFTProfile, Grid};

use feos_core::{Contributions, FeosResult, ReferenceSystem, State};

use ndarray::prelude::*;

use quantity::{Energy, Length};

use std::ops::Deref;

/// The underlying pair potential, that the Helmholtz energy functional

/// models.

pub trait PairPotential {

/// Return the pair potential of particle i with all other particles.

fn pair_potential(&self, i: usize, r: &Array1<f64>, temperature: f64) -> Array2<f64>;

}

impl<C: Deref<Target = T>, T: PairPotential> PairPotential for C {

fn pair_potential(&self, i: usize, r: &Array1<f64>, temperature: f64) -> Array2<f64> {

T::pair_potential(self, i, r, temperature)

}

}

/// Density profile and properties of a test particle system.

pub struct PairCorrelation<F> {

pub profile: DFTProfile<Ix1, F>,

pub pair_correlation_function: Option<Array2<f64>>,

pub self_solvation_free_energy: Option<Energy>,

pub structure_factor: Option<f64>,

}

impl<F: HelmholtzEnergyFunctional + PairPotential> PairCorrelation<F> {

pub fn new(bulk: &State<F>, test_particle: usize, n_grid: usize, width: Length) -> Self {

let dft = &bulk.eos;

// generate grid

let axis = Axis::new_spherical(n_grid, width);

// calculate external potential

let t = bulk.temperature.to_reduced();

let mut external_potential = dft.pair_potential(test_particle, &axis.grid, t) / t;

external_potential.map_inplace(|x| {

if *x > MAX_POTENTIAL {

*x = MAX_POTENTIAL

}

});

let grid = Grid::Spherical(axis);

Self {

profile: DFTProfile::new(grid, bulk, Some(external_potential), None, Some(1)),

pair_correlation_function: None,

self_solvation_free_energy: None,

structure_factor: None,

}

}

pub fn solve_inplace(&mut self, solver: Option<&DFTSolver>, debug: bool) -> FeosResult<()> {

// Solve the profile

self.profile.solve(solver, debug)?;

// calculate pair correlation function

let partial_density = self.profile.bulk.partial_density();

self.pair_correlation_function = Some(Array::from_shape_fn(

self.profile.density.raw_dim(),

|(i, j)| (self.profile.density.get((i, j)) / partial_density.get(i)).into_value(),

));

// calculate self solvation free energy

self.self_solvation_free_energy = Some(self.profile.integrate(

&(self.profile.grand_potential_density()?

+ self.profile.bulk.pressure(Contributions::Total)),

));

// calculate structure factor

self.structure_factor = Some(

(self.profile.total_moles() - self.profile.bulk.density * self.profile.volume())

.to_reduced()

+ 1.0,

);

Ok(())

}

pub fn solve(mut self, solver: Option<&DFTSolver>) -> FeosResult<Self> {

self.solve_inplace(solver, false)?;

Ok(self)

}

}