rustytorch/crates/specialized/rtx-neuro-pinn/src/solver.rs

//! PINN solver for neural source localization.
//!
//! Implements the training loop and optimization for physics-informed
//! neural networks solving the MEG/EEG inverse problem.

use crate::error::PinnResult;
use crate::head_model::{HeadModel, SensorArray, SourceSpace};
use crate::maxwell::{CurrentDensity, MagneticMaxwell, QuasiStaticMaxwell};
use crate::network::{SourceNetwork, SourceNetworkConfig};
use ndarray::{Array1, Array2};
use serde::{Deserialize, Serialize};

/// Source PINN configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourcePINNConfig {
    /// Network configuration
    pub network: SourceNetworkConfig,
    /// Number of collocation points for physics
    pub n_collocation: usize,
    /// Learning rate
    pub learning_rate: f64,
    /// Weight for data fitting loss
    pub data_weight: f64,
    /// Weight for physics loss
    pub physics_weight: f64,
    /// Weight for sparsity regularization
    pub sparsity_weight: f64,
    /// Weight for smoothness regularization
    pub smoothness_weight: f64,
    /// Whether to learn conductivity jointly
    pub learn_conductivity: bool,
    /// Maximum number of iterations
    pub max_iterations: usize,
    /// Convergence tolerance
    pub tolerance: f64,
    /// Print frequency (0 = no printing)
    pub print_every: usize,
}

impl Default for SourcePINNConfig {
    fn default() -> Self {
        Self {
            network: SourceNetworkConfig::default(),
            n_collocation: 1000,
            learning_rate: 1e-3,
            data_weight: 1.0,
            physics_weight: 0.1,
            sparsity_weight: 0.01,
            smoothness_weight: 0.001,
            learn_conductivity: false,
            max_iterations: 5000,
            tolerance: 1e-6,
            print_every: 100,
        }
    }
}

/// Training result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingResult {
    /// Final total loss
    pub final_loss: f64,
    /// Final data loss
    pub data_loss: f64,
    /// Final physics loss
    pub physics_loss: f64,
    /// Number of iterations
    pub iterations: usize,
    /// Loss history
    pub loss_history: Vec<f64>,
    /// Whether converged
    pub converged: bool,
}

/// Source estimate from PINN
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourceEstimate {
    /// Source positions [n_sources, 3]
    pub positions: Array2<f64>,
    /// Current density components [n_sources, 3]
    pub current_density: Array2<f64>,
    /// Source magnitudes [n_sources]
    pub magnitudes: Array1<f64>,
    /// Estimated conductivity field (if learned)
    pub conductivity: Option<Array1<f64>>,
    /// Goodness of fit (1 - relative residual)
    pub gof: f64,
}

/// PINN-based source localizer
pub struct SourcePINN {
    /// Configuration
    config: SourcePINNConfig,
    /// Neural network
    network: SourceNetwork,
    /// Maxwell physics
    maxwell: QuasiStaticMaxwell,
    /// Magnetic field computation (for MEG)
    magnetic: MagneticMaxwell,
    /// Current training state
    iteration: usize,
}

impl SourcePINN {
    /// Create a new source PINN
    pub fn new(config: SourcePINNConfig) -> PinnResult<Self> {
        // Create network with conductivity output if learning
        let mut network_config = config.network.clone();
        network_config.output_conductivity = config.learn_conductivity;

        let network = SourceNetwork::new(network_config);
        let maxwell = QuasiStaticMaxwell::new(0.33)
            .with_sparsity(config.sparsity_weight)
            .with_smoothness(config.smoothness_weight);
        let magnetic = MagneticMaxwell::new();

        Ok(Self {
            config,
            network,
            maxwell,
            magnetic,
            iteration: 0,
        })
    }

    /// Get configuration
    pub fn config(&self) -> &SourcePINNConfig {
        &self.config
    }

