atmosim 0.1.0

//! Handles the search of the best maxcap by offloading it to argmin framework.
//! You probably only want to touch this if there's a bug, or you want to change the algorythm
//! Also 65% of code here is just converting normal data to matrixes and back. I hate it.

use core::{fmt::Display, time::Duration};

use argmin::{
    core::{CostFunction, Error, Executor},
    solver::particleswarm::ParticleSwarm,
};
use enum_map::Enum;
use log::{debug, info, warn};
use micromath::F32Ext;
use nalgebra::{ArrayStorage, Const, Matrix};
use thiserror::Error;

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

use crate::{
    Atmosim,
    config::R,
    gases::{GAS_COUNT, Gas, GasFlags, GasMixture, Moles, gases},
    maxcap::{GasContainer, GasContainerPrototype, GasContainerState},
};

const EXP_LIMIT: f32 = 6.;
const DUPLICATE_GAS_PENALTY: f32 = 0.1;
const EACH_GAS_PENALTY: f32 = 0.01;

#[derive(Debug, Error)]
#[error("Optimizer gave no plausible results")]
struct OptimizationError;

#[derive(Debug, Clone, Copy)]
pub enum OptimizeFor {
    BiggerExplosion,
    SlowerFuse,
    FasterFuse,
}

#[derive(Debug, Clone, Copy)]
pub struct Optimizer<'a> {
    pub engine: &'a Atmosim,
    pub allowed_gases0: GasFlags,
    // The reason I'm not using Range<T> is because it for whatever reason doesn't implement Copy
    pub temp_range0: (f32, f32),
    pub allowed_gases1: GasFlags,
    pub temp_range1: (f32, f32),
    pub optimize_for: OptimizeFor,
    pub tick_range: (usize, usize),
    pub min_explosion_radius: f32,
    pub minimal_part: f32,
    pub pso_particles: usize,
    pub timeout: Duration,
}

impl<'a> Optimizer<'a> {
    #[must_use]
    pub fn from_engine(engine: &'a Atmosim) -> Self {
        Self {
            engine,
            allowed_gases0: gases! {
                Gas::Ammonia | Gas::CarbonDioxide | Gas::Frezon |
                Gas::Nitrogen | Gas::NitrousOxide | Gas::Plasma |
                Gas::Tritium | Gas::WaterVapor => true,
                _ => false,
            },
            temp_range0: (71.2, 595.1),
            allowed_gases1: gases! {
                Gas::Oxygen => true,
                _ => false,
            },
            temp_range1: (293.15, 293.15),
            optimize_for: OptimizeFor::BiggerExplosion,
            tick_range: (0, 7200), // one hour
            min_explosion_radius: 4.,
            minimal_part: 0.05,
            pso_particles: 7777, // unironically good
            timeout: Duration::from_secs(3),
        }
    }
}

// Matrix GAS_COUNT + 1 (because temperature) tall and 2 wide (because mixing 2 containers)
type GasMatrix =
    Matrix<f32, Const<{ GAS_COUNT + 1 }>, Const<2>, ArrayStorage<f32, { GAS_COUNT + 1 }, 2>>;

impl CostFunction for Optimizer<'_> {
    type Param = GasMatrix;
    type Output = f32;

    fn cost(&self, param: &Self::Param) -> Result<Self::Output, Error> {
        let (mut mix0, mut mix1) = self.matrix_to_mixes(param);

        // Check if both mixes are stable
        if mix0.react() || mix1.react() {
            return Ok(f32::MAX); // they're not
        }

        // Check if two mixes are of the same order
        let total0 = mix0.total_moles();
        let total1 = mix1.total_moles();
        let control = (total0 + total1) * self.minimal_part;
        if total0 < control || total1 < control {
            return Ok(f32::MAX); // they're not
        }

        // Penalize mixes so we use smallest number of gases and no duplicates
        let penalty = {
            let each_penalty = mix0
                .moles()
                .values()
                .zip(mix1.moles().values())
                .filter(|&(&first, &second)| first > 0. || second > 0.)
                .count() as f32
                * EACH_GAS_PENALTY;

            let duplicate_penalty = mix0
                .moles()
                .values()
                .zip(mix1.moles().values())
                .filter(|&(&first, &second)| first > 0. && second > 0.)
                .count() as f32
                * DUPLICATE_GAS_PENALTY;

            each_penalty + duplicate_penalty
        };

        let (radius, ticks) = calculate_maxcap(mix0, mix1, self.tick_range.1);

        // Kinda cursed but whatever
        Ok(if radius > 0. {
            (match (self.tick_range.0, self.min_explosion_radius) {
                (min_ticks, _) if ticks < min_ticks => f32::MAX, // Exploded too soon
                (_, min_radius) if radius < min_radius => f32::MAX, // Explosion too small
                _ => match self.optimize_for {
                    OptimizeFor::BiggerExplosion => -radius,
                    OptimizeFor::SlowerFuse => -(ticks as f32),
                    OptimizeFor::FasterFuse => ticks as f32,
                },
            } + penalty) // negative is good
        } else {
            f32::MAX
        })
    }
}

/// Returns a tuple of explosion radius and ticks to achieve it
fn calculate_maxcap(mix0: GasMixture, mix1: GasMixture, tick_limit: usize) -> (f32, usize) {
    // TODO: canisters should work a bit differently bc they're connected to pipes
    let mut container = GasContainer::new(mix0 + mix1);
    container.step_n(tick_limit);

    (
        match container.state() {
            GasContainerState::Exploded { radius } => radius,
            _ => 0.,
        },
        container.ticks(),
    )
}

