atmosim 0.1.0

//! Handles gas types. If you want to add more gases, just put them into Gas enum
//! and add its default properties into impl Default for `GasProperties`

use core::{fmt::Display, ops::Add, ptr};

pub use enum_map::enum_map as gases;
use enum_map::{Enum, EnumMap};
use strum::{Display, EnumIter, EnumString};

#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

use crate::{
    Atmosim,
    config::{MIN_MOLES, TCMB},
};

#[derive(Debug, Enum, Clone, Copy, EnumIter, EnumString, Display, PartialEq, Eq)]
#[strum(ascii_case_insensitive)]
#[strum(serialize_all = "snake_case")]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize, Hash))]
pub enum Gas {
    // common (wizden)
    #[strum(serialize = "O2")]
    Oxygen,
    #[strum(serialize = "N2")]
    Nitrogen,
    #[strum(serialize = "CO2")]
    CarbonDioxide,
    #[strum(serialize = "PLS")]
    Plasma,
    #[strum(serialize = "TRT")]
    Tritium,
    #[strum(serialize = "H2O")]
    WaterVapor,
    #[strum(serialize = "NH3")]
    Ammonia,
    #[strum(serialize = "N2O")]
    NitrousOxide,
    #[strum(serialize = "FRZ")]
    Frezon,
    // funky/goob/mono
    #[strum(serialize = "BZ")]
    BZ,
    #[strum(serialize = "HEA")]
    Healium,
    #[strum(serialize = "NTR")]
    Nitrium,
    #[strum(serialize = "PLX")]
    Pluoxium,
    // funky only
    #[strum(serialize = "H2")]
    Hydrogen, // this is literally tritium dumbfucks
    #[strum(serialize = "HNB")]
    HyperNoblium,
    #[strum(serialize = "PNR")]
    ProtoNitrate,
    #[strum(serialize = "ZK")]
    Zauker,
    #[strum(serialize = "HAL")]
    Halon,
    #[strum(serialize = "HE")]
    Helium, // there is healium already holy shit
    #[strum(serialize = "ANB")]
    AntiNoblium,
    // klovn
    #[strum(serialize = "ZIP")]
    Zipion,
    #[strum(serialize = "ARG")]
    Argon,
}

pub(crate) const GAS_COUNT: usize = Gas::LENGTH;

#[derive(Debug, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct MoleProperties {
    pub specific_heat: f32,
    pub mass: f32,
}

impl Default for MoleProperties {
    fn default() -> Self {
        Self {
            specific_heat: 0.,
            mass: 1.,
        }
    }
}

pub type Moles = EnumMap<Gas, f32>;
pub type GasFlags = EnumMap<Gas, bool>;

#[cfg(feature = "serde")]
use serde_stuff::{GasPropertiesHashMap, MolesHashMap};

#[derive(Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(from = "GasPropertiesHashMap"))]
pub struct GasProperties(pub EnumMap<Gas, MoleProperties>);

#[derive(Debug, Default, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(from = "MolesHashMap"))]
pub struct MolesWrapper(pub Moles);

#[cfg(feature = "serde")]
mod serde_stuff {
    use std::collections::HashMap;

    use crate::gases::{Gas, GasProperties, MoleProperties, MolesWrapper};

    pub(super) type GasPropertiesHashMap = HashMap<Gas, MoleProperties>;

    impl From<GasPropertiesHashMap> for GasProperties {
        fn from(value: GasPropertiesHashMap) -> Self {
            Self(Self::default().0.into_iter().chain(value).collect())
        }
    }

    pub(super) type MolesHashMap = HashMap<Gas, f32>;

    impl From<MolesHashMap> for MolesWrapper {
        fn from(value: MolesHashMap) -> Self {
            Self(Self::default().0.into_iter().chain(value).collect())
        }
    }
}

impl Default for GasProperties {
    fn default() -> Self {
        Self(gases! {
            Gas::Oxygen => MoleProperties { specific_heat: 20., mass: 32. },
            Gas::Nitrogen => MoleProperties { specific_heat: 30., mass: 28. },
            Gas::CarbonDioxide => MoleProperties { specific_heat: 30., mass: 44. },
            Gas::Plasma => MoleProperties { specific_heat: 200., mass: 120. },
            Gas::Tritium => MoleProperties { specific_heat: 10., mass: 6. },
            Gas::WaterVapor => MoleProperties { specific_heat: 40., mass: 18. },
            Gas::Ammonia => MoleProperties { specific_heat: 20., mass: 44. },
            Gas::NitrousOxide => MoleProperties { specific_heat: 40., mass: 44. },
            Gas::Frezon => MoleProperties { specific_heat: 600., mass: 50. },

            Gas::BZ => MoleProperties { specific_heat: 20., mass: 12. },
            Gas::Healium => MoleProperties { specific_heat: 10., mass: 15. },
            Gas::Nitrium => MoleProperties { specific_heat: 10., mass: 8. },
            Gas::Pluoxium => MoleProperties { specific_heat: 80., mass: 32. },

            Gas::Hydrogen => MoleProperties { specific_heat: 15., mass: 32. },
            Gas::HyperNoblium => MoleProperties { specific_heat: 2000., mass: 8. },
            Gas::ProtoNitrate => MoleProperties { specific_heat: 30., mass: 120. },
            Gas::Zauker => MoleProperties { specific_heat: 350., mass: 110. },
            Gas::Halon => MoleProperties { specific_heat: 1.4, mass: 150. },
            Gas::Helium => MoleProperties { specific_heat: 15., mass: 4. },
            Gas::AntiNoblium => MoleProperties { specific_heat: 1., mass: 200. },

           Gas::Zipion => MoleProperties { specific_heat: 5., mass: 20. },
           Gas::Argon => MoleProperties { specific_heat: 25., mass: 40. },
        })
    }
}