    /// Train the PINN on sensor measurements
    ///
    /// # Arguments
    /// * `sensor_data` - Sensor measurements [n_sensors] or [n_sensors, n_times]
    /// * `head` - Head model with geometry and conductivity
    /// * `sensors` - Sensor array configuration
    ///
    /// # Returns
    /// Training result with loss history
    pub fn train(
        &mut self,
        sensor_data: &Array1<f64>,
        head: &HeadModel,
        sensors: &SensorArray,
    ) -> PinnResult<TrainingResult> {
        let mut loss_history = Vec::with_capacity(self.config.max_iterations);
        let mut prev_loss = f64::MAX;

        // Sample collocation points in brain
        let collocation_points = head.sample_brain_points(self.config.n_collocation);

        for iter in 0..self.config.max_iterations {
            self.iteration = iter;

            // Forward pass: get current density at collocation points
            let (current, learned_sigma) = self.network.forward(&collocation_points);

            // Get conductivity at collocation points
            let conductivity = if self.config.learn_conductivity {
                learned_sigma.unwrap_or_else(|| head.conductivity_at_points(&collocation_points))
            } else {
                head.conductivity_at_points(&collocation_points)
            };

            // Create current density from network output
            let current_density = self.array_to_current_density(&collocation_points, &current);

            // Compute predicted sensor measurements
            let predicted = self.compute_sensor_predictions(&current_density, sensors);

            // Compute losses
            let data_loss = self.compute_data_loss(&predicted, sensor_data);
            let physics_loss =
                self.compute_physics_loss(&collocation_points, &current, &conductivity);
            let sparsity_loss = self.maxwell.sparsity_loss(&current_density);
            let smoothness_loss = self.maxwell.smoothness_loss(&current_density);

            let total_loss = self.config.data_weight * data_loss
                + self.config.physics_weight * physics_loss
                + sparsity_loss
                + smoothness_loss;

            loss_history.push(total_loss);

            // Check convergence
            if (prev_loss - total_loss).abs() < self.config.tolerance {
                return Ok(TrainingResult {
                    final_loss: total_loss,
                    data_loss,
                    physics_loss,
                    iterations: iter + 1,
                    loss_history,
                    converged: true,
                });
            }

            // Update network parameters (simplified gradient descent)
            self.update_parameters(&collocation_points, sensor_data, sensors, &conductivity)?;

            prev_loss = total_loss;
        }

        // Did not converge
        let final_loss = *loss_history.last().unwrap_or(&f64::MAX);
        Ok(TrainingResult {
            final_loss,
            data_loss: 0.0,
            physics_loss: 0.0,
            iterations: self.config.max_iterations,
            loss_history,
            converged: false,
        })
    }

    /// Estimate sources after training
    pub fn estimate_sources(&self, source_space: &SourceSpace) -> PinnResult<SourceEstimate> {
        let n_sources = source_space.n_sources();
        let positions = &source_space.positions;

        // Get current density at source positions
        let (current, sigma) = self.network.forward(positions);

        // Compute magnitudes
        let mut magnitudes = Array1::zeros(n_sources);
        for i in 0..n_sources {
            let jx = current[[i, 0]];
            let jy = current[[i, 1]];
            let jz = current[[i, 2]];
            magnitudes[i] = (jx * jx + jy * jy + jz * jz).sqrt();
        }

        // Compute goodness of fit (simplified)
        let total_power: f64 = magnitudes.iter().map(|&m| m * m).sum();
        let gof = if total_power > 0.0 {
            1.0 - 1.0 / (1.0 + total_power)
        } else {
            0.0
        };

        Ok(SourceEstimate {
            positions: positions.clone(),
            current_density: current,
            magnitudes,
            conductivity: sigma,
            gof,
        })
    }

    /// Convert array output to CurrentDensity
    fn array_to_current_density(
        &self,
        positions: &Array2<f64>,
        current: &Array2<f64>,
    ) -> CurrentDensity {
        let n = positions.nrows();
        let mut magnitudes = Array1::zeros(n);

        for i in 0..n {
            let jx = current[[i, 0]];
            let jy = current[[i, 1]];
            let jz = current[[i, 2]];
            magnitudes[i] = (jx * jx + jy * jy + jz * jz).sqrt();
        }

        CurrentDensity {
            components: current.clone(),
            positions: positions.clone(),
            magnitudes,
        }
    }

    /// Compute predicted sensor measurements using forward model
    fn compute_sensor_predictions(
        &self,
        current_density: &CurrentDensity,
        sensors: &SensorArray,
    ) -> Array1<f64> {
        // Use magnetic field for MEG-like sensors
        let b_field = self
            .magnetic
            .magnetic_field(current_density, sensors.positions());

        // Project onto sensor orientations if available
        if let Some(orientations) = sensors.orientations() {
            let n = sensors.n_sensors();
            let mut measurements = Array1::zeros(n);

            for i in 0..n {
                measurements[i] = b_field[[i, 0]] * orientations[[i, 0]]
                    + b_field[[i, 1]] * orientations[[i, 1]]
                    + b_field[[i, 2]] * orientations[[i, 2]];
            }
            measurements
        } else {
            // For EEG, use magnitude
            let n = sensors.n_sensors();
            let mut measurements = Array1::zeros(n);
            for i in 0..n {
                measurements[i] =
                    (b_field[[i, 0]].powi(2) + b_field[[i, 1]].powi(2) + b_field[[i, 2]].powi(2))
                        .sqrt();
            }
            measurements
        }
    }

