//! Main bioheat solver combining all components

use crate::{
    config::BioheatConfig,
    network::{Coordinate4D, ThermalPinn},
    pennes::{PennesResidual, TemperatureDerivatives},
    probe::ProbeHeatSource,
    tissue::TissueDomain,
};
use bioheat_shared::{
    AblationZone, BioheatLossRecord, BioheatSnapshot, BioheatStatus, Point3D, ProbeGeometry,
    TemperatureField,
};
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;
use std::time::Instant;

/// AdamW optimizer for neural network training
#[derive(Debug, Clone)]
struct AdamW {
    /// First moment estimates (momentum) for weights
    m_weights: Vec<Vec<Vec<f64>>>,
    /// Second moment estimates (variance) for weights
    v_weights: Vec<Vec<Vec<f64>>>,
    /// First moment estimates for biases
    m_biases: Vec<Vec<f64>>,
    /// Second moment estimates for biases
    v_biases: Vec<Vec<f64>>,
    /// First moment estimates for Fourier features
    m_fourier: Vec<Vec<f64>>,
    /// Second moment estimates for Fourier features
    v_fourier: Vec<Vec<f64>>,
    /// Timestep counter for bias correction
    t: u64,
    /// Learning rate
    lr: f64,
    /// First moment decay rate (β₁)
    beta1: f64,
    /// Second moment decay rate (β₂)
    beta2: f64,
    /// Numerical stability constant (ε)
    eps: f64,
    /// Weight decay coefficient (λ)
    weight_decay: f64,
}

impl AdamW {
    /// Create a new AdamW optimizer matching the network architecture
    fn new(
        weights_shape: &[Vec<Vec<f64>>],
        biases_shape: &[Vec<f64>],
        fourier_shape: &[Vec<f64>],
        lr: f64,
    ) -> Self {
        // Initialize momentum and variance to zero
        let m_weights = weights_shape
            .iter()
            .map(|layer| layer.iter().map(|row| vec![0.0; row.len()]).collect())
            .collect();
        let v_weights = weights_shape
            .iter()
            .map(|layer| layer.iter().map(|row| vec![0.0; row.len()]).collect())
            .collect();

        let m_biases = biases_shape.iter().map(|b| vec![0.0; b.len()]).collect();
        let v_biases = biases_shape.iter().map(|b| vec![0.0; b.len()]).collect();

        let m_fourier = fourier_shape
            .iter()
            .map(|row| vec![0.0; row.len()])
            .collect();
        let v_fourier = fourier_shape
            .iter()
            .map(|row| vec![0.0; row.len()])
            .collect();

        Self {
            m_weights,
            v_weights,
            m_biases,
            v_biases,
            m_fourier,
            v_fourier,
            t: 0,
            lr,
            beta1: 0.9,
            beta2: 0.999,
            eps: 1e-8,
            weight_decay: 0.01,
        }
    }

    /// Perform one optimization step
    fn step(&mut self, grads: &NetworkGradients, pinn: &mut ThermalPinn) {
        self.t += 1;

        // Bias correction factors
        let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
        let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);

        // Update weights
        let weights = pinn.weights_mut();
        for (layer_idx, (layer_grads, layer_weights)) in
            grads.weights.iter().zip(weights.iter_mut()).enumerate()
        {
            for (row_idx, (grad_row, weight_row)) in
                layer_grads.iter().zip(layer_weights.iter_mut()).enumerate()
            {
                for (col_idx, (grad, weight)) in
                    grad_row.iter().zip(weight_row.iter_mut()).enumerate()
                {
                    // Update first moment: m = β₁ * m + (1 - β₁) * g
                    self.m_weights[layer_idx][row_idx][col_idx] = self.beta1
                        * self.m_weights[layer_idx][row_idx][col_idx]
                        + (1.0 - self.beta1) * grad;

                    // Update second moment: v = β₂ * v + (1 - β₂) * g²
                    self.v_weights[layer_idx][row_idx][col_idx] = self.beta2
                        * self.v_weights[layer_idx][row_idx][col_idx]
                        + (1.0 - self.beta2) * grad * grad;

                    // Bias-corrected estimates
                    let m_hat = self.m_weights[layer_idx][row_idx][col_idx] / bias_correction1;
                    let v_hat = self.v_weights[layer_idx][row_idx][col_idx] / bias_correction2;

                    // AdamW update with decoupled weight decay
                    let adam_update = m_hat / (v_hat.sqrt() + self.eps);
                    *weight -= self.lr * (adam_update + self.weight_decay * *weight);
                }
            }
        }

        // Update biases
        let biases = pinn.biases_mut();
        for (layer_idx, (layer_grads, layer_biases)) in
            grads.biases.iter().zip(biases.iter_mut()).enumerate()
        {
            for (idx, (grad, bias)) in layer_grads.iter().zip(layer_biases.iter_mut()).enumerate() {
                self.m_biases[layer_idx][idx] =
                    self.beta1 * self.m_biases[layer_idx][idx] + (1.0 - self.beta1) * grad;

                self.v_biases[layer_idx][idx] =
                    self.beta2 * self.v_biases[layer_idx][idx] + (1.0 - self.beta2) * grad * grad;

                let m_hat = self.m_biases[layer_idx][idx] / bias_correction1;
                let v_hat = self.v_biases[layer_idx][idx] / bias_correction2;

                let adam_update = m_hat / (v_hat.sqrt() + self.eps);
                *bias -= self.lr * adam_update; // No weight decay on biases
            }
        }

