rustytorch/demos/server/src/bioheat_service.rs

//! Thermal Ablation Bioheat service implementation
//!
//! This module provides the main service interface for 3D thermal ablation
//! simulation using the Pennes bioheat equation. It handles initialization,
//! training steps, and snapshot generation for visualization.

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

use bioheat_shared::{
    BioheatLossRecord, BioheatSnapshot, BioheatStatus, ProbeGeometry, SimulationParams, SliceAxis,
    SliceData,
};
use rtx_bioheat::config::BioheatConfig;
use rtx_bioheat::solver::BioheatSolver;

use crate::error::{ServerError, ServerResult};

/// Bioheat service configuration
#[derive(Debug, Clone)]
pub struct BioheatServiceConfig {
    /// Number of Fourier features
    pub fourier_features: usize,
    /// Hidden layer sizes
    pub hidden_layers: Vec<usize>,
    /// Number of collocation points for physics
    pub num_collocation: usize,
    /// Snapshot resolution (nx, ny, nz)
    pub snapshot_resolution: (usize, usize, usize),
    /// Slice resolution for 2D visualization
    pub slice_resolution: (usize, usize),
}

impl Default for BioheatServiceConfig {
    fn default() -> Self {
        Self {
            fourier_features: 64,
            hidden_layers: vec![128, 128, 128, 64],
            num_collocation: 4096,
            snapshot_resolution: (32, 32, 32),
            slice_resolution: (64, 64),
        }
    }
}

impl BioheatServiceConfig {
    /// Fast configuration for testing
    #[must_use]
    pub fn fast() -> Self {
        Self {
            fourier_features: 16,
            hidden_layers: vec![32, 32],
            num_collocation: 512,
            snapshot_resolution: (16, 16, 16),
            slice_resolution: (32, 32),
        }
    }

    /// Creates a `BioheatConfig` from service config
    fn to_bioheat_config(&self, params: &SimulationParams) -> BioheatConfig {
        let mut config = BioheatConfig::from_tissue_type(params.tissue_type);
        config.probe = params.probe.clone();
        config.probe_power = params.probe_power;
        config.domain = params.bounds();
        config.time.t_end = params.duration;
        config.network.fourier_features = self.fourier_features;
        config.network.hidden_layers = self.hidden_layers.clone();
        config.training.num_collocation = self.num_collocation;
        config
    }
}

/// Bioheat simulation handle
struct BioheatHandle {
    /// Solver instance
    solver: BioheatSolver,
    /// Whether training is in progress
    training: bool,
    /// Last step time in milliseconds
    last_step_ms: f32,
    /// Total training steps
    total_steps: usize,
    /// Start time
    start_time: Instant,
    /// Simulation parameters
    params: SimulationParams,
}

impl BioheatHandle {
    /// Create a new handle
    fn new(solver: BioheatSolver, params: SimulationParams) -> Self {
        Self {
            solver,
            training: false,
            last_step_ms: 0.0,
            total_steps: 0,
            start_time: Instant::now(),
            params,
        }
    }

    /// Get status struct
    fn status(&self) -> BioheatStatus {
        let elapsed = self.start_time.elapsed().as_secs_f32();
        let steps_per_second = if elapsed > 0.0 {
            self.total_steps as f32 / elapsed
        } else {
            0.0
        };

        let solver_status = self.solver.status();

        BioheatStatus {
            initialized: true,
            training: self.training,
            step: self.solver.current_step(),
            total_steps: self.total_steps,
            simulation_time: self.solver.current_time(),
            steps_per_second,
            max_temperature: solver_status.max_temperature,
            ablation_volume_mm3: solver_status.ablation_volume_mm3,
        }
    }
}