    /// Compute data fitting loss
    fn compute_data_loss(&self, predicted: &Array1<f64>, measured: &Array1<f64>) -> f64 {
        if predicted.len() != measured.len() {
            return f64::MAX;
        }

        predicted
            .iter()
            .zip(measured.iter())
            .map(|(&p, &m)| (p - m).powi(2))
            .sum::<f64>()
            / predicted.len() as f64
    }

    /// Compute physics loss (Maxwell residual)
    fn compute_physics_loss(
        &self,
        points: &Array2<f64>,
        current: &Array2<f64>,
        conductivity: &Array1<f64>,
    ) -> f64 {
        let n = points.nrows();
        let mut residual_sum = 0.0;

        // Simplified physics loss: penalize non-physical current patterns
        for i in 0..n {
            let jx = current[[i, 0]];
            let jy = current[[i, 1]];
            let jz = current[[i, 2]];
            let sigma = conductivity[i];

            // Current magnitude should be reasonable
            let j_mag = (jx * jx + jy * jy + jz * jz).sqrt();

            // Penalize unreasonably large currents (> 1 A/m² is very high for brain)
            if j_mag > 1.0 {
                residual_sum += (j_mag - 1.0).powi(2);
            }

            // Penalize negative conductivity (shouldn't happen but safety)
            if sigma < 0.0 {
                residual_sum += sigma.powi(2);
            }
        }

        residual_sum / n as f64
    }

    /// Update network parameters using finite differences
    fn update_parameters(
        &mut self,
        points: &Array2<f64>,
        measured: &Array1<f64>,
        sensors: &SensorArray,
        conductivity: &Array1<f64>,
    ) -> PinnResult<()> {
        let lr = self.config.learning_rate;
        let eps = 1e-5;

        let params = self.network.get_parameters();
        let n_params = params.len();

        // Limit number of parameters to update per iteration for efficiency
        let max_updates = n_params.min(100);
        let step = (n_params / max_updates).max(1);

        let mut new_params = params.clone();

        // Compute current loss
        let (current, _) = self.network.forward(points);
        let current_density = self.array_to_current_density(points, &current);
        let predicted = self.compute_sensor_predictions(&current_density, sensors);
        let base_loss = self.compute_data_loss(&predicted, measured)
            + self.config.physics_weight
                * self.compute_physics_loss(points, &current, conductivity);

        // Numerical gradient for subset of parameters
        for i in (0..n_params).step_by(step) {
            // Perturb parameter
            let mut perturbed = params.clone();
            perturbed[i] += eps;
            self.network.set_parameters(&perturbed)?;

            // Compute perturbed loss
            let (current, _) = self.network.forward(points);
            let current_density = self.array_to_current_density(points, &current);
            let predicted = self.compute_sensor_predictions(&current_density, sensors);
            let perturbed_loss = self.compute_data_loss(&predicted, measured)
                + self.config.physics_weight
                    * self.compute_physics_loss(points, &current, conductivity);

            // Gradient
            let grad = (perturbed_loss - base_loss) / eps;

            // Update
            new_params[i] -= lr * grad;
        }

        self.network.set_parameters(&new_params)?;
        Ok(())
    }

    /// Get network
    pub fn network(&self) -> &SourceNetwork {
        &self.network
    }

    /// Get mutable network
    pub fn network_mut(&mut self) -> &mut SourceNetwork {
        &mut self.network
    }

    /// Current iteration
    pub fn iteration(&self) -> usize {
        self.iteration
    }
}