#[derive(Debug, Clone, Copy)]
/// An abstract gas mix
pub struct GasMixture<'a> {
    // these two shouldn't be changed directly hence private
    temp: f32,
    moles: Moles,
    // This is perfectly derivable from moles but we cache it for performance
    pub(crate) cached_heat_capacity: f32,
    pub(crate) engine: &'a Atmosim,
}

impl Display for GasMixture<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        for (gas, amount) in &self.moles {
            if *amount > 0. {
                writeln!(f, "{gas:?}: {amount}")?;
            }
        }

        write!(f, "T = {}K", self.temp)
    }
}

impl GasMixture<'_> {
    #[must_use]
    pub fn moles(&self) -> Moles {
        self.moles
    }

    #[must_use]
    pub fn temperature(&self) -> f32 {
        self.temp
    }

    pub(crate) fn set_temperature(&mut self, value: f32) {
        assert!(!value.is_nan());
        self.temp = value.clamp(TCMB, self.engine.game.atmospherics.tmax);
    }

    #[must_use]
    pub fn total_moles(&self) -> f32 {
        self.moles.values().sum()
    }

    #[must_use]
    pub fn get_moles(&self, gas: Gas) -> f32 {
        self.moles[gas]
    }

    pub(crate) fn set_moles(&mut self, gas: Gas, quantity: f32) {
        assert!(quantity.is_finite());
        assert!(!quantity.is_sign_negative());

        let old = self.moles[gas];
        self.moles[gas] = quantity;

        // It should be only ever the job of this function to handle heat capacity.
        // All other functions changing moles should call it.
        let changed = quantity - old;
        self.cached_heat_capacity = {
            let mhc = self.engine.game.atmospherics.minimum_heat_capacity;
            let old_heat_capacity = self.cached_heat_capacity;

            if old_heat_capacity <= mhc {
                // shit's fucked up, recalculate
                GasMixture::calculate_heat_capacity(self.moles, self.engine)
            } else {
                (old_heat_capacity + changed * self.engine.game.gas_properties.0[gas].specific_heat)
                    .max(mhc)
            }
        }
    }

    pub(crate) fn adjust_moles(&mut self, gas: Gas, quantity: f32) {
        let old = self.moles[gas];
        self.set_moles(gas, (old + quantity).max(0.));
    }

    pub(crate) fn modify_moles_each(&mut self, how: impl Fn(f32) -> f32) {
        for (gas, moles) in self.moles {
            self.set_moles(gas, how(moles));
        }
    }

    // The game also returns a mixture of removed gases but I think it's not actually needed
    pub(crate) fn remove_ratio(&mut self, mut ratio: f32) {
        match ratio {
            ..=0. => return self.modify_moles_each(|_| 0.),
            1.0.. => ratio = 1.,
            _ => (),
        }

        let factor = 1. - ratio;
        self.modify_moles_each(|mut amount| {
            amount *= factor;
            if amount < MIN_MOLES || amount.is_nan() {
                amount = 0.;
            }
            amount
        });
    }

    // this is a hotspot in modern maxcaps
    pub(crate) fn remove(&mut self, amount: f32) {
        self.remove_ratio(amount / self.total_moles());
    }

    // this is a hotspot in the game so we cache it
    #[must_use]
    pub fn get_heat_capacity(&self) -> f32 {
        self.cached_heat_capacity
    }

    #[must_use]
    pub fn get_thermal_energy(&self) -> f32 {
        self.temp * self.cached_heat_capacity
    }

    /// If you already have a gas mixture, use `get_heat_capacity` instead
    #[must_use]
    pub fn calculate_heat_capacity(moles: Moles, engine: &Atmosim) -> f32 {
        (
            moles
                .into_values()
                .zip(engine.game.gas_properties.0.values())
                // we'll be having a lot of zeroes so let's just filter them out
                .filter(|&(amt, _)| amt != 0.)
                .map(|(amt, mp)| amt * mp.specific_heat)
                .sum::<f32>()
            // for some reason it wants turbofish
        )
        .max(engine.game.atmospherics.minimum_heat_capacity)
    }

    // might want to cache this too but we'll see
    #[must_use]
    pub(crate) fn get_mass(&self) -> f32 {
        self.moles
            .into_values()
            .zip(self.engine.game.gas_properties.0.values())
            .filter(|&(amt, _)| amt != 0.)
            .map(|(amt, mp)| amt * mp.mass)
            .sum()
    }
}

impl<'a> Add for GasMixture<'a> {
    type Output = GasMixture<'a>;
    fn add(self, rhs: Self) -> Self::Output {
        assert!(
            ptr::eq(self.engine, rhs.engine),
            "Tried to add mixtures from different engines"
        );

        let moles = self
            .moles
            .iter()
            .zip(rhs.moles)
            .map(|((gas, mole0), (_, mole1))| (gas, mole0 + mole1))
            .collect();

        let temperature = (self.cached_heat_capacity * self.temp
            + rhs.cached_heat_capacity * rhs.temp)
            / (self.cached_heat_capacity + rhs.cached_heat_capacity);

        self.engine.create_mixture(moles, temperature)
    }
}

impl<'a> Atmosim {
    #[must_use]
    pub fn create_mixture(&'a self, moles: Moles, temperature: f32) -> GasMixture<'a> {
        GasMixture {
            temp: temperature,
            moles,
            cached_heat_capacity: GasMixture::calculate_heat_capacity(moles, self),
            engine: self,
        }
    }
}