//! MRE inverse solver combining Wave Net and Stiffness Texture
//!
//! The solver trains both components jointly to minimize:
//! - Data loss: Match the measured wave field
//! - Physics loss: Satisfy the Helmholtz equation

use crate::config::MreConfig;
use crate::helmholtz::{HelmholtzResidual, compute_data_loss};
use crate::phantom::{generate_grid_coords, generate_random_coords};
use crate::stiffness_texture::StiffnessTexture;
use crate::wave_net::WaveNet;
use anyhow::Result;
use mre_shared::{LossRecord, MreSnapshot, StiffnessField, WaveField};
use rtx_tensor::{Device, Tensor};

/// MRE inverse solver
pub struct MreSolver {
    /// Configuration
    config: MreConfig,
    /// Wave field neural network
    wave_net: WaveNet,
    /// Learnable stiffness texture
    stiffness: StiffnessTexture,
    /// Helmholtz residual computer
    helmholtz: HelmholtzResidual,
    /// Measured wave field (input data)
    measured_wave: Option<WaveField>,
    /// Current training step
    step: usize,
    /// Loss history
    loss_history: Vec<LossRecord>,
    /// Device
    device: Device,
}

impl MreSolver {
    /// Create a new MRE solver
    pub fn new(config: MreConfig) -> Result<Self> {
        let device = Device::try_default()?; // Use default device (CUDA if available)

        let wave_net = WaveNet::new(&config, &device)?;
        let stiffness = StiffnessTexture::new(&config, &device)?;
        let helmholtz = HelmholtzResidual::new(&config, &device);

        Ok(Self {
            config,
            wave_net,
            stiffness,
            helmholtz,
            measured_wave: None,
            step: 0,
            loss_history: Vec::new(),
            device,
        })
    }

    /// Initialize with measured wave field data
    pub fn set_measured_wave(&mut self, wave: WaveField) {
        self.measured_wave = Some(wave);
    }

    /// Initialize stiffness from a known field (for comparison)
    pub fn set_initial_stiffness(&mut self, field: &StiffnessField) -> Result<()> {
        self.stiffness = StiffnessTexture::from_field(field, &self.config, &self.device)?;
        Ok(())
    }

    /// Run a single training step
    ///
    /// # Returns
    /// LossRecord containing total, data, and physics loss
    pub fn step(&mut self) -> Result<LossRecord> {
        let wave = self
            .measured_wave
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("No measured wave field set"))?;

        // 1. Sample collocation points for physics loss
        let (x_coll, y_coll) = self.sample_collocation_points()?;

        // 2. Forward pass through Wave Net with derivatives
        let wave_derivs = self.wave_net.forward_with_derivs(&x_coll, &y_coll)?;

        // 3. Sample stiffness and gradients at collocation points
        let (mu, mu_x, mu_y) = self.stiffness.sample_with_gradients(&x_coll, &y_coll)?;

        // 4. Compute physics residual
        let (r_real, r_imag) = self
            .helmholtz
            .compute_residual(&wave_derivs, &mu, &mu_x, &mu_y)?;
        let physics_loss = self.helmholtz.physics_loss(&r_real, &r_imag)?;

        // 5. Sample data points and compute data loss
        let (x_data, y_data, meas_real, meas_imag) = self.sample_data_points(wave)?;
        let (pred_real, pred_imag) = self.wave_net.forward(&x_data, &y_data)?;
        let data_loss = compute_data_loss(&pred_real, &pred_imag, &meas_real, &meas_imag)?;

        // 6. Total loss (weighted sum)
        let total_loss = self.config.physics_weight * physics_loss
            + self.config.data_weight * data_loss
            + self.config.tv_weight * self.stiffness.tv_loss()?;

        // 7. Update stiffness texture using gradient descent
        // Compute gradient of physics loss w.r.t stiffness
        let stiffness_grad =
            self.helmholtz
                .compute_stiffness_gradient(&wave_derivs, &r_real, &r_imag)?;

        // Scatter gradients back to texture
        let texture_grad =
            self.stiffness
                .compute_texture_gradient(&x_coll, &y_coll, &stiffness_grad)?;

        // Apply gradient update
        self.stiffness
            .apply_gradient(&texture_grad, self.config.learning_rate_stiffness)?;

        // 8. Record loss
        self.step += 1;
        let record = LossRecord {
            step: self.step,
            total_loss,
            data_loss,
            physics_loss,
        };
        self.loss_history.push(record.clone());

