rustytorch/demos/server/src/state.rs

//! Simulation state management
//!
//! This module manages the state of hemodynamics simulations including
//! model instances, training progress, and performance metrics.

use std::time::{Duration, Instant};

use rtx_hemodynamics::config::PinnConfig;
use rtx_hemodynamics::inference::InferenceEngine;
use rtx_hemodynamics::network::VesselPinn;
use rtx_hemodynamics::training::Trainer;
use rtx_hemodynamics::vessel::VesselSdf;
use rtx_hemodynamics_shared::geometry::VesselGeometry;
use rtx_hemodynamics_shared::ipc::{PerformanceMetrics, SimulationState};

/// Handle to an active simulation
#[derive(Debug)]
pub struct SimulationHandle {
    /// Vessel geometry
    geometry: VesselGeometry,
    /// Vessel SDF wrapper
    vessel: VesselSdf,
    /// PINN model
    model: VesselPinn,
    /// Inference engine (if trained)
    inference: Option<InferenceEngine>,
    /// Trainer instance
    trainer: Trainer,
    /// Whether the model has been trained
    trained: bool,
    /// Creation time
    created_at: Instant,
    /// Last inference time
    last_inference_at: Option<Instant>,
    /// Total inference count
    inference_count: u64,
    /// Total inference time
    total_inference_time: Duration,
}

impl SimulationHandle {
    /// Creates a new simulation handle
    pub fn new(geometry: VesselGeometry, config: PinnConfig) -> Result<Self, String> {
        let vessel = VesselSdf::new(geometry.clone());
        let model = VesselPinn::new(config.clone());
        let trainer = Trainer::new(config, vessel.clone())?;

        Ok(Self {
            geometry,
            vessel,
            model,
            inference: None,
            trainer,
            trained: false,
            created_at: Instant::now(),
            last_inference_at: None,
            inference_count: 0,
            total_inference_time: Duration::ZERO,
        })
    }

    /// Returns the vessel geometry
    #[must_use]
    pub fn geometry(&self) -> &VesselGeometry {
        &self.geometry
    }

    /// Returns the vessel SDF
    #[must_use]
    pub fn vessel(&self) -> &VesselSdf {
        &self.vessel
    }

    /// Returns the PINN model
    #[must_use]
    pub fn model(&self) -> &VesselPinn {
        &self.model
    }

    /// Returns the trainer
    #[must_use]
    pub fn trainer(&self) -> &Trainer {
        &self.trainer
    }

    /// Returns whether the model has been trained
    #[must_use]
    pub const fn is_trained(&self) -> bool {
        self.trained
    }

    /// Marks the model as trained and creates the inference engine
    pub fn mark_trained(&mut self) {
        self.trained = true;
        self.inference = Some(InferenceEngine::new(
            self.model.clone(),
            self.vessel.clone(),
        ));
    }

    /// Returns the inference engine (if trained)
    #[must_use]
    pub fn inference_engine(&self) -> Option<&InferenceEngine> {
        self.inference.as_ref()
    }

    /// Records an inference operation
    pub fn record_inference(&mut self, duration: Duration) {
        self.inference_count += 1;
        self.total_inference_time += duration;
        self.last_inference_at = Some(Instant::now());
    }

    /// Returns current simulation status
    #[must_use]
    pub fn status(&self) -> SimulationStatus {
        SimulationStatus {
            initialized: true,
            trained: self.trained,
            inference_count: self.inference_count,
            avg_inference_time_ms: if self.inference_count > 0 {
                self.total_inference_time.as_millis() as f64 / self.inference_count as f64
            } else {
                0.0
            },
            uptime_secs: self.created_at.elapsed().as_secs(),
        }
    }

    /// Returns performance metrics
    #[must_use]
    pub fn metrics(&self) -> PerformanceMetrics {
        let avg_inference_ms = if self.inference_count > 0 {
            self.total_inference_time.as_millis() as f64 / self.inference_count as f64
        } else {
            0.0
        };

        PerformanceMetrics::new(
            avg_inference_ms,
            0.0, // GPU utilization not tracked yet
            0,   // GPU memory not tracked yet
        )
    }

    /// Returns the current simulation state
    #[must_use]
    pub fn simulation_state(&self) -> SimulationState {
        SimulationState::new()
    }
}

/// Status of a simulation
#[derive(Debug, Clone, Copy)]
pub struct SimulationStatus {
    /// Whether simulation is initialized
    pub initialized: bool,
    /// Whether model is trained
    pub trained: bool,
    /// Total inference count
    pub inference_count: u64,
    /// Average inference time in milliseconds
    pub avg_inference_time_ms: f64,
    /// Uptime in seconds
    pub uptime_secs: u64,
}

impl Default for SimulationStatus {
    fn default() -> Self {
        Self {
            initialized: false,
            trained: false,
            inference_count: 0,
            avg_inference_time_ms: 0.0,
            uptime_secs: 0,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rtx_hemodynamics_shared::geometry::VesselGeometry;

    fn create_test_handle() -> SimulationHandle {
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();
        let config = PinnConfig::default()
            .with_layers(2)
            .with_hidden_dim(16)
            .with_collocation_points(100)
            .with_boundary_points(50);
        SimulationHandle::new(geometry, config).unwrap()
    }

    #[test]
    fn test_simulation_handle_creation() {
        let handle = create_test_handle();
        assert!(!handle.is_trained());
        assert!(handle.inference_engine().is_none());
    }

    #[test]
    fn test_mark_trained() {
        let mut handle = create_test_handle();
        handle.mark_trained();
        assert!(handle.is_trained());
        assert!(handle.inference_engine().is_some());
    }

    #[test]
    fn test_record_inference() {
        let mut handle = create_test_handle();
        handle.record_inference(Duration::from_millis(10));
        handle.record_inference(Duration::from_millis(20));

        let status = handle.status();
        assert_eq!(status.inference_count, 2);
        assert!((status.avg_inference_time_ms - 15.0).abs() < 1.0);
    }

    #[test]
    fn test_simulation_status() {
        let handle = create_test_handle();
        let status = handle.status();
        assert!(status.initialized);
        assert!(!status.trained);
    }

    #[test]
    fn test_performance_metrics() {
        let mut handle = create_test_handle();
        handle.record_inference(Duration::from_millis(10));

        let metrics = handle.metrics();
        assert!(metrics.inference_time_ms() > 0.0);
    }
}