Entropyk/bindings/fmi/src/model.rs

//! Rust-level FMU instance: parse-once, re-solve-each-step lifecycle.
//!
//! The FMI 2.0 Co-Simulation C ABI in [`crate::fmi2`] wraps this struct. The
//! lifecycle is:
//!
//! 1. [`FmuInstance::new`] — parse the model JSON and the IO-map JSON once.
//!    The model JSON is an Entropyk `ScenarioConfig`; the IO-map JSON declares
//!    which component parameters are PLC inputs and which result fields are
//!    PLC outputs (see [`crate::io_map`]).
//! 2. `set_real` — the host writes input values (ambient temperature, water
//!    temperature, EXV opening, setpoints, ...).
//! 3. [`FmuInstance::do_step`] — apply the inputs to the already-parsed config
//!    and re-solve the steady cycle via [`entropyk_cli::run::run_from_config`].
//! 4. `get_real` — the host reads outputs (COP, capacities, power, pressures).
//!
//! The model JSON is parsed exactly once (at instantiation); each `do_step`
//! only mutates boundary parameters and re-solves. The `System` graph is
//! currently rebuilt every step — warm-start (reusing the previous state
//! vector and the built `System`) is tracked as a TODO in `run_from_config`.

use entropyk_cli::config::ScenarioConfig;
use entropyk_cli::run::{run_from_config, SimulationResult, SimulationStatus};

use crate::io_map::{FmuIoSpec, IoInput, IoOutput};

/// FMI 2.0 status codes (mirrors `fmi2Status`).
#[repr(i32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FmiStatus {
    Ok = 0,
    Warning = 1,
    Discard = 2,
    Error = 3,
    Fatal = 4,
    Pending = 5,
}

impl FmiStatus {
    pub fn as_i32(self) -> i32 {
        self as i32
    }
}

/// One FMU Co-Simulation instance. Independent and self-contained so the
/// exported C ABI can hand the host a raw pointer per instance.
pub struct FmuInstance {
    config: ScenarioConfig,
    io: FmuIoSpec,
    n_inputs: usize,
    n_outputs: usize,
    input_values: Vec<f64>,
    output_values: Vec<f64>,
    last_status: FmiStatus,
    last_error: Option<String>,
}

impl FmuInstance {
    /// Parse the model JSON and the IO-map JSON. Both are bundled inside the
    /// `.fmu` (model under `resources/`, IO-map under `resources/fmu_io.json`).
    pub fn new(config_json: &str, io_json: &str) -> Result<Self, String> {
        let config =
            ScenarioConfig::from_json(config_json).map_err(|e| format!("model JSON: {e}"))?;
        let io: FmuIoSpec =
            serde_json::from_str(io_json).map_err(|e| format!("IO-map JSON: {e}"))?;

        let n_inputs = io.inputs.len();
        let n_outputs = io.outputs.len();

        Ok(Self {
            config,
            io,
            n_inputs,
            n_outputs,
            input_values: vec![0.0; n_inputs],
            output_values: vec![f64::NAN; n_outputs],
            last_status: FmiStatus::Ok,
            last_error: None,
        })
    }

    /// Number of input value references.
    pub fn n_inputs(&self) -> usize {
        self.n_inputs
    }

    /// Number of output value references.
    pub fn n_outputs(&self) -> usize {
        self.n_outputs
    }

    /// Input VRs occupy `[0, n_inputs)`; output VRs occupy
    /// `[n_inputs, n_inputs + n_outputs)`.
    fn is_input_vr(&self, vr: u32) -> bool {
        (vr as usize) < self.n_inputs
    }

    fn output_index(&self, vr: u32) -> Option<usize> {
        let v = vr as usize;
        if v >= self.n_inputs && v < self.n_inputs + self.n_outputs {
            Some(v - self.n_inputs)
        } else {
            None
        }
    }

    /// Write an input value. Outputs are read-only.
    pub fn set_real(&mut self, vr: u32, value: f64) -> FmiStatus {
        if self.is_input_vr(vr) {
            self.input_values[vr as usize] = value;
            FmiStatus::Ok
        } else {
            self.last_status = FmiStatus::Error;
            self.last_error = Some(format!("set_real: VR {vr} is not an input"));
            FmiStatus::Error
        }
    }

