use super::{BarkerHenderson, Perturbation, WeeksChandlerAndersen};

use crate::hard_sphere::{HardSphereProperties, MonomerShape};

use feos_core::parameter::Parameters;

use nalgebra::{DMatrix, DVector, SMatrix, matrix, stack, vector};

use num_dual::DualNum;

use serde::{Deserialize, Serialize};

/// uv-theory parameters for a pure substance

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]

pub struct UVTheoryRecord {

rep: f64,

att: f64,

sigma: f64,

epsilon_k: f64,

}

impl UVTheoryRecord {

/// Single substance record for uv-theory

pub fn new(rep: f64, att: f64, sigma: f64, epsilon_k: f64) -> Self {

Self {

rep,

att,

sigma,

epsilon_k,

}

}

}

/// Constants for BH temperature dependent HS diameter.

const CD_BH: SMatrix<f64, 4, 3> = matrix![

0.0, 1.09360455168912E-02, 0.0;

-2.00897880971934E-01, -1.27074910870683E-02, 0.0;

1.40422470174053E-02, 7.35946850956932E-02, 1.28463973950737E-02;

3.71527116894441E-03, 5.05384813757953E-03, 4.91003312452622E-02];

#[inline]

pub fn mie_prefactor<D: DualNum<f64> + Copy>(rep: D, att: D) -> D {

rep / (rep - att) * (rep / att).powd(att / (rep - att))

}

#[inline]

pub fn mean_field_constant<D: DualNum<f64> + Copy>(rep: D, att: D, x: D) -> D {

mie_prefactor(rep, att) * (x.powd(-att + 3.0) / (att - 3.0) - x.powd(-rep + 3.0) / (rep - 3.0))

}

/// Parameters for all substances for uv-theory equation of state and Helmholtz energy functional

pub type UVTheoryParameters = Parameters<UVTheoryRecord, f64, ()>;

/// Parameters for all substances for uv-theory equation of state and Helmholtz energy functional

#[derive(Debug, Clone)]

pub struct UVTheoryPars {

pub perturbation: Perturbation,

pub rep: DVector<f64>,

pub att: DVector<f64>,

pub sigma: DVector<f64>,

pub epsilon_k: DVector<f64>,

pub rep_ij: DMatrix<f64>,

pub att_ij: DMatrix<f64>,

pub sigma_ij: DMatrix<f64>,

pub eps_k_ij: DMatrix<f64>,

pub cd_bh_pure: Vec<[f64; 5]>,

pub cd_bh_binary: DMatrix<[f64; 5]>,

}

impl UVTheoryPars {

pub fn new(parameters: &UVTheoryParameters, perturbation: Perturbation) -> Self {

let n = parameters.pure.len();

let [rep, att, sigma, epsilon_k] =

parameters.collate(|pr| [pr.rep, pr.att, pr.sigma, pr.epsilon_k]);

let mut rep_ij = DMatrix::zeros(n, n);

let mut att_ij = DMatrix::zeros(n, n);

let mut sigma_ij = DMatrix::zeros(n, n);

let mut eps_k_ij = DMatrix::zeros(n, n);

let [k_ij] = parameters.collate_binary(|&br| [br]);

for i in 0..n {

rep_ij[(i, i)] = rep[i];

att_ij[(i, i)] = att[i];

sigma_ij[(i, i)] = sigma[i];

eps_k_ij[(i, i)] = epsilon_k[i];

for j in i + 1..n {

rep_ij[(i, j)] = (rep[i] * rep[j]).sqrt();

rep_ij[(j, i)] = rep_ij[(i, j)];

att_ij[(i, j)] = (att[i] * att[j]).sqrt();

att_ij[(j, i)] = att_ij[(i, j)];

sigma_ij[(i, j)] = 0.5 * (sigma[i] + sigma[j]);

sigma_ij[(j, i)] = sigma_ij[(i, j)];

eps_k_ij[(i, j)] = (1.0 - k_ij[(i, j)]) * (epsilon_k[i] * epsilon_k[j]).sqrt();

eps_k_ij[(j, i)] = eps_k_ij[(i, j)];

}

}

// BH temperature dependent HS diameter, eq. 21

let cd_bh_pure: Vec<_> = rep.iter().map(|&mi| bh_coefficients(mi, 6.0)).collect();

let cd_bh_binary = DMatrix::from_fn(n, n, |i, j| bh_coefficients(rep_ij[(i, j)], 6.0));

Self {

perturbation,

rep,

att,

sigma,

epsilon_k,

rep_ij,

att_ij,

sigma_ij,

eps_k_ij,

cd_bh_pure,

cd_bh_binary,

}

}

}

