use super::{DataSet, EstimatorError, Phase};

use feos_core::{DensityInitialization, EntropyScaling, EosUnit, State, Residual};

use itertools::izip;

use ndarray::{arr1, Array1};

use quantity::si::{SIArray1, SIUnit};

use std::collections::HashMap;

use std::sync::Arc;

/// Store experimental diffusion data.

#[derive(Clone)]

pub struct Diffusion {

pub target: SIArray1,

temperature: SIArray1,

pressure: SIArray1,

initial_density: Vec<DensityInitialization>,

}

impl Diffusion {

/// Create a new data set for experimental diffusion data.

pub fn new(

target: SIArray1,

temperature: SIArray1,

pressure: SIArray1,

phase: Option<&Vec<Phase>>,

) -> Result<Self, EstimatorError> {

let n = temperature.len();

Ok(Self {

target,

temperature,

pressure,

initial_density: phase.map_or(vec![DensityInitialization::None; n], |phase| {

phase.iter().map(|&p| p.into()).collect()

}),

})

}

/// Return temperature.

pub fn temperature(&self) -> SIArray1 {

self.temperature.clone()

}

/// Return pressure.

pub fn pressure(&self) -> SIArray1 {

self.pressure.clone()

}

}

impl<E: Residual + EntropyScaling> DataSet<E> for Diffusion {

fn target(&self) -> &SIArray1 {

&self.target

}

fn target_str(&self) -> &str {

"diffusion"

}

fn input_str(&self) -> Vec<&str> {

vec!["temperature", "pressure"]

}

fn predict(&self, eos: &Arc<E>) -> Result<SIArray1, EstimatorError> {

let moles = arr1(&[1.0]) * SIUnit::reference_moles();

let ts = self

.temperature

.to_reduced(SIUnit::reference_temperature())

.unwrap();

let ps = self

.pressure

.to_reduced(SIUnit::reference_pressure())

.unwrap();

let res = izip!(&ts, &ps, &self.initial_density)

.map(|(&t, &p, &initial_density)| {

State::new_npt(

eos,

t * SIUnit::reference_temperature(),

p * SIUnit::reference_pressure(),

&moles,

initial_density,

)?

.diffusion()?

.to_reduced(SIUnit::reference_diffusion())

.map_err(EstimatorError::from)

})

.collect::<Result<Vec<f64>, EstimatorError>>();

Ok(Array1::from_vec(res?) * SIUnit::reference_diffusion())

}

fn get_input(&self) -> HashMap<String, SIArray1> {

let mut m = HashMap::with_capacity(1);

m.insert("temperature".to_owned(), self.temperature());

m.insert("pressure".to_owned(), self.pressure());

m

}

}