rustytorch/demos/server/src/service.rs

//! Hemodynamics service implementation
//!
//! This module provides the main service interface for hemodynamics
//! simulation. It handles initialization, inference, and geometry modifications.

use std::sync::Arc;
use std::time::Instant;
use tokio::sync::RwLock;

use rtx_hemodynamics::config::PinnConfig;
use rtx_hemodynamics_shared::fields::{FieldResponse, PressureField, VelocityField, WssField};
use rtx_hemodynamics_shared::geometry::{GeometryModification, Point2D, VesselGeometry};
use rtx_hemodynamics_shared::ipc::{PerformanceMetrics, SimulationState};

use crate::error::{ServerError, ServerResult};
use crate::state::{SimulationHandle, SimulationStatus};

/// Service configuration
#[derive(Debug, Clone)]
pub struct ServiceConfig {
    /// Number of hidden layers
    pub num_layers: usize,
    /// Hidden dimension
    pub hidden_dim: usize,
    /// Fourier feature dimension
    pub fourier_dim: usize,
    /// Number of collocation points for physics loss
    pub num_collocation_points: usize,
    /// Number of boundary points
    pub num_boundary_points: usize,
    /// Inference batch size
    pub inference_batch_size: usize,
}

impl Default for ServiceConfig {
    fn default() -> Self {
        Self {
            num_layers: 4,
            hidden_dim: 128,
            fourier_dim: 32,
            num_collocation_points: 1000,
            num_boundary_points: 250,
            inference_batch_size: 1000,
        }
    }
}

impl ServiceConfig {
    /// Creates a PINN configuration from service config
    fn to_pinn_config(&self) -> PinnConfig {
        PinnConfig::default()
            .with_layers(self.num_layers)
            .with_hidden_dim(self.hidden_dim)
            .with_fourier_dim(self.fourier_dim)
            .with_collocation_points(self.num_collocation_points)
            .with_boundary_points(self.num_boundary_points)
    }
}

/// Hemodynamics simulation service
///
/// This is the main service that manages simulation lifecycle,
/// handles inference requests, and provides performance metrics.
pub struct HemodynamicsService {
    /// Service configuration
    config: ServiceConfig,
    /// Active simulation handle (if any)
    handle: Arc<RwLock<Option<SimulationHandle>>>,
}

impl HemodynamicsService {
    /// Creates a new hemodynamics service
    #[must_use]
    pub fn new(config: ServiceConfig) -> Self {
        Self {
            config,
            handle: Arc::new(RwLock::new(None)),
        }
    }

    /// Creates a service with default configuration
    #[must_use]
    pub fn with_defaults() -> Self {
        Self::new(ServiceConfig::default())
    }

    /// Initializes a new simulation with the given geometry
    ///
    /// # Arguments
    ///
    /// * `geometry` - Vessel geometry parameters
    ///
    /// # Errors
    ///
    /// Returns error if simulation is already initialized or geometry is invalid
    pub async fn initialize(&self, geometry: VesselGeometry) -> ServerResult<()> {
        let mut handle_guard = self.handle.write().await;

        if handle_guard.is_some() {
            return Err(ServerError::AlreadyInitialized);
        }

        let pinn_config = self.config.to_pinn_config();
        let handle = SimulationHandle::new(geometry, pinn_config).map_err(ServerError::geometry)?;

        *handle_guard = Some(handle);
        Ok(())
    }

    /// Resets the simulation state
    pub async fn reset(&self) -> ServerResult<()> {
        let mut handle_guard = self.handle.write().await;
        *handle_guard = None;
        Ok(())
    }

    /// Checks if simulation is initialized
    pub async fn is_initialized(&self) -> bool {
        self.handle.read().await.is_some()
    }

    /// Queries velocity and pressure fields at given points
    ///
    /// # Arguments
    ///
    /// * `points` - Query points
    /// * `time` - Simulation time
    ///
    /// # Errors
    ///
    /// Returns error if simulation not initialized
    pub async fn query_fields(&self, points: &[Point2D], time: f64) -> ServerResult<FieldResponse> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        let start = Instant::now();

        // Use the model directly for inference (even if not formally "trained")
        let predictions = handle.model().forward_batch(points, time);

        let u: Vec<f64> = predictions.iter().map(|(u, _, _)| *u).collect();
        let v: Vec<f64> = predictions.iter().map(|(_, v, _)| *v).collect();
        let p: Vec<f64> = predictions.iter().map(|(_, _, p)| *p).collect();

        let velocity = VelocityField::new(points.to_vec(), u, v)
            .map_err(|e| ServerError::inference(e.to_string()))?;
        let pressure = PressureField::new(points.to_vec(), p)
            .map_err(|e| ServerError::inference(e.to_string()))?;

        handle.record_inference(start.elapsed());