#[expect(clippy::toplevel_ref_arg)]

fn bh_coefficients(rep: f64, att: f64) -> [f64; 5] {

let inv_a76 = 1.0 / mean_field_constant(7.0, att, 1.0);

let am6 = mean_field_constant(rep, att, 1.0);

let alpha = 1.0 / am6 - inv_a76;

let c0 = vector![-2.0 * rep / ((att - rep) * mie_prefactor(rep, att))];

let x = stack![c0; CD_BH * vector![1.0, alpha, alpha * alpha]];

x.data.0[0]

}

impl HardSphereProperties for UVTheoryPars {

fn monomer_shape<D: DualNum<f64> + Copy>(&self, _: D) -> MonomerShape<'_, D> {

MonomerShape::Spherical(self.sigma.len())

}

fn hs_diameter<D: DualNum<f64> + Copy>(&self, temperature: D) -> DVector<D> {

match self.perturbation {

Perturbation::BarkerHenderson => BarkerHenderson::diameter_bh(self, temperature),

Perturbation::WeeksChandlerAndersen => {

WeeksChandlerAndersen::diameter_wca(self, temperature)

}

Perturbation::WeeksChandlerAndersenB3 => {

WeeksChandlerAndersen::diameter_wca(self, temperature)

}

}

}

}

#[cfg(test)]

pub mod utils {

use super::*;

use feos_core::parameter::{Identifier, PureRecord};

use std::f64;

pub fn new_simple(rep: f64, att: f64, sigma: f64, epsilon_k: f64) -> UVTheoryParameters {

UVTheoryParameters::new_pure(PureRecord::new(

Default::default(),

0.0,

UVTheoryRecord::new(rep, att, sigma, epsilon_k),

))

.unwrap()

}

pub fn test_parameters(

rep: f64,

att: f64,

sigma: f64,

epsilon: f64,

p: Perturbation,

) -> UVTheoryPars {

let identifier = Identifier::new(Some("1"), None, None, None, None, None);

let model_record = UVTheoryRecord::new(rep, att, sigma, epsilon);

let pr = PureRecord::new(identifier, 1.0, model_record);

UVTheoryPars::new(&UVTheoryParameters::new_pure(pr).unwrap(), p)

}

pub fn test_parameters_mixture(

rep: DVector<f64>,

att: DVector<f64>,

sigma: DVector<f64>,

epsilon: DVector<f64>,

) -> UVTheoryParameters {

let identifier = Identifier::new(Some("1"), None, None, None, None, None);

let model_record = UVTheoryRecord::new(rep[0], att[0], sigma[0], epsilon[0]);

let pr1 = PureRecord::new(identifier, 1.0, model_record);

//

let identifier2 = Identifier::new(Some("1"), None, None, None, None, None);

let model_record2 = UVTheoryRecord::new(rep[1], att[1], sigma[1], epsilon[1]);

let pr2 = PureRecord::new(identifier2, 1.0, model_record2);

UVTheoryParameters::new_binary([pr1, pr2], None, vec![]).unwrap()

}

pub fn methane_parameters(rep: f64, att: f64) -> UVTheoryParameters {

let identifier = Identifier::new(Some("1"), None, None, None, None, None);

let model_record = UVTheoryRecord::new(rep, att, 3.7039, 150.03);

let pr = PureRecord::new(identifier, 1.0, model_record);

UVTheoryParameters::new_pure(pr).unwrap()

}

}