    /// Read any value (input or computed output).
    pub fn get_real(&self, vr: u32) -> Result<f64, FmiStatus> {
        let v = vr as usize;
        if v < self.n_inputs {
            return Ok(self.input_values[v]);
        }
        if let Some(idx) = self.output_index(vr) {
            return Ok(self.output_values[idx]);
        }
        Err(FmiStatus::Error)
    }

    /// Apply the current inputs to the config and re-solve the steady cycle.
    pub fn do_step(&mut self) -> FmiStatus {
        // 1. Push input values into the config's component params.
        for (i, input) in self.io.inputs.iter().enumerate() {
            let value = self.input_values[i];
            if !apply_input(&mut self.config, input, value) {
                self.last_status = FmiStatus::Warning;
                self.last_error = Some(format!(
                    "input '{}' -> component '{}' param '{}' not found",
                    input.name, input.component, input.param
                ));
                // Keep going: an unbound input is a warning, not fatal.
            }
        }

        // 2. Re-solve.
        let result: SimulationResult = run_from_config(&self.config);

        // 3. Extract outputs.
        match result.status {
            SimulationStatus::Converged => {
                self.last_status = FmiStatus::Ok;
                self.last_error = None;
            }
            SimulationStatus::Timeout | SimulationStatus::NonConverged => {
                self.last_status = FmiStatus::Discard;
                self.last_error = Some(format!("solver did not converge: {:?}", result.status));
            }
            SimulationStatus::Error => {
                self.last_status = FmiStatus::Error;
                self.last_error = result.error.clone();
            }
        }

        for (idx, out) in self.io.outputs.iter().enumerate() {
            self.output_values[idx] = extract_output(&result, out);
        }

        self.last_status
    }

    /// Enter initialization mode: run one cold solve so outputs are valid
    /// before the host reads them during initialization.
    pub fn enter_init(&mut self) -> FmiStatus {
        self.do_step()
    }

    pub fn last_error(&self) -> Option<&str> {
        self.last_error.as_deref()
    }
}

/// Write `value` into `config.circuits[*].components[name].params[param]`.
fn apply_input(config: &mut ScenarioConfig, input: &IoInput, value: f64) -> bool {
    for circuit in &mut config.circuits {
        for comp in &mut circuit.components {
            if comp.name == input.component {
                comp.params
                    .insert(input.param.clone(), serde_json::Value::from(value));
                return true;
            }
        }
    }
    false
}

/// Pull a single output from the simulation result.
fn extract_output(result: &SimulationResult, out: &IoOutput) -> f64 {
    let perf = result.performance.as_ref();
    match out.kind.as_str() {
        "cop" => perf.and_then(|p| p.cop).unwrap_or(f64::NAN),
        "q_cooling_kw" => perf.and_then(|p| p.q_cooling_kw).unwrap_or(f64::NAN),
        "q_heating_kw" => perf.and_then(|p| p.q_heating_kw).unwrap_or(f64::NAN),
        "compressor_power_kw" => perf.and_then(|p| p.compressor_power_kw).unwrap_or(f64::NAN),
        "pressure_bar" => find_edge(result, out.edge)
            .map(|e| e.pressure_bar)
            .unwrap_or(f64::NAN),
        "enthalpy_kj_kg" => find_edge(result, out.edge)
            .map(|e| e.enthalpy_kj_kg)
            .unwrap_or(f64::NAN),
        "mass_flow_kg_s" => find_edge(result, out.edge)
            .and_then(|e| e.mass_flow_kg_s)
            .unwrap_or(f64::NAN),
        other => {
            let _ = other;
            f64::NAN
        }
    }
}

fn find_edge(
    result: &SimulationResult,
    edge: Option<usize>,
) -> Option<&entropyk_cli::run::StateEntry> {
    let edge = edge?;
    result.state.as_ref()?.iter().find(|e| e.edge == edge)
}