        Ok(FieldResponse::new(velocity, pressure, None))
    }

    /// Queries fields with wall shear stress
    ///
    /// # Arguments
    ///
    /// * `interior_points` - Interior query points
    /// * `wall_points` - Wall query points for WSS
    /// * `time` - Simulation time
    pub async fn query_fields_with_wss(
        &self,
        interior_points: &[Point2D],
        wall_points: &[Point2D],
        time: f64,
    ) -> ServerResult<FieldResponse> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        let start = Instant::now();

        // Query interior fields
        let interior_preds = handle.model().forward_batch(interior_points, time);
        let u: Vec<f64> = interior_preds.iter().map(|(u, _, _)| *u).collect();
        let v: Vec<f64> = interior_preds.iter().map(|(_, v, _)| *v).collect();
        let p: Vec<f64> = interior_preds.iter().map(|(_, _, p)| *p).collect();

        let velocity = VelocityField::new(interior_points.to_vec(), u, v)
            .map_err(|e| ServerError::inference(e.to_string()))?;
        let pressure = PressureField::new(interior_points.to_vec(), p)
            .map_err(|e| ServerError::inference(e.to_string()))?;

        // Compute WSS at wall points
        let wss_computer = rtx_hemodynamics::wss::WssComputer::blood();
        let wss_values = wss_computer.compute_wss_batch(wall_points, handle.vessel(), |point| {
            let (u, v, _) = handle.model().forward(point.x, point.y, time);
            (u, v)
        });

        let wss = WssField::new(wall_points.to_vec(), wss_values)
            .map_err(|e| ServerError::inference(e.to_string()))?;

        handle.record_inference(start.elapsed());

        Ok(FieldResponse::new(velocity, pressure, Some(wss)))
    }

    /// Evaluates fields on a regular grid for visualization
    pub async fn query_grid(
        &self,
        nx: usize,
        ny: usize,
        time: f64,
    ) -> ServerResult<GridQueryResult> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        let start = Instant::now();

        // Get bounding box
        let (bbox_min, bbox_max) = handle.geometry().bounding_box();
        let length = handle.geometry().length();

        let dx = length / (nx - 1).max(1) as f64;
        let dy = (bbox_max.y - bbox_min.y) / (ny - 1).max(1) as f64;

        let mut points = Vec::with_capacity(nx * ny);
        let mut mask = Vec::with_capacity(nx * ny);

        for j in 0..ny {
            for i in 0..nx {
                let x = i as f64 * dx;
                let y = bbox_min.y + j as f64 * dy;
                let p = Point2D::new(x, y);

                points.push(p);
                mask.push(handle.vessel().is_inside(&p));
            }
        }

        let predictions = handle.model().forward_batch(&points, time);

        let u: Vec<f64> = predictions.iter().map(|(u, _, _)| *u).collect();
        let v: Vec<f64> = predictions.iter().map(|(_, v, _)| *v).collect();
        let p: Vec<f64> = predictions.iter().map(|(_, _, p)| *p).collect();

        handle.record_inference(start.elapsed());

        Ok(GridQueryResult {
            nx,
            ny,
            dx,
            dy,
            x_min: 0.0,
            y_min: bbox_min.y,
            u,
            v,
            p,
            mask,
        })
    }

    /// Modifies the vessel geometry
    pub async fn modify_geometry(&self, modification: GeometryModification) -> ServerResult<()> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        // Get current geometry and apply modification
        let new_geometry = match modification {
            GeometryModification::AddStenosis(params) => {
                handle.geometry().clone().with_stenosis(params)
            }
            GeometryModification::AddAneurysm(params) => {
                handle.geometry().clone().with_aneurysm(params)
            }
            GeometryModification::PlaceStent(_) => {
                // Stent application is more complex - for now just return current
                handle.geometry().clone()
            }
            GeometryModification::ModifyStenosisSeverity(_severity) => {
                // Would need to modify existing stenosis - for now just return current
                handle.geometry().clone()
            }
            GeometryModification::Reset => VesselGeometry::straight(
                handle.geometry().length(),
                handle.geometry().base_radius(),
            )
            .map_err(|e| ServerError::geometry(e.to_string()))?,
        };

        // Recreate handle with new geometry
        let pinn_config = self.config.to_pinn_config();
        let new_handle =
            SimulationHandle::new(new_geometry, pinn_config).map_err(ServerError::geometry)?;

        *handle = new_handle;
        Ok(())
    }

    /// Returns current simulation status
    pub async fn status(&self) -> ServerResult<SimulationStatus> {
        let handle_guard = self.handle.read().await;
        match handle_guard.as_ref() {
            Some(handle) => Ok(handle.status()),
            None => Ok(SimulationStatus::default()),
        }
    }

    /// Returns performance metrics
    pub async fn metrics(&self) -> ServerResult<PerformanceMetrics> {
        let handle_guard = self.handle.read().await;
        match handle_guard.as_ref() {
            Some(handle) => Ok(handle.metrics()),
            None => Ok(PerformanceMetrics::new(0.0, 0.0, 0)),
        }
    }

    /// Returns simulation state
    pub async fn simulation_state(&self) -> ServerResult<SimulationState> {
        let handle_guard = self.handle.read().await;
        match handle_guard.as_ref() {
            Some(handle) => Ok(handle.simulation_state()),
            None => Ok(SimulationState::new()),
        }
    }

    /// Samples interior points from the vessel
    pub async fn sample_interior(&self, n: usize) -> ServerResult<Vec<Point2D>> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        Ok(handle.vessel().sample_interior(n, 42))
    }

    /// Samples boundary points from the vessel
    pub async fn sample_boundary(&self, n: usize) -> ServerResult<Vec<Point2D>> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        Ok(handle.vessel().sample_boundary(n))
    }
}