/// Thermal Ablation Bioheat service
///
/// This service manages the bioheat solver lifecycle, handles training steps,
/// time advancement, and provides snapshots for 3D visualization.
pub struct BioheatService {
    /// Service configuration
    config: BioheatServiceConfig,
    /// Active handle (if any)
    handle: Arc<RwLock<Option<BioheatHandle>>>,
}

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

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

    /// Initialize with simulation parameters
    ///
    /// # Errors
    ///
    /// Returns error if already initialized
    pub async fn initialize(&self, params: SimulationParams) -> ServerResult<BioheatStatus> {
        let mut handle_guard = self.handle.write().await;

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

        let bioheat_config = self.config.to_bioheat_config(&params);
        let solver = BioheatSolver::new(bioheat_config);

        let handle = BioheatHandle::new(solver, params);
        let status = handle.status();
        *handle_guard = Some(handle);

        Ok(status)
    }

    /// Initialize with default liver ablation parameters
    pub async fn initialize_default(&self) -> ServerResult<BioheatStatus> {
        self.initialize(SimulationParams::liver_default()).await
    }

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

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

    /// Run a single training step
    ///
    /// # Returns
    ///
    /// Loss record from the step
    pub async fn step(&self) -> ServerResult<BioheatLossRecord> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        let start = Instant::now();
        handle.training = true;

        let loss = handle.solver.step();

        handle.last_step_ms = start.elapsed().as_secs_f32() * 1000.0;
        handle.total_steps += 1;

        Ok(loss)
    }

    /// Run multiple training steps
    ///
    /// # Arguments
    ///
    /// * `num_steps` - Number of steps to run
    ///
    /// # Returns
    ///
    /// Vector of loss records
    pub async fn train(&self, num_steps: usize) -> ServerResult<Vec<BioheatLossRecord>> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        let start = Instant::now();
        handle.training = true;

        let losses = handle.solver.train(num_steps);

        handle.last_step_ms = start.elapsed().as_secs_f32() * 1000.0;
        handle.total_steps += num_steps;

        Ok(losses)
    }

    /// Get a snapshot of current state
    ///
    /// Returns temperature field, ablation zone, and probe for visualization
    pub async fn snapshot(&self) -> ServerResult<BioheatSnapshot> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        let res = self.config.snapshot_resolution;
        Ok(handle.solver.snapshot(res))
    }

    /// Get a 2D slice at the specified position
    ///
    /// # Arguments
    ///
    /// * `axis` - Axis perpendicular to the slice (X, Y, or Z)
    /// * `position` - Position along the axis (normalized 0.0 to 1.0)
    pub async fn get_slice(&self, axis: SliceAxis, position: f32) -> ServerResult<SliceData> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        // Get full snapshot and extract slice
        let res = self.config.snapshot_resolution;
        let snapshot = handle.solver.snapshot(res);

        // Convert normalized position to grid index
        let index = match axis {
            SliceAxis::X => (position * (res.0 - 1) as f32) as usize,
            SliceAxis::Y => (position * (res.1 - 1) as f32) as usize,
            SliceAxis::Z => (position * (res.2 - 1) as f32) as usize,
        };

        Ok(SliceData::from_field(&snapshot.temperature, axis, index))
    }

    /// Update probe position
    pub async fn update_probe(&self, probe: ProbeGeometry) -> ServerResult<()> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        handle.solver.update_probe(probe.clone());
        handle.params.probe = probe;
        Ok(())
    }

    /// Update probe power
    pub async fn update_power(&self, power: f32) -> ServerResult<()> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        handle.solver.update_probe_power(power);
        handle.params.probe_power = power;
        Ok(())
    }

    /// Advance simulation time
    pub async fn advance_time(&self, dt: f32) -> ServerResult<f32> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        handle.solver.advance_time(dt);
        Ok(handle.solver.current_time())
    }

    /// Set simulation time directly
    pub async fn set_time(&self, t: f32) -> ServerResult<f32> {
        let mut handle_guard = self.handle.write().await;
        let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;

        handle.solver.set_time(t);
        Ok(handle.solver.current_time())
    }

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

    /// Get loss history
    pub async fn loss_history(&self) -> ServerResult<Vec<BioheatLossRecord>> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        Ok(handle.solver.loss_history().to_vec())
    }

    /// Get current step count
    pub async fn current_step(&self) -> ServerResult<usize> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        Ok(handle.solver.current_step())
    }

    /// Get current simulation time
    pub async fn current_time(&self) -> ServerResult<f32> {
        let handle_guard = self.handle.read().await;
        let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;

        Ok(handle.solver.current_time())
    }

    /// Get performance metrics
    pub async fn metrics(&self) -> ServerResult<BioheatMetrics> {
        let handle_guard = self.handle.read().await;
        match handle_guard.as_ref() {
            Some(handle) => Ok(BioheatMetrics {
                step_time_ms: handle.last_step_ms,
                total_steps: handle.total_steps,
                elapsed_secs: handle.start_time.elapsed().as_secs_f32(),
            }),
            None => Ok(BioheatMetrics::default()),
        }
    }
}

/// Performance metrics for bioheat service
#[derive(Debug, Clone, Default)]
pub struct BioheatMetrics {
    /// Last step time in milliseconds
    pub step_time_ms: f32,
    /// Total steps executed
    pub total_steps: usize,
    /// Total elapsed time in seconds
    pub elapsed_secs: f32,
}

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

    async fn create_test_service() -> BioheatService {
        BioheatService::new(BioheatServiceConfig::fast())
    }

    #[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 status = service.initialize_default().await.unwrap();
        assert!(status.initialized);
        assert!(service.is_initialized().await);
    }

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

        let params = SimulationParams {
            tissue_type: TissueType::Tumor,
            probe_power: 25.0,
            ..SimulationParams::default()
        };

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

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

        service.initialize_default().await.unwrap();
        let result = service.initialize_default().await;

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

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

        let loss = service.step().await.unwrap();
        assert!(loss.total_loss.is_finite());
        assert_eq!(loss.step, 1);
    }

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

        let losses = service.train(5).await.unwrap();
        assert_eq!(losses.len(), 5);
    }

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

        let snapshot = service.snapshot().await.unwrap();
        assert!(!snapshot.temperature.values.is_empty());
    }

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

        let slice = service.get_slice(SliceAxis::Z, 0.5).await.unwrap();
        assert!(!slice.values.is_empty());
        assert_eq!(slice.axis, SliceAxis::Z);
    }

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

        let new_probe = ProbeGeometry::rf_needle(Point3D::new(0.01, 0.02, 0.0));
        service.update_probe(new_probe).await.unwrap();

        // Should not error
        let _ = service.snapshot().await.unwrap();
    }

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

        service.update_power(30.0).await.unwrap();
        // Should not error
        let _ = service.snapshot().await.unwrap();
    }

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

        let t1 = service.advance_time(10.0).await.unwrap();
        assert!((t1 - 10.0).abs() < 1e-6);

        let t2 = service.set_time(50.0).await.unwrap();
        assert!((t2 - 50.0).abs() < 1e-6);
    }

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

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

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

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

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

        service.initialize_default().await.unwrap();

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

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

        service.train(3).await.unwrap();

        let history = service.loss_history().await.unwrap();
        assert_eq!(history.len(), 3);
    }

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

        service.train(5).await.unwrap();

        let metrics = service.metrics().await.unwrap();
        assert_eq!(metrics.total_steps, 5);
        assert!(metrics.elapsed_secs > 0.0);
    }
}