//! IPC types for MRE demo communication

use crate::fields::{StiffnessField, WaveField};
use serde::{Deserialize, Serialize};

/// Training loss record for a single step
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct LossRecord {
    /// Training step number
    pub step: usize,
    /// Total combined loss
    pub total_loss: f32,
    /// Data fitting loss (wave field match)
    pub data_loss: f32,
    /// Physics loss (Helmholtz residual)
    pub physics_loss: f32,
}

/// Snapshot of solver state for visualization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MreSnapshot {
    /// Current stiffness map
    pub stiffness: StiffnessField,
    /// Predicted wave field (real component)
    pub wave_real: Vec<f32>,
    /// Predicted wave field (imaginary component)
    pub wave_imag: Vec<f32>,
    /// Physics residual magnitude at grid points
    pub residual: Vec<f32>,
    /// Grid resolution for wave/residual (nx, ny)
    pub grid_resolution: (usize, usize),
    /// Current training step
    pub step: usize,
    /// Current loss values
    pub loss: LossRecord,
}

/// Solver status information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MreStatus {
    /// Whether the solver is initialized
    pub initialized: bool,
    /// Whether training is in progress
    pub training: bool,
    /// Current training step
    pub step: usize,
    /// Total steps completed
    pub total_steps: usize,
    /// Steps per second (training throughput)
    pub steps_per_second: f32,
}

impl Default for MreStatus {
    fn default() -> Self {
        Self {
            initialized: false,
            training: false,
            step: 0,
            total_steps: 0,
            steps_per_second: 0.0,
        }
    }
}

/// Request to initialize the MRE solver
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InitializeRequest {
    /// Phantom type to use for synthetic data
    pub phantom_type: PhantomType,
    /// Optional custom wave field data (overrides phantom)
    pub wave_data: Option<WaveField>,
}

/// Types of synthetic phantoms available
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
pub enum PhantomType {
    /// Single centered tumor
    #[default]
    SingleTumor,
    /// Multiple lesions of varying stiffness
    MultipleLesions,
    /// Layered tissue (e.g., liver with capsule)
    Layered,
}

/// Request to run training steps
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainRequest {
    /// Number of steps to run
    pub num_steps: usize,
    /// Whether to return snapshots at intervals
    pub return_snapshots: bool,
    /// Snapshot interval (every N steps)
    pub snapshot_interval: usize,
}

/// Response from training request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainResponse {
    /// Loss records for each step
    pub losses: Vec<LossRecord>,
    /// Final snapshot (if requested)
    pub final_snapshot: Option<MreSnapshot>,
}

/// Request for a visualization snapshot
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotRequest {
    /// Grid resolution for wave field visualization
    pub grid_nx: usize,
    pub grid_ny: usize,
}

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

    #[test]
    fn test_loss_record_default() {
        let loss = LossRecord::default();
        assert_eq!(loss.step, 0);
        assert!((loss.total_loss - 0.0).abs() < 1e-6);
    }

    #[test]
    fn test_serialize_phantom_type() {
        let phantom = PhantomType::SingleTumor;
        let json = serde_json::to_string(&phantom).unwrap();
        assert_eq!(json, "\"SingleTumor\"");
    }
}