        // Update Fourier features
        let fourier_b = pinn.fourier_b_mut();
        for (row_idx, (grad_row, fourier_row)) in
            grads.fourier_b.iter().zip(fourier_b.iter_mut()).enumerate()
        {
            for (col_idx, (grad, fourier_val)) in
                grad_row.iter().zip(fourier_row.iter_mut()).enumerate()
            {
                self.m_fourier[row_idx][col_idx] =
                    self.beta1 * self.m_fourier[row_idx][col_idx] + (1.0 - self.beta1) * grad;

                self.v_fourier[row_idx][col_idx] = self.beta2 * self.v_fourier[row_idx][col_idx]
                    + (1.0 - self.beta2) * grad * grad;

                let m_hat = self.m_fourier[row_idx][col_idx] / bias_correction1;
                let v_hat = self.v_fourier[row_idx][col_idx] / bias_correction2;

                let adam_update = m_hat / (v_hat.sqrt() + self.eps);
                *fourier_val -= self.lr * (adam_update + self.weight_decay * *fourier_val);
            }
        }
    }
}

/// Gradients for all network parameters
#[derive(Debug, Clone)]
struct NetworkGradients {
    /// Gradients for weight matrices
    weights: Vec<Vec<Vec<f64>>>,
    /// Gradients for bias vectors
    biases: Vec<Vec<f64>>,
    /// Gradients for Fourier feature matrix
    fourier_b: Vec<Vec<f64>>,
}

impl NetworkGradients {
    /// Create zero gradients matching network architecture
    fn zeros_like(pinn: &ThermalPinn) -> Self {
        // Clone the network to access shapes
        let pinn_clone = pinn.clone();
        let mut temp_pinn = pinn_clone;

        // Get shapes by accessing each mutable reference separately
        let grad_weights = {
            let weights = temp_pinn.weights_mut();
            weights
                .iter()
                .map(|layer| layer.iter().map(|row| vec![0.0; row.len()]).collect())
                .collect()
        };

        let grad_biases = {
            let biases = temp_pinn.biases_mut();
            biases.iter().map(|b| vec![0.0; b.len()]).collect()
        };

        let grad_fourier = {
            let fourier_b = temp_pinn.fourier_b_mut();
            fourier_b.iter().map(|row| vec![0.0; row.len()]).collect()
        };

        Self {
            weights: grad_weights,
            biases: grad_biases,
            fourier_b: grad_fourier,
        }
    }
}

/// Main bioheat PINN solver
pub struct BioheatSolver {
    /// Configuration
    config: BioheatConfig,
    /// Neural network for temperature prediction
    pinn: ThermalPinn,
    /// Pennes equation residual computer
    pennes: PennesResidual,
    /// Heat source from ablation probe
    probe_source: ProbeHeatSource,
    /// Tissue domain
    domain: TissueDomain,
    /// Current training step
    step: usize,
    /// Current simulation time
    current_time: f32,
    /// Training loss history
    loss_history: Vec<BioheatLossRecord>,
    /// Random number generator (seeded for reproducibility)
    rng: ChaCha8Rng,
    /// Training start time (for throughput calculation)
    train_start: Option<Instant>,
    /// AdamW optimizer
    optimizer: Option<AdamW>,
}

impl BioheatSolver {
    /// Create a new bioheat solver with given configuration
    #[must_use]
    pub fn new(config: BioheatConfig) -> Self {
        let pinn = ThermalPinn::new(config.network.clone());
        let pennes = PennesResidual::new(config.physics.clone());
        let probe_source = ProbeHeatSource::new(config.probe.clone(), config.probe_power);
        let domain = TissueDomain::new(config.domain, (config.time.t_start, config.time.t_end));

        Self {
            config,
            pinn,
            pennes,
            probe_source,
            domain,
            step: 0,
            current_time: 0.0,
            loss_history: Vec::new(),
            rng: ChaCha8Rng::seed_from_u64(42),
            train_start: None,
            optimizer: None,
        }
    }

    /// Create solver with default liver ablation configuration
    #[must_use]
    pub fn liver_default() -> Self {
        Self::new(BioheatConfig::liver_default())
    }