/// Quick source estimation without full PINN training
pub fn quick_estimate(
    sensor_data: &Array1<f64>,
    head: &HeadModel,
    sensors: &SensorArray,
    n_sources: usize,
) -> PinnResult<SourceEstimate> {
    // Create a simple source space
    let source_space = SourceSpace::grid(head, 0.01);

    // Limit to n_sources
    let n = source_space.n_sources().min(n_sources);
    let positions = source_space
        .positions
        .slice(ndarray::s![..n, ..])
        .to_owned();

    // Simple beamformer-like estimation
    let n_sensors = sensors.n_sensors();
    let magnetic = MagneticMaxwell::new();

    // Compute lead field
    let mut current_density = Array2::zeros((n, 3));
    let mut magnitudes = Array1::zeros(n);

    for i in 0..n {
        // Create test dipole at each location
        let pos = Array2::from_shape_vec(
            (1, 3),
            vec![positions[[i, 0]], positions[[i, 1]], positions[[i, 2]]],
        )
        .unwrap();

        // Test each orientation
        let mut best_corr = 0.0_f64;
        let mut best_orient = [0.0, 0.0, 1.0];

        for orient in &[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]] {
            let test_current = CurrentDensity {
                components: Array2::from_shape_vec((1, 3), vec![orient[0], orient[1], orient[2]])
                    .unwrap(),
                positions: pos.clone(),
                magnitudes: Array1::from_vec(vec![1.0]),
            };

            let forward = magnetic.magnetic_field(&test_current, sensors.positions());

            // Compute correlation with measured data
            let mut pred = Array1::zeros(n_sensors);
            if let Some(orientations) = sensors.orientations() {
                for s in 0..n_sensors {
                    pred[s] = forward[[s, 0]] * orientations[[s, 0]]
                        + forward[[s, 1]] * orientations[[s, 1]]
                        + forward[[s, 2]] * orientations[[s, 2]];
                }
            } else {
                for s in 0..n_sensors {
                    pred[s] = (forward[[s, 0]].powi(2)
                        + forward[[s, 1]].powi(2)
                        + forward[[s, 2]].powi(2))
                    .sqrt();
                }
            }

            // Correlation
            let pred_mean = pred.mean().unwrap_or(0.0);
            let data_mean = sensor_data.mean().unwrap_or(0.0);

            let mut cov = 0.0;
            let mut var_pred = 0.0;
            let mut var_data = 0.0;

            for s in 0..n_sensors {
                let dp = pred[s] - pred_mean;
                let dd = sensor_data[s] - data_mean;
                cov += dp * dd;
                var_pred += dp * dp;
                var_data += dd * dd;
            }

            let corr = if var_pred > 1e-10 && var_data > 1e-10 {
                (cov / (var_pred.sqrt() * var_data.sqrt())).abs()
            } else {
                0.0
            };

            if corr > best_corr {
                best_corr = corr;
                best_orient = *orient;
            }
        }

        current_density[[i, 0]] = best_orient[0] * best_corr;
        current_density[[i, 1]] = best_orient[1] * best_corr;
        current_density[[i, 2]] = best_orient[2] * best_corr;
        magnitudes[i] = best_corr;
    }

    // Compute GOF
    let total_power: f64 = magnitudes.iter().map(|&m| m * m).sum();
    let gof = if total_power > 0.0 {
        magnitudes.iter().copied().fold(0.0f64, f64::max).min(1.0)
    } else {
        0.0
    };

    Ok(SourceEstimate {
        positions,
        current_density,
        magnitudes,
        conductivity: None,
        gof,
    })
}

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

    #[test]
    fn test_pinn_creation() {
        let config = SourcePINNConfig::default();
        let pinn = SourcePINN::new(config).unwrap();
        assert_eq!(pinn.iteration(), 0);
    }

    #[test]
    fn test_source_estimate() {
        let config = SourcePINNConfig {
            max_iterations: 10,
            n_collocation: 50,
            ..Default::default()
        };
        let pinn = SourcePINN::new(config).unwrap();

        let head = HeadModel::spherical(3);
        let source_space = SourceSpace::grid(&head, 0.02);

        let estimate = pinn.estimate_sources(&source_space).unwrap();

        assert!(estimate.positions.nrows() > 0);
        assert_eq!(estimate.current_density.ncols(), 3);
        assert!(estimate.gof >= 0.0 && estimate.gof <= 1.0);
    }

    #[test]
    fn test_quick_estimate() {
        let head = HeadModel::spherical(3);
        let sensors = SensorArray::meg_helmet(50, 0.1, 0.02);
        let sensor_data = Array1::zeros(sensors.n_sensors());

        let estimate = quick_estimate(&sensor_data, &head, &sensors, 100).unwrap();

        assert!(estimate.positions.nrows() > 0);
    }

    #[test]
    fn test_training_small() {
        let config = SourcePINNConfig {
            max_iterations: 5,
            n_collocation: 20,
            network: SourceNetworkConfig {
                hidden_layers: vec![16, 16],
                n_fourier_features: 8,
                ..Default::default()
            },
            ..Default::default()
        };

        let mut pinn = SourcePINN::new(config).unwrap();
        let head = HeadModel::spherical(3);
        let sensors = SensorArray::meg_helmet(20, 0.1, 0.02);
        let sensor_data = Array1::from_vec(vec![1e-12; 20]);

        let result = pinn.train(&sensor_data, &head, &sensors).unwrap();

        assert!(result.iterations > 0);
        assert!(result.loss_history.len() > 0);
    }
}