/// Result of a grid query
#[derive(Debug, Clone)]
pub struct GridQueryResult {
    /// Number of grid points in x
    pub nx: usize,
    /// Number of grid points in y
    pub ny: usize,
    /// Grid spacing in x
    pub dx: f64,
    /// Grid spacing in y
    pub dy: f64,
    /// Minimum x coordinate
    pub x_min: f64,
    /// Minimum y coordinate
    pub y_min: f64,
    /// X-velocity component (row-major)
    pub u: Vec<f64>,
    /// Y-velocity component
    pub v: Vec<f64>,
    /// Pressure
    pub p: Vec<f64>,
    /// Interior mask
    pub mask: Vec<bool>,
}

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

    async fn create_test_service() -> HemodynamicsService {
        let config = ServiceConfig {
            num_layers: 2,
            hidden_dim: 16,
            fourier_dim: 8,
            num_collocation_points: 100,
            num_boundary_points: 50,
            inference_batch_size: 100,
        };
        HemodynamicsService::new(config)
    }

    #[tokio::test]
    async fn test_service_creation() {
        let service = create_test_service().await;
        assert!(!service.is_initialized().await);
    }

    #[tokio::test]
    async fn test_initialize() {
        let service = create_test_service().await;
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();

        service.initialize(geometry).await.unwrap();
        assert!(service.is_initialized().await);
    }

    #[tokio::test]
    async fn test_double_initialize_fails() {
        let service = create_test_service().await;
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();

        service.initialize(geometry.clone()).await.unwrap();
        let result = service.initialize(geometry).await;

        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            ServerError::AlreadyInitialized
        ));
    }

    #[tokio::test]
    async fn test_reset() {
        let service = create_test_service().await;
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();

        service.initialize(geometry.clone()).await.unwrap();
        service.reset().await.unwrap();

        assert!(!service.is_initialized().await);

        // Should be able to initialize again
        service.initialize(geometry).await.unwrap();
        assert!(service.is_initialized().await);
    }

    #[tokio::test]
    async fn test_query_fields() {
        let service = create_test_service().await;
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();
        service.initialize(geometry).await.unwrap();

        let points = vec![Point2D::new(0.05, 0.0), Point2D::new(0.05, 0.002)];

        let response = service.query_fields(&points, 0.0).await.unwrap();

        assert_eq!(response.velocity().len(), 2);
        assert_eq!(response.pressure().len(), 2);
    }

    #[tokio::test]
    async fn test_query_not_initialized() {
        let service = create_test_service().await;
        let points = vec![Point2D::new(0.05, 0.0)];

        let result = service.query_fields(&points, 0.0).await;
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), ServerError::NotInitialized));
    }

    #[tokio::test]
    async fn test_query_grid() {
        let service = create_test_service().await;
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();
        service.initialize(geometry).await.unwrap();

        let result = service.query_grid(10, 5, 0.0).await.unwrap();

        assert_eq!(result.nx, 10);
        assert_eq!(result.ny, 5);
        assert_eq!(result.u.len(), 50);
        assert_eq!(result.mask.len(), 50);
    }

    #[tokio::test]
    async fn test_modify_geometry() {
        let service = create_test_service().await;
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();
        service.initialize(geometry).await.unwrap();

        let stenosis = StenosisParams::new(0.5, 0.02, 0.05).unwrap();
        let modification = GeometryModification::AddStenosis(stenosis);

        service.modify_geometry(modification).await.unwrap();

        // Should still be initialized
        assert!(service.is_initialized().await);
    }

    #[tokio::test]
    async fn test_status() {
        let service = create_test_service().await;

        // Not initialized
        let status = service.status().await.unwrap();
        assert!(!status.initialized);

        // After initialization
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();
        service.initialize(geometry).await.unwrap();

        let status = service.status().await.unwrap();
        assert!(status.initialized);
    }

    #[tokio::test]
    async fn test_metrics() {
        let service = create_test_service().await;
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();
        service.initialize(geometry).await.unwrap();

        // Do some inference
        let points = vec![Point2D::new(0.05, 0.0)];
        service.query_fields(&points, 0.0).await.unwrap();

        let metrics = service.metrics().await.unwrap();
        assert!(metrics.inference_time_ms() >= 0.0);
    }

    #[tokio::test]
    async fn test_sample_interior() {
        let service = create_test_service().await;
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();
        service.initialize(geometry).await.unwrap();

        let points = service.sample_interior(100).await.unwrap();
        assert_eq!(points.len(), 100);
    }

    #[tokio::test]
    async fn test_sample_boundary() {
        let service = create_test_service().await;
        let geometry = VesselGeometry::straight(0.1, 0.005).unwrap();
        service.initialize(geometry).await.unwrap();

        let points = service.sample_boundary(50).await.unwrap();
        assert_eq!(points.len(), 50);
    }
}