    /// Perform a single training step
    pub fn step(&mut self) -> BioheatLossRecord {
        if self.train_start.is_none() {
            self.train_start = Some(Instant::now());
        }

        // Initialize optimizer on first step
        if self.optimizer.is_none() {
            // Create optimizer by cloning shapes
            let pinn_clone = self.pinn.clone();
            let mut temp_pinn = pinn_clone;

            let weights_shape = {
                let w = temp_pinn.weights_mut();
                w.clone()
            };
            let biases_shape = {
                let b = temp_pinn.biases_mut();
                b.clone()
            };
            let fourier_shape = {
                let f = temp_pinn.fourier_b_mut();
                f.clone()
            };

            self.optimizer = Some(AdamW::new(
                &weights_shape,
                &biases_shape,
                &fourier_shape,
                self.config.training.learning_rate as f64,
            ));
        }

        // Sample training points
        let collocation_points = self
            .domain
            .sample_interior_batch(&mut self.rng, self.config.training.num_collocation);
        let boundary_points = self
            .domain
            .sample_boundary_batch(&mut self.rng, self.config.training.num_boundary);
        let initial_points = self
            .domain
            .sample_initial_batch(&mut self.rng, self.config.training.num_initial);

        // Compute losses
        let physics_loss = self.compute_physics_loss(&collocation_points);
        let boundary_loss = self.compute_boundary_loss(&boundary_points);
        let initial_loss = self.compute_initial_loss(&initial_points);
        let probe_loss = self.compute_probe_loss(&collocation_points);

        // Weight and combine losses
        let weights = &self.config.training.weights;
        let _total_loss = weights.physics * physics_loss
            + weights.boundary * boundary_loss
            + weights.initial * initial_loss
            + weights.probe * probe_loss;

        // Compute gradients using finite differences
        let grads = self.compute_network_gradients();

        // Update weights using AdamW
        if let Some(optimizer) = &mut self.optimizer {
            optimizer.step(&grads, &mut self.pinn);
        }

        self.step += 1;

        let record = BioheatLossRecord::new(
            self.step,
            physics_loss,
            boundary_loss,
            initial_loss,
            probe_loss,
        );
        self.loss_history.push(record.clone());

        record
    }

    /// Train for multiple steps
    pub fn train(&mut self, num_steps: usize) -> Vec<BioheatLossRecord> {
        let mut losses = Vec::with_capacity(num_steps);
        for _ in 0..num_steps {
            losses.push(self.step());
        }
        losses
    }

    /// Compute physics (PDE residual) loss
    fn compute_physics_loss(&self, points: &[(Point3D, f32)]) -> f32 {
        let eps = 1e-4; // For finite difference derivatives
        let mut total_residual_sq = 0.0;

        for (point, t) in points {
            let coord = Coordinate4D::from_point_and_time(*point, *t);

            // Get temperature and derivatives
            let temp = self.pinn.forward(&coord);
            let (dt_dx, dt_dy, dt_dz, dt_dt) = self.pinn.gradients(&coord, eps);
            let (d2t_dx2, d2t_dy2, d2t_dz2) = self.pinn.second_derivatives(&coord, eps);

            let derivs = TemperatureDerivatives {
                t: temp as f32,
                dt_dt: dt_dt as f32,
                dt_dx: dt_dx as f32,
                dt_dy: dt_dy as f32,
                dt_dz: dt_dz as f32,
                d2t_dx2: d2t_dx2 as f32,
                d2t_dy2: d2t_dy2 as f32,
                d2t_dz2: d2t_dz2 as f32,
            };

            // Get heat source at this point
            let heat_source = self.probe_source.heat_source_at(point);

            // Compute residual
            let residual = self.pennes.compute(&derivs, heat_source);
            total_residual_sq += residual * residual;
        }

        (total_residual_sq / points.len() as f32).sqrt()
    }

    /// Compute boundary condition loss (Dirichlet: T = T_body at boundary)
    fn compute_boundary_loss(&self, points: &[(Point3D, f32)]) -> f32 {
        let target_temp = self.config.physics.body_temperature;
        let mut total_error_sq = 0.0;

        for (point, t) in points {
            let coord = Coordinate4D::from_point_and_time(*point, *t);
            let predicted_temp = self.pinn.forward(&coord) as f32;
            let error = predicted_temp - target_temp;
            total_error_sq += error * error;
        }

        (total_error_sq / points.len() as f32).sqrt()
    }

    /// Compute initial condition loss (T = T_body at t=0)
    fn compute_initial_loss(&self, points: &[(Point3D, f32)]) -> f32 {
        let target_temp = self.config.physics.body_temperature;
        let mut total_error_sq = 0.0;

        for (point, t) in points {
            let coord = Coordinate4D::from_point_and_time(*point, *t);
            let predicted_temp = self.pinn.forward(&coord) as f32;
            let error = predicted_temp - target_temp;
            total_error_sq += error * error;
        }

        (total_error_sq / points.len() as f32).sqrt()
    }

    /// Compute probe heat source condition loss
    fn compute_probe_loss(&self, points: &[(Point3D, f32)]) -> f32 {
        // For points near the probe, temperature should be elevated
        // This is a soft constraint to help convergence
        let mut total_loss = 0.0;
        let mut count = 0;

        for (point, t) in points {
            let heat_source = self.probe_source.heat_source_at(point);
            if heat_source > 1e3 {
                // Near probe (significant heat)
                let coord = Coordinate4D::from_point_and_time(*point, *t);
                let predicted_temp = self.pinn.forward(&coord) as f32;

                // Temperature should be above body temp when probe is active
                if predicted_temp < self.config.physics.body_temperature {
                    let error = self.config.physics.body_temperature - predicted_temp;
                    total_loss += error * error;
                }
                count += 1;
            }
        }

        if count > 0 {
            (total_loss / count as f32).sqrt()
        } else {
            0.0
        }
    }