        Ok(record)
    }

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

    /// Get current snapshot for visualization
    pub fn snapshot(&self, grid_nx: usize, grid_ny: usize) -> Result<MreSnapshot> {
        // Generate regular grid in non-dimensional coordinates
        let (x_grid, y_grid) =
            generate_grid_coords(grid_nx, grid_ny, 0.0, 1.0, 0.0, 1.0, &self.device)?;

        // Evaluate wave field
        let (wave_real, wave_imag) = self.wave_net.forward(&x_grid, &y_grid)?;
        let wave_real_data = wave_real.to_cpu()?;
        let wave_imag_data = wave_imag.to_cpu()?;

        // Evaluate residual
        let wave_derivs = self.wave_net.forward_with_derivs(&x_grid, &y_grid)?;
        let (mu, mu_x, mu_y) = self.stiffness.sample_with_gradients(&x_grid, &y_grid)?;
        let (r_real, r_imag) = self
            .helmholtz
            .compute_residual(&wave_derivs, &mu, &mu_x, &mu_y)?;

        // Residual magnitude
        let r_real_data = r_real.to_cpu()?;
        let r_imag_data = r_imag.to_cpu()?;
        let residual: Vec<f32> = r_real_data
            .iter()
            .zip(&r_imag_data)
            .map(|(&r, &i)| (r * r + i * i).sqrt())
            .collect();

        // Get stiffness field
        let stiffness = self.stiffness.to_field(&self.config)?;

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

        Ok(MreSnapshot {
            stiffness,
            wave_real: wave_real_data,
            wave_imag: wave_imag_data,
            residual,
            grid_resolution: (grid_nx, grid_ny),
            step: self.step,
            loss,
        })
    }

    /// Sample random collocation points for physics loss
    fn sample_collocation_points(&self) -> Result<(Tensor, Tensor)> {
        generate_random_coords(
            self.config.num_collocation,
            0.0,
            1.0,
            0.0,
            1.0,
            &self.device,
        )
    }

    /// Sample data points from measured wave field
    fn sample_data_points(&self, wave: &WaveField) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
        let n = wave.len().min(self.config.num_data_points);

        // Random sample indices using rand crate
        use rand::seq::SliceRandom;
        let mut rng = rand::thread_rng();
        let mut indices: Vec<usize> = (0..wave.len()).collect();
        indices.shuffle(&mut rng);
        indices.truncate(n);

        // Extract points and values
        let mut x_data = Vec::with_capacity(n);
        let mut y_data = Vec::with_capacity(n);
        let mut real_data = Vec::with_capacity(n);
        let mut imag_data = Vec::with_capacity(n);

        for &idx in &indices {
            let pt = &wave.points[idx];
            // Convert to non-dimensional
            let x_nd = self.config.nondim.nondim_length(pt.x);
            let y_nd = self.config.nondim.nondim_length(pt.y);

            x_data.push(x_nd);
            y_data.push(y_nd);
            real_data.push(wave.real[idx]);
            imag_data.push(wave.imag[idx]);
        }

        let x = Tensor::from_data(x_data, vec![n], &self.device)?;
        let y = Tensor::from_data(y_data, vec![n], &self.device)?;
        let real = Tensor::from_data(real_data, vec![n], &self.device)?;
        let imag = Tensor::from_data(imag_data, vec![n], &self.device)?;

        Ok((x, y, real, imag))
    }

    /// Get current training step
    pub fn current_step(&self) -> usize {
        self.step
    }

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

    /// Reset solver to initial state
    pub fn reset(&mut self) -> Result<()> {
        self.wave_net = WaveNet::new(&self.config, &self.device)?;
        self.stiffness = StiffnessTexture::new(&self.config, &self.device)?;
        self.step = 0;
        self.loss_history.clear();
        Ok(())
    }

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

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

    fn get_test_config() -> MreConfig {
        MreConfig::fast()
            .with_stiffness_resolution(16, 16)
            .with_wave_net_layers(2)
            .with_wave_net_hidden(32)
            .with_fourier_features(8)
    }

    #[test]
    fn test_solver_creation() {
        let config = get_test_config();
        let solver = MreSolver::new(config).unwrap();
        assert_eq!(solver.current_step(), 0);
    }

    #[test]
    fn test_solver_with_phantom() {
        let config = get_test_config();
        let mut solver = MreSolver::new(config.clone()).unwrap();

        // Generate phantom
        let phantom = PhantomGenerator::tumor_phantom(config);
        let (_stiffness, wave) = phantom.generate();

        solver.set_measured_wave(wave);

        // Run a few steps
        for _ in 0..3 {
            let loss = solver.step().unwrap();
            assert!(loss.total_loss.is_finite());
        }

        assert_eq!(solver.current_step(), 3);
    }

    #[test]
    fn test_snapshot() {
        let config = get_test_config();
        let mut solver = MreSolver::new(config.clone()).unwrap();

        let phantom = PhantomGenerator::tumor_phantom(config);
        let (_stiffness, wave) = phantom.generate();
        solver.set_measured_wave(wave);

        let snapshot = solver.snapshot(16, 16).unwrap();

        assert_eq!(snapshot.grid_resolution, (16, 16));
        assert_eq!(snapshot.wave_real.len(), 256);
        assert_eq!(snapshot.stiffness.resolution, (16, 16));
    }
}