zhl16 0.1.0

use crate::gas::{Gas, GasMixture};
use crate::quantities::{Frequency, FrequencyUnit, Pressure, PressureUnit, Time, TimeUnit};
use crate::tissue_constants::TISSUE_CONSTANTS;
use core::array::from_fn;
use core::fmt;

const WATER_VAPOR_PRESSURE: Pressure = Pressure::new_unchecked(0.0627, PressureUnit::Bar);

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum BodyModelError {
    NonPositiveTimeStep,
    InvalidGradientFactor,
}

impl fmt::Display for BodyModelError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BodyModelError::NonPositiveTimeStep => f.write_str("time step must be positive"),
            BodyModelError::InvalidGradientFactor => {
                f.write_str("gradient factor must be between 0.0 and 1.0")
            }
        }
    }
}

#[derive(Copy, Clone)]
pub(crate) struct TissueModel {
    gases: [TissueModelPerGas; 2],
}

impl TissueModel {
    pub(crate) const fn new(nitrogen: TissueModelPerGas, helium: TissueModelPerGas) -> Self {
        Self {
            gases: [nitrogen, helium],
        }
    }

    #[cfg(test)]
    pub(crate) fn get_gas(&self, gas: Gas) -> TissueModelPerGas {
        self.gases[gas as usize]
    }
}

#[derive(Copy, Clone)]
pub(crate) struct TissueModelPerGas {
    k: Frequency,
    a: Pressure,
    b: f64,
}

impl TissueModelPerGas {
    pub const fn new(half_time: Time, a: Pressure, b: f64) -> Self {
        Self {
            k: Frequency::new_unchecked(
                core::f64::consts::LN_2 / half_time.to(TimeUnit::Seconds),
                FrequencyUnit::Hertz,
            ),
            a,
            b,
        }
    }
    #[cfg(test)]
    pub(crate) fn get_k(&self) -> Frequency {
        self.k
    }
    pub fn get_a(&self) -> Pressure {
        self.a
    }
    pub fn get_b(&self) -> f64 {
        self.b
    }
}

#[derive(Copy, Clone)]
pub struct CompartmentState {
    pressures: [Pressure; 2],
    tissue_model: TissueModel,
}

impl CompartmentState {
    pub(crate) fn new(
        nitrogen_pressures: Pressure,
        helium_pressures: Pressure,
        tissue_model: TissueModel,
    ) -> Self {
        let mut pressures = [Pressure::zero(); 2];
        pressures[Gas::Nitrogen as usize] = nitrogen_pressures;
        pressures[Gas::Helium as usize] = helium_pressures;
        Self {
            pressures,
            tissue_model,
        }
    }

    pub(crate) fn evolve(
        &mut self,
        constants: &TissueModel,
        ambient_pressure: Pressure,
        gas_fractions: &GasMixture,
        dt: Time,
    ) {
        for gas in Gas::ALL {
            let idx = gas as usize;
            let tissue = &constants.gases[idx];
            let pressure = &mut self.pressures[idx];
            let fraction = gas_fractions.inert_fraction(gas);

            let inspired = (ambient_pressure - WATER_VAPOR_PRESSURE) * fraction;
            let rate = (inspired - *pressure) * tissue.k;
            *pressure += rate * dt;
        }
    }

    pub fn get_nitrogen_pressure(&self) -> Pressure {
        self.pressures[Gas::Nitrogen as usize]
    }

    pub fn get_helium_pressure(&self) -> Pressure {
        self.pressures[Gas::Helium as usize]
    }

    pub fn safe_ambient_pressure(&self, gf: f64) -> Result<Pressure, BodyModelError> {
        if !(0.0..=1.0).contains(&gf) {
            return Err(BodyModelError::InvalidGradientFactor);
        }

        let p_n2 = self.pressures[0];
        let p_he = self.pressures[1];
        let p_total = p_n2 + p_he;

        if p_total < Pressure::new_unchecked(0.001, PressureUnit::Bar) {
            return Ok(Pressure::zero());
        }

        let model_n2 = self.tissue_model.gases[0];
        let model_he = self.tissue_model.gases[1];

        let a = (p_n2 / p_total) * model_n2.get_a() + (p_he / p_total) * model_he.get_a();

        let b = (model_n2.get_b() * p_n2 + model_he.get_b() * p_he) / p_total;

        let numerator = p_total - a * gf;
        let denominator = gf / b + (1.0 - gf);

        Ok(numerator / denominator)
    }
}

pub struct BodyState {
    pub(crate) compartments: [CompartmentState; 16],
}

impl BodyState {
    pub fn new(nitrogen_pressure: Pressure, helium_pressure: Pressure) -> Self {
        Self {
            compartments: from_fn(|i| {
                CompartmentState::new(nitrogen_pressure, helium_pressure, TISSUE_CONSTANTS[i])
            }),
        }
    }

    pub fn compartments(&self) -> &[CompartmentState; 16] {
        &self.compartments
    }

    pub fn evolve(
        &mut self,
        ambient_pressure: Pressure,
        gas_mixture: &GasMixture,
        dt: Time,
    ) -> Result<(), BodyModelError> {
        if !(dt.to(TimeUnit::Seconds) > 0.0) {
            return Err(BodyModelError::NonPositiveTimeStep);
        }

        for (i, compartment) in self.compartments.iter_mut().enumerate() {
            let constants = &TISSUE_CONSTANTS[i];
            compartment.evolve(constants, ambient_pressure, gas_mixture, dt);
        }

        Ok(())
    }
}