    /// Generate a snapshot of the current state for visualization
    #[must_use]
    pub fn snapshot(&self, resolution: (usize, usize, usize)) -> BioheatSnapshot {
        let grid_points = self.domain.regular_grid(resolution);
        let mut values = Vec::with_capacity(grid_points.len());

        // Evaluate temperature at each grid point for current time
        for point in &grid_points {
            let coord = Coordinate4D::from_point_and_time(*point, self.current_time);
            values.push(self.pinn.forward(&coord) as f32);
        }

        let temperature = TemperatureField {
            resolution,
            values,
            bounds: self.domain.bounds,
        };

        let max_temperature = temperature.max_temperature();
        let ablation_volume = temperature.ablated_volume();

        // Create ablation zone (simplified - no mesh generation)
        let ablation_zone = AblationZone::from_volume(ablation_volume, 60.0);

        let loss = self.loss_history.last().cloned().unwrap_or_default();

        BioheatSnapshot {
            temperature,
            ablation_zone,
            probe: self.probe_source.geometry.clone(),
            step: self.step,
            time: self.current_time,
            loss,
            max_temperature,
        }
    }

    /// Get current solver status
    #[must_use]
    pub fn status(&self) -> BioheatStatus {
        let steps_per_second = if let Some(start) = self.train_start {
            let elapsed = start.elapsed().as_secs_f32();
            if elapsed > 0.0 {
                self.step as f32 / elapsed
            } else {
                0.0
            }
        } else {
            0.0
        };

        // Quick estimate of max temp and ablation volume
        let center_coord = Coordinate4D::from_point_and_time(Point3D::origin(), self.current_time);
        let max_temperature = self.pinn.forward(&center_coord) as f32;

        BioheatStatus {
            initialized: true,
            training: self.train_start.is_some(),
            step: self.step,
            total_steps: self.step,
            simulation_time: self.current_time,
            steps_per_second,
            max_temperature,
            ablation_volume_mm3: 0.0, // Would need full grid evaluation
        }
    }

    /// Advance simulation time
    pub fn advance_time(&mut self, dt: f32) {
        self.current_time = (self.current_time + dt).min(self.config.time.t_end);
    }

    /// Set simulation time directly
    pub fn set_time(&mut self, t: f32) {
        self.current_time = t.clamp(self.config.time.t_start, self.config.time.t_end);
    }

    /// Update probe position
    pub fn update_probe_position(&mut self, position: Point3D) {
        self.probe_source.geometry.position = position;
        self.config.probe.position = position;
    }

    /// Update probe power
    pub fn update_probe_power(&mut self, power: f32) {
        self.probe_source.set_power(power);
        self.config.probe_power = power;
    }

    /// Update probe geometry
    pub fn update_probe(&mut self, probe: ProbeGeometry) {
        self.probe_source.geometry = probe.clone();
        self.config.probe = probe;
    }

    /// Compute gradients of total loss with respect to network parameters using finite differences
    fn compute_network_gradients(&self) -> NetworkGradients {
        // Use finite difference method to compute gradients
        const EPS: f64 = 1e-4;

        // Initialize gradients
        let mut grads = NetworkGradients::zeros_like(&self.pinn);

        // Sample a batch of points for gradient computation
        let mut rng = ChaCha8Rng::seed_from_u64(self.step as u64 + 42);
        let collocation_points = self
            .domain
            .sample_interior_batch(&mut rng, self.config.training.num_collocation);
        let boundary_points = self
            .domain
            .sample_boundary_batch(&mut rng, self.config.training.num_boundary);
        let initial_points = self
            .domain
            .sample_initial_batch(&mut rng, self.config.training.num_initial);

        // Compute base loss (for reference, not used in gradient computation)
        let _base_loss =
            self.compute_total_loss(&collocation_points, &boundary_points, &initial_points);

        // Clone network for perturbation
        let mut pinn_perturbed = self.pinn.clone();

        // Compute gradients for weights using stochastic parameter sampling
        // For efficiency, we sample a subset of parameters per batch
        let total_weight_params: usize = {
            let weights = pinn_perturbed.weights_mut();
            weights
                .iter()
                .map(|layer| layer.iter().map(std::vec::Vec::len).sum::<usize>())
                .sum()
        };

        // Sample up to 256 parameters per step for gradient computation
        let max_samples = 256.min(total_weight_params);
        let sample_scale = total_weight_params as f64 / max_samples as f64;

        for sample_idx in 0..max_samples {
            // Deterministic sampling using step counter
            let flat_idx = (self.step * 31 + sample_idx * 17) % total_weight_params;

            // Convert flat index to (layer, row, col)
            let (layer_idx, row_idx, col_idx) = {
                let weights = pinn_perturbed.weights_mut();
                let mut remaining = flat_idx;
                let mut layer_idx = 0;
                while layer_idx < weights.len() {
                    let layer_size: usize = weights[layer_idx].iter().map(std::vec::Vec::len).sum();
                    if remaining < layer_size {
                        break;
                    }
                    remaining -= layer_size;
                    layer_idx += 1;
                }

                if layer_idx >= weights.len() {
                    continue;
                }

                // Find row and column within layer
                let mut row_idx = 0;
                while row_idx < weights[layer_idx].len()
                    && remaining >= weights[layer_idx][row_idx].len()
                {
                    remaining -= weights[layer_idx][row_idx].len();
                    row_idx += 1;
                }

                if row_idx >= weights[layer_idx].len() {
                    continue;
                }

                let col_idx = remaining;
                (layer_idx, row_idx, col_idx)
            };

            // Compute gradient using central differences
            let original_val = {
                let weights = pinn_perturbed.weights_mut();
                weights[layer_idx][row_idx][col_idx]
            };

            // Perturb +
            {
                let weights = pinn_perturbed.weights_mut();
                weights[layer_idx][row_idx][col_idx] = original_val + EPS;
            }
            let loss_plus = self.compute_total_loss_with_pinn(
                &pinn_perturbed,
                &collocation_points,
                &boundary_points,
                &initial_points,
            );

            // Perturb -
            {
                let weights = pinn_perturbed.weights_mut();
                weights[layer_idx][row_idx][col_idx] = original_val - EPS;
            }
            let loss_minus = self.compute_total_loss_with_pinn(
                &pinn_perturbed,
                &collocation_points,
                &boundary_points,
                &initial_points,
            );

            // Restore original value
            {
                let weights = pinn_perturbed.weights_mut();
                weights[layer_idx][row_idx][col_idx] = original_val;
            }

            // Compute gradient (scaled to account for sampling)
            let grad = ((loss_plus - loss_minus) / (2.0 * EPS)) * sample_scale;
            grads.weights[layer_idx][row_idx][col_idx] = grad;
        }

        // Compute gradients for biases (full gradient, they're small)
        let num_bias_layers = {
            let biases = pinn_perturbed.biases_mut();
            biases.len()
        };

        for layer_idx in 0..num_bias_layers {
            let num_biases = {
                let biases = pinn_perturbed.biases_mut();
                biases[layer_idx].len()
            };

            for bias_idx in 0..num_biases {
                let original_val = {
                    let biases = pinn_perturbed.biases_mut();
                    biases[layer_idx][bias_idx]
                };

                {
                    let biases = pinn_perturbed.biases_mut();
                    biases[layer_idx][bias_idx] = original_val + EPS;
                }
                let loss_plus = self.compute_total_loss_with_pinn(
                    &pinn_perturbed,
                    &collocation_points,
                    &boundary_points,
                    &initial_points,
                );

                {
                    let biases = pinn_perturbed.biases_mut();
                    biases[layer_idx][bias_idx] = original_val - EPS;
                }
                let loss_minus = self.compute_total_loss_with_pinn(
                    &pinn_perturbed,
                    &collocation_points,
                    &boundary_points,
                    &initial_points,
                );

                {
                    let biases = pinn_perturbed.biases_mut();
                    biases[layer_idx][bias_idx] = original_val;
                }

                let grad = (loss_plus - loss_minus) / (2.0 * EPS);
                grads.biases[layer_idx][bias_idx] = grad;
            }
        }

        // Compute gradients for Fourier features (sample-based)
        let total_fourier_params: usize = {
            let fourier_b = pinn_perturbed.fourier_b_mut();
            fourier_b.iter().map(std::vec::Vec::len).sum()
        };
        let fourier_samples = 64.min(total_fourier_params);
        let fourier_scale = total_fourier_params as f64 / fourier_samples as f64;

        let fourier_cols = {
            let fourier_b = pinn_perturbed.fourier_b_mut();
            if fourier_b.is_empty() {
                0
            } else {
                fourier_b[0].len()
            }
        };

        for sample_idx in 0..fourier_samples {
            let flat_idx = (self.step * 37 + sample_idx * 23) % total_fourier_params;

            let row_idx = flat_idx / fourier_cols;
            let col_idx = flat_idx % fourier_cols;

            let fourier_rows = {
                let fourier_b = pinn_perturbed.fourier_b_mut();
                fourier_b.len()
            };

            if row_idx >= fourier_rows {
                continue;
            }

            let original_val = {
                let fourier_b = pinn_perturbed.fourier_b_mut();
                fourier_b[row_idx][col_idx]
            };

            {
                let fourier_b = pinn_perturbed.fourier_b_mut();
                fourier_b[row_idx][col_idx] = original_val + EPS;
            }
            let loss_plus = self.compute_total_loss_with_pinn(
                &pinn_perturbed,
                &collocation_points,
                &boundary_points,
                &initial_points,
            );

            {
                let fourier_b = pinn_perturbed.fourier_b_mut();
                fourier_b[row_idx][col_idx] = original_val - EPS;
            }
            let loss_minus = self.compute_total_loss_with_pinn(
                &pinn_perturbed,
                &collocation_points,
                &boundary_points,
                &initial_points,
            );

            {
                let fourier_b = pinn_perturbed.fourier_b_mut();
                fourier_b[row_idx][col_idx] = original_val;
            }

            let grad = ((loss_plus - loss_minus) / (2.0 * EPS)) * fourier_scale;
            grads.fourier_b[row_idx][col_idx] = grad;
        }

        grads
    }