// magic, do not touch
fn bound(gases: GasFlags, i: usize, positive: bool) -> f32 {
    if gases[Gas::from_usize(i)] {
        if positive { EXP_LIMIT } else { -EXP_LIMIT }
    } else {
        f32::MIN
    }
}

impl Optimizer<'_> {
    pub fn run(&self) -> Result<MaxcapData, Error> {
        let tr0 = self.temp_range0;
        let tr1 = self.temp_range1;
        if tr0.0 > tr0.1 || tr1.0 > tr1.1 {
            Err(OptimizationError)?; // temp min is greater than temp max for some mix
        }

        // Construct a matrix the optimizer understands
        let lower_bounds = {
            GasMatrix::from_fn(|i, j| match (i, j) {
                (..GAS_COUNT, 0) => bound(self.allowed_gases0, i, false),
                (..GAS_COUNT, 1) => bound(self.allowed_gases1, i, false),
                (GAS_COUNT, 0) => self.temp_range0.0,
                (GAS_COUNT, 1) => self.temp_range1.0,
                _ => unreachable!(),
            })
        };

        let upper_bounds = {
            GasMatrix::from_fn(|i, j| match (i, j) {
                (..GAS_COUNT, 0) => bound(self.allowed_gases0, i, true),
                (..GAS_COUNT, 1) => bound(self.allowed_gases1, i, true),
                (GAS_COUNT, 0) => self.temp_range0.1,
                (GAS_COUNT, 1) => self.temp_range1.1,
                _ => unreachable!(),
            })
        };

        let solver = ParticleSwarm::new((lower_bounds, upper_bounds), self.pso_particles);
        let res = Executor::new(*self, solver).timeout(self.timeout).run()?;
        debug!("{res}");
        let cur = res.state.iter * self.pso_particles as u64;
        let best = res.state.last_best_iter * self.pso_particles as u64;
        if best * 3 / 2 > cur {
            warn!(
                "Processed {cur} maxcaps, but the best result was only found recently at {best}. Consider trying again with more time."
            );
        }

        // Now let's make the data readable again...
        let (mix0, mix1) = match res.state.best_individual {
            Some(particle) if particle.cost.is_sign_negative() => {
                self.matrix_to_mixes(&particle.position)
            }
            _ => Err(OptimizationError)?, // no results
        };

        // This is shitcode but let's now simulate this one again so we know its radius and ticks
        let (radius, ticks) = calculate_maxcap(mix0, mix1, self.tick_range.1);

        info!("The result mixes down to {}K.", (mix0 + mix1).temperature());

        Ok(MaxcapData {
            moles0: mix0.moles(),
            temp0: mix0.temperature(),
            moles1: mix1.moles(),
            temp1: mix1.temperature(),
            radius,
            ticks,
            container_proto: self.engine.container(),
        })
    }

    /// Note that gases stored in matrix are exponents
    fn matrix_to_mixes<'a>(&'a self, matrix: &GasMatrix) -> (GasMixture<'a>, GasMixture<'a>) {
        let mut mix0 = self.row_to_mix(&matrix.data.0[0]);
        let mut mix1 = self.row_to_mix(&matrix.data.0[1]);

        // Normalize everything because argmin spews random numbers out
        let ratio = {
            let container = GasContainer::new(mix0 + mix1);
            self.engine.container().normal_pressure / container.pressure()
        };

        mix0.modify_moles_each(|mole| mole * ratio);
        mix1.modify_moles_each(|mole| mole * ratio);

        (mix0, mix1)
    }

    fn row_to_mix<'a>(&'a self, row: &[f32]) -> GasMixture<'a> {
        assert_eq!(row.len(), GAS_COUNT + 1); // +1 for temperature
        let (&temp, slice) = row.split_last().unwrap();
        // exp is actually giga expensive so micromath
        let mut moles = Moles::from_fn(|i| F32Ext::exp(slice[i.into_usize()]));

        // Cut non-existant values
        let total = moles.values().sum::<f32>();
        for mole in moles.values_mut() {
            if *mole < total * self.minimal_part {
                *mole = 0.;
            }
        }

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

#[derive(Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct MaxcapData {
    pub moles0: Moles,
    pub temp0: f32,
    pub moles1: Moles,
    pub temp1: f32,
    pub radius: f32,
    pub ticks: usize,
    pub container_proto: GasContainerPrototype,
}

impl Display for MaxcapData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let moles0_sum = self.moles0.values().sum::<f32>();
        let moles1_sum = self.moles1.values().sum::<f32>();
        let pressure0 = moles0_sum * R * self.temp0 / self.container_proto.volume;

        writeln!(
            f,
            "- MIX 1: {}K, fill up to {} kPa -",
            self.temp0, pressure0
        )?;
        for (gas, amount) in self.moles0 {
            if amount > f32::MIN_POSITIVE {
                writeln!(
                    f,
                    "{gas:?}: {}% ({amount} moles)",
                    amount / moles0_sum * 100.
                )?;
            }
        }

        writeln!(
            f,
            "- MIX 2: {}K, fill up to {} kPa -",
            self.temp1, self.container_proto.normal_pressure
        )?;
        for (gas, amount) in self.moles1 {
            if amount > f32::MIN_POSITIVE {
                writeln!(
                    f,
                    "{gas:?}: {}% ({amount} moles)",
                    amount / moles1_sum * 100.
                )?;
            }
        }

        write!(
            f,
            "- Explodes after {} ticks with {} tile radius -",
            self.ticks, self.radius
        )
    }
}

// P.S. Yeah, maybe Ilya was right after all.
// At this point it could've been easier to write the optimizer from scratch.