    /// Compute total weighted loss
    fn compute_total_loss(
        &self,
        collocation_points: &[(Point3D, f32)],
        boundary_points: &[(Point3D, f32)],
        initial_points: &[(Point3D, f32)],
    ) -> f64 {
        let physics_loss = self.compute_physics_loss(collocation_points);
        let boundary_loss = self.compute_boundary_loss(boundary_points);
        let initial_loss = self.compute_initial_loss(initial_points);
        let probe_loss = self.compute_probe_loss(collocation_points);

        let weights = &self.config.training.weights;
        (weights.physics * physics_loss
            + weights.boundary * boundary_loss
            + weights.initial * initial_loss
            + weights.probe * probe_loss) as f64
    }

    /// Compute total weighted loss with a specific PINN instance
    fn compute_total_loss_with_pinn(
        &self,
        pinn: &ThermalPinn,
        collocation_points: &[(Point3D, f32)],
        boundary_points: &[(Point3D, f32)],
        initial_points: &[(Point3D, f32)],
    ) -> f64 {
        let physics_loss = self.compute_physics_loss_with_pinn(pinn, collocation_points);
        let boundary_loss = self.compute_boundary_loss_with_pinn(pinn, boundary_points);
        let initial_loss = self.compute_initial_loss_with_pinn(pinn, initial_points);
        let probe_loss = self.compute_probe_loss_with_pinn(pinn, collocation_points);

        let weights = &self.config.training.weights;
        (weights.physics * physics_loss
            + weights.boundary * boundary_loss
            + weights.initial * initial_loss
            + weights.probe * probe_loss) as f64
    }

    /// Compute physics loss with a specific PINN instance
    fn compute_physics_loss_with_pinn(&self, pinn: &ThermalPinn, points: &[(Point3D, f32)]) -> f32 {
        let eps = 1e-4;
        let mut total_residual_sq = 0.0;

        for (point, t) in points {
            let coord = Coordinate4D::from_point_and_time(*point, *t);
            let temp = pinn.forward(&coord);
            let (dt_dx, dt_dy, dt_dz, dt_dt) = pinn.gradients(&coord, eps);
            let (d2t_dx2, d2t_dy2, d2t_dz2) = pinn.second_derivatives(&coord, eps);

            let derivs = TemperatureDerivatives {
                t: temp as f32,
                dt_dt: dt_dt as f32,
                dt_dx: dt_dx as f32,
                dt_dy: dt_dy as f32,
                dt_dz: dt_dz as f32,
                d2t_dx2: d2t_dx2 as f32,
                d2t_dy2: d2t_dy2 as f32,
                d2t_dz2: d2t_dz2 as f32,
            };

            let heat_source = self.probe_source.heat_source_at(point);
            let residual = self.pennes.compute(&derivs, heat_source);
            total_residual_sq += residual * residual;
        }

        (total_residual_sq / points.len() as f32).sqrt()
    }

    /// Compute boundary loss with a specific PINN instance
    fn compute_boundary_loss_with_pinn(
        &self,
        pinn: &ThermalPinn,
        points: &[(Point3D, f32)],
    ) -> f32 {
        let target_temp = self.config.physics.body_temperature;
        let mut total_error_sq = 0.0;

        for (point, t) in points {
            let coord = Coordinate4D::from_point_and_time(*point, *t);
            let predicted_temp = pinn.forward(&coord) as f32;
            let error = predicted_temp - target_temp;
            total_error_sq += error * error;
        }

        (total_error_sq / points.len() as f32).sqrt()
    }

    /// Compute initial loss with a specific PINN instance
    fn compute_initial_loss_with_pinn(&self, pinn: &ThermalPinn, points: &[(Point3D, f32)]) -> f32 {
        let target_temp = self.config.physics.body_temperature;
        let mut total_error_sq = 0.0;

        for (point, t) in points {
            let coord = Coordinate4D::from_point_and_time(*point, *t);
            let predicted_temp = pinn.forward(&coord) as f32;
            let error = predicted_temp - target_temp;
            total_error_sq += error * error;
        }

        (total_error_sq / points.len() as f32).sqrt()
    }

    /// Compute probe loss with a specific PINN instance
    fn compute_probe_loss_with_pinn(&self, pinn: &ThermalPinn, points: &[(Point3D, f32)]) -> f32 {
        let mut total_loss = 0.0;
        let mut count = 0;

        for (point, t) in points {
            let heat_source = self.probe_source.heat_source_at(point);
            if heat_source > 1e3 {
                let coord = Coordinate4D::from_point_and_time(*point, *t);
                let predicted_temp = pinn.forward(&coord) as f32;

                if predicted_temp < self.config.physics.body_temperature {
                    let error = self.config.physics.body_temperature - predicted_temp;
                    total_loss += error * error;
                }
                count += 1;
            }
        }

        if count > 0 {
            (total_loss / count as f32).sqrt()
        } else {
            0.0
        }
    }

    /// Reset the solver to initial state
    pub fn reset(&mut self) {
        self.step = 0;
        self.current_time = self.config.time.t_start;
        self.loss_history.clear();
        self.train_start = None;
        self.pinn = ThermalPinn::new(self.config.network.clone());
        self.optimizer = None;
    }

    /// Get current step
    #[must_use]
    pub fn current_step(&self) -> usize {
        self.step
    }

    /// Get current time
    #[must_use]
    pub fn current_time(&self) -> f32 {
        self.current_time
    }

    /// Get loss history
    #[must_use]
    pub fn loss_history(&self) -> &[BioheatLossRecord] {
        &self.loss_history
    }

    /// Get the configuration
    #[must_use]
    pub fn config(&self) -> &BioheatConfig {
        &self.config
    }
}

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

    #[test]
    fn test_solver_creation() {
        let solver = BioheatSolver::liver_default();
        assert_eq!(solver.current_step(), 0);
    }

    #[test]
    fn test_single_step() {
        let mut solver = BioheatSolver::new(BioheatConfig::benchmark());
        let loss = solver.step();
        assert_eq!(loss.step, 1);
        assert!(loss.total_loss >= 0.0);
    }

    #[test]
    fn test_train_multiple() {
        let mut solver = BioheatSolver::new(BioheatConfig::benchmark());
        let losses = solver.train(10);
        assert_eq!(losses.len(), 10);
        assert_eq!(solver.current_step(), 10);
    }

    #[test]
    fn test_snapshot() {
        let solver = BioheatSolver::new(BioheatConfig::benchmark());
        let snapshot = solver.snapshot((8, 8, 8));
        assert_eq!(snapshot.temperature.len(), 512);
    }

    #[test]
    fn test_status() {
        let solver = BioheatSolver::liver_default();
        let status = solver.status();
        assert!(status.initialized);
        assert!(!status.training);
    }

    #[test]
    fn test_time_advancement() {
        let mut solver = BioheatSolver::liver_default();
        solver.advance_time(10.0);
        assert!((solver.current_time() - 10.0).abs() < 1e-6);

        // Should clamp at t_end
        solver.advance_time(1000.0);
        assert!(solver.current_time() <= solver.config().time.t_end);
    }

    #[test]
    fn test_reset() {
        let mut solver = BioheatSolver::new(BioheatConfig::benchmark());
        solver.train(5);
        solver.advance_time(50.0);
        solver.reset();

        assert_eq!(solver.current_step(), 0);
        assert!((solver.current_time()).abs() < 1e-6);
        assert!(solver.loss_history().is_empty());
    }

    #[test]
    fn test_probe_update() {
        let mut solver = BioheatSolver::liver_default();
        let new_pos = Point3D::new(0.01, 0.02, 0.03);
        solver.update_probe_position(new_pos);
        assert!((solver.config().probe.position.x - 0.01).abs() < 1e-6);
    }

    // ============================================================================
    // TDD Tests for Weight Updates (RED phase - these should FAIL initially)
    // ============================================================================

    #[test]
    fn test_weights_change_after_training_step() {
        // Test that neural network weights actually change after a training step
        let mut solver = BioheatSolver::new(BioheatConfig::benchmark());

        // Capture initial weight snapshot
        let initial_weights = capture_network_weights(&solver.pinn);

        // Perform a training step
        solver.step();

        // Capture weights after training
        let updated_weights = capture_network_weights(&solver.pinn);

        // At least some weights should have changed
        let weights_changed = weights_are_different(&initial_weights, &updated_weights);
        assert!(
            weights_changed,
            "Neural network weights must change after a training step"
        );
    }

    #[test]
    fn test_loss_decreases_over_training() {
        // Test that loss decreases over multiple training epochs
        let mut solver = BioheatSolver::new(BioheatConfig::benchmark());

        // Train for several steps
        let losses = solver.train(20);

        // Loss should generally decrease (allowing for some noise)
        // Check that average loss in second half is less than first half
        let first_half_avg: f32 = losses.iter().take(10).map(|r| r.total_loss).sum::<f32>() / 10.0;
        let second_half_avg: f32 = losses.iter().skip(10).map(|r| r.total_loss).sum::<f32>() / 10.0;

        assert!(
            second_half_avg < first_half_avg,
            "Loss should decrease over training: first_half={}, second_half={}",
            first_half_avg,
            second_half_avg
        );
    }

    #[test]
    fn test_gradient_computation_is_nonzero() {
        // Test that computed gradients are non-zero for typical loss values
        let solver = BioheatSolver::new(BioheatConfig::benchmark());

        // Sample a few training points
        let mut rng = ChaCha8Rng::seed_from_u64(123);
        let points = solver.domain.sample_interior_batch(&mut rng, 5);

        // Compute loss
        let loss = solver.compute_physics_loss(&points);

        // Loss should be positive (network starts random)
        assert!(
            loss > 0.0,
            "Physics loss should be positive for untrained network"
        );

        // Compute gradients
        let grads = solver.compute_network_gradients();

        // At least some gradients should be non-zero
        let has_nonzero_grads = grads.weights.iter().any(|layer_weights: &Vec<Vec<f64>>| {
            layer_weights
                .iter()
                .any(|row: &Vec<f64>| row.iter().any(|&g: &f64| g.abs() > 1e-8))
        });

        assert!(
            has_nonzero_grads,
            "Gradients should be non-zero for non-zero loss"
        );
    }

    #[test]
    fn test_optimizer_state_updates() {
        // Test that optimizer maintains state (momentum, variance) across steps
        let mut solver = BioheatSolver::new(BioheatConfig::benchmark());

        // Initial optimizer state should exist after first step
        solver.step();

        // Check optimizer has non-zero state (this will be implemented)
        assert!(
            solver.optimizer.is_some(),
            "Optimizer should be initialized after first training step"
        );

        let opt = solver.optimizer.as_ref().unwrap();

        // After first step, optimizer should have accumulated some state
        assert_eq!(opt.t, 1, "Optimizer timestep should be 1 after one step");
    }

    #[test]
    fn test_learning_rate_is_applied() {
        // Test that different learning rates produce different weight updates
        let mut config1 = BioheatConfig::benchmark();
        config1.training.learning_rate = 0.001;

        let mut config2 = BioheatConfig::benchmark();
        config2.training.learning_rate = 0.01; // 10x larger

        let mut solver1 = BioheatSolver::new(config1);
        let mut solver2 = BioheatSolver::new(config2);

        // Set same random seed for reproducibility
        solver1.rng = ChaCha8Rng::seed_from_u64(42);
        solver2.rng = ChaCha8Rng::seed_from_u64(42);

        // Capture initial weights (should be same due to seed in new())
        let initial1 = capture_network_weights(&solver1.pinn);
        let initial2 = capture_network_weights(&solver2.pinn);

        // Perform one training step
        solver1.step();
        solver2.step();

        // Capture updated weights
        let updated1 = capture_network_weights(&solver1.pinn);
        let updated2 = capture_network_weights(&solver2.pinn);

        // Compute weight change magnitudes
        let change1 = compute_weight_change_magnitude(&initial1, &updated1);
        let change2 = compute_weight_change_magnitude(&initial2, &updated2);

        // Higher learning rate should produce larger weight changes
        assert!(
            change2 > change1,
            "Higher learning rate should produce larger weight updates: lr=0.01 change={} vs lr=0.001 change={}",
            change2,
            change1
        );
    }

    // ============================================================================
    // Helper functions for tests
    // ============================================================================

    /// Capture a snapshot of network weights for comparison
    fn capture_network_weights(pinn: &ThermalPinn) -> NetworkWeightSnapshot {
        // Clone the network to access weights
        let mut pinn_clone = pinn.clone();
        let weights = pinn_clone.weights_mut().clone();
        let biases = pinn_clone.biases_mut().clone();
        let fourier_b = pinn_clone.fourier_b_mut().clone();

        NetworkWeightSnapshot {
            weights,
            biases,
            fourier_b,
        }
    }

    /// Check if two weight snapshots are different
    fn weights_are_different(w1: &NetworkWeightSnapshot, w2: &NetworkWeightSnapshot) -> bool {
        // Check if any weights have changed
        for (layer1, layer2) in w1.weights.iter().zip(&w2.weights) {
            for (row1, row2) in layer1.iter().zip(layer2) {
                for (val1, val2) in row1.iter().zip(row2) {
                    if (val1 - val2).abs() > 1e-10 {
                        return true;
                    }
                }
            }
        }

        // Check biases
        for (b1, b2) in w1.biases.iter().zip(&w2.biases) {
            for (val1, val2) in b1.iter().zip(b2) {
                if (val1 - val2).abs() > 1e-10 {
                    return true;
                }
            }
        }

        // Check Fourier features
        for (f1, f2) in w1.fourier_b.iter().zip(&w2.fourier_b) {
            for (val1, val2) in f1.iter().zip(f2) {
                if (val1 - val2).abs() > 1e-10 {
                    return true;
                }
            }
        }

        false
    }

    /// Compute magnitude of weight change
    fn compute_weight_change_magnitude(
        before: &NetworkWeightSnapshot,
        after: &NetworkWeightSnapshot,
    ) -> f64 {
        let mut total_change = 0.0;

        // Sum squared differences for weights
        for (layer1, layer2) in before.weights.iter().zip(&after.weights) {
            for (row1, row2) in layer1.iter().zip(layer2) {
                for (val1, val2) in row1.iter().zip(row2) {
                    let diff = val1 - val2;
                    total_change += diff * diff;
                }
            }
        }

        // Sum squared differences for biases
        for (b1, b2) in before.biases.iter().zip(&after.biases) {
            for (val1, val2) in b1.iter().zip(b2) {
                let diff = val1 - val2;
                total_change += diff * diff;
            }
        }

        // Sum squared differences for Fourier features
        for (f1, f2) in before.fourier_b.iter().zip(&after.fourier_b) {
            for (val1, val2) in f1.iter().zip(f2) {
                let diff = val1 - val2;
                total_change += diff * diff;
            }
        }

        total_change.sqrt()
    }

    /// Snapshot of network weights for testing
    #[derive(Clone)]
    struct NetworkWeightSnapshot {
        weights: Vec<Vec<Vec<f64>>>,
        biases: Vec<Vec<f64>>,
        fourier_b: Vec<Vec<f64>>,
    }
}