//! Learnable stiffness texture for MRE inverse problem
//!
//! The stiffness field mu(x,y) is represented as a 2D texture (grid) rather
//! than a neural network. This allows for sharp discontinuities (tumor edges)
//! that neural networks tend to over-smooth.
//!
//! Key features:
//! - Bilinear interpolation for smooth sampling between grid points
//! - Finite difference gradients (mu_x, mu_y) for physics loss
//! - Direct gradient descent on texture pixels

use crate::config::MreConfig;
use anyhow::Result;
use mre_shared::StiffnessField;
use rtx_tensor::{Device, Tensor};

/// Learnable stiffness texture
pub struct StiffnessTexture {
    /// Stiffness values [ny, nx] in non-dimensional units
    values: Tensor,
    /// Grid dimensions
    nx: usize,
    ny: usize,
    /// Physical domain bounds (non-dimensional: [0, 1] x [0, 1])
    x_min: f32,
    x_max: f32,
    y_min: f32,
    y_max: f32,
    /// Minimum allowed stiffness (non-dimensional)
    min_stiffness: f32,
    /// Device
    device: Device,
}

impl StiffnessTexture {
    /// Create a new stiffness texture with uniform initial value
    pub fn new(config: &MreConfig, device: &Device) -> Result<Self> {
        let nx = config.stiffness_nx;
        let ny = config.stiffness_ny;

        // Initialize with uniform value = 1.0 (reference stiffness in non-dim)
        let init_value = 1.0;
        let n = nx * ny;
        let data = vec![init_value; n];
        let values = Tensor::from_data(data, vec![ny, nx], device)?;

        Ok(Self {
            values,
            nx,
            ny,
            x_min: 0.0,
            x_max: 1.0, // Non-dimensional domain
            y_min: 0.0,
            y_max: 1.0,
            min_stiffness: config.min_stiffness,
            device: device.clone(),
        })
    }

    /// Initialize from a ground truth stiffness field
    pub fn from_field(field: &StiffnessField, config: &MreConfig, device: &Device) -> Result<Self> {
        let (nx, ny) = field.resolution;

        // Convert to non-dimensional values
        let values_nondim: Vec<f32> = field
            .values
            .iter()
            .map(|&v| config.nondim.nondim_stiffness_kpa(v))
            .collect();

        let values = Tensor::from_data(values_nondim, vec![ny, nx], device)?;

        Ok(Self {
            values,
            nx,
            ny,
            x_min: 0.0,
            x_max: 1.0,
            y_min: 0.0,
            y_max: 1.0,
            min_stiffness: config.min_stiffness,
            device: device.clone(),
        })
    }

    /// Sample stiffness at arbitrary points using bilinear interpolation
    ///
    /// # Arguments
    /// * `x` - x-coordinates [batch] (non-dimensional, in [0, 1])
    /// * `y` - y-coordinates [batch]
    ///
    /// # Returns
    /// Tensor [batch] with interpolated stiffness values
    pub fn sample(&self, x: &Tensor, y: &Tensor) -> Result<Tensor> {
        let batch_size = x.shape().dims()[0];
        let x_data = x.to_cpu()?;
        let y_data = y.to_cpu()?;
        let values_data = self.values.to_cpu()?;

        let mut result = Vec::with_capacity(batch_size);

        for i in 0..batch_size {
            let xi = x_data[i];
            let yi = y_data[i];

            let mu = self.bilinear_interpolate(&values_data, xi, yi);
            result.push(mu.max(self.min_stiffness));
        }

        Ok(Tensor::from_data(result, vec![batch_size], &self.device)?)
    }

    /// Sample stiffness and its spatial gradients
    ///
    /// Uses finite differences on the texture for gradient computation.
    ///
    /// # Returns
    /// (mu, mu_x, mu_y) - stiffness and its gradients
    pub fn sample_with_gradients(
        &self,
        x: &Tensor,
        y: &Tensor,
    ) -> Result<(Tensor, Tensor, Tensor)> {
        let batch_size = x.shape().dims()[0];
        let x_data = x.to_cpu()?;
        let y_data = y.to_cpu()?;
        let values_data = self.values.to_cpu()?;

        let dx = (self.x_max - self.x_min) / (self.nx - 1).max(1) as f32;
        let dy = (self.y_max - self.y_min) / (self.ny - 1).max(1) as f32;

        let mut mu_result = Vec::with_capacity(batch_size);
        let mut mu_x_result = Vec::with_capacity(batch_size);
        let mut mu_y_result = Vec::with_capacity(batch_size);

        for i in 0..batch_size {
            let xi = x_data[i];
            let yi = y_data[i];

            // Sample mu at current point
            let mu = self.bilinear_interpolate(&values_data, xi, yi);

            // Central difference for gradients
            let mu_xp = self.bilinear_interpolate(&values_data, xi + dx, yi);
            let mu_xm = self.bilinear_interpolate(&values_data, xi - dx, yi);
            let mu_yp = self.bilinear_interpolate(&values_data, xi, yi + dy);
            let mu_ym = self.bilinear_interpolate(&values_data, xi, yi - dy);

            let mu_x = (mu_xp - mu_xm) / (2.0 * dx);
            let mu_y = (mu_yp - mu_ym) / (2.0 * dy);

            mu_result.push(mu.max(self.min_stiffness));
            mu_x_result.push(mu_x);
            mu_y_result.push(mu_y);
        }

        let mu = Tensor::from_data(mu_result, vec![batch_size], &self.device)?;
        let mu_x = Tensor::from_data(mu_x_result, vec![batch_size], &self.device)?;
        let mu_y = Tensor::from_data(mu_y_result, vec![batch_size], &self.device)?;

        Ok((mu, mu_x, mu_y))
    }

    /// Bilinear interpolation helper
    fn bilinear_interpolate(&self, data: &[f32], x: f32, y: f32) -> f32 {
        // Convert to texture coordinates
        let tx = (x - self.x_min) / (self.x_max - self.x_min);
        let ty = (y - self.y_min) / (self.y_max - self.y_min);

        // Scale to grid indices
        let fx = tx * (self.nx - 1) as f32;
        let fy = ty * (self.ny - 1) as f32;

        // Floor indices
        let ix0 = (fx.floor() as usize).min(self.nx - 1);
        let iy0 = (fy.floor() as usize).min(self.ny - 1);
        let ix1 = (ix0 + 1).min(self.nx - 1);
        let iy1 = (iy0 + 1).min(self.ny - 1);

        // Fractional parts (weights)
        let wx1 = fx - ix0 as f32;
        let wy1 = fy - iy0 as f32;
        let wx0 = 1.0 - wx1;
        let wy0 = 1.0 - wy1;

        // Gather values at four corners
        let v00 = data[iy0 * self.nx + ix0];
        let v01 = data[iy0 * self.nx + ix1];
        let v10 = data[iy1 * self.nx + ix0];
        let v11 = data[iy1 * self.nx + ix1];

        // Bilinear interpolation
        v00 * wx0 * wy0 + v01 * wx1 * wy0 + v10 * wx0 * wy1 + v11 * wx1 * wy1
    }

    /// Apply gradient update to texture values
    ///
    /// # Arguments
    /// * `grad` - Gradient tensor [ny, nx]
    /// * `lr` - Learning rate
    pub fn apply_gradient(&mut self, grad: &Tensor, lr: f32) -> Result<()> {
        // values = values - lr * grad
        let update = grad.mul_scalar(lr)?;
        self.values = self.values.sub(&update)?;

        // Clamp to minimum stiffness
        self.clamp_values()?;

        Ok(())
    }

    /// Compute gradient of stiffness texture from physics residual
    ///
    /// Given the physics residual R and its gradient dR/dmu, compute
    /// the gradient dL/d(texture) by scattering residual contributions
    /// back to texture pixels.
    ///
    /// # Arguments
    /// * `x` - sample x-coordinates [batch]
    /// * `y` - sample y-coordinates [batch]
    /// * `grad_mu` - gradient dL/dmu at each sample point [batch]
    pub fn compute_texture_gradient(
        &self,
        x: &Tensor,
        y: &Tensor,
        grad_mu: &Tensor,
    ) -> Result<Tensor> {
        let batch_size = x.shape().dims()[0];
        let x_data = x.to_cpu()?;
        let y_data = y.to_cpu()?;
        let grad_data = grad_mu.to_cpu()?;

        // Accumulate gradients using bilinear splatting (inverse of interpolation)
        let mut texture_grad = vec![0.0f32; self.ny * self.nx];

        for i in 0..batch_size {
            let xi = x_data[i];
            let yi = y_data[i];
            let grad = grad_data[i];

            // Convert to texture coordinates
            let tx = (xi - self.x_min) / (self.x_max - self.x_min);
            let ty = (yi - self.y_min) / (self.y_max - self.y_min);

            let fx = tx * (self.nx - 1) as f32;
            let fy = ty * (self.ny - 1) as f32;

            let ix0 = (fx.floor() as usize).min(self.nx - 1);
            let iy0 = (fy.floor() as usize).min(self.ny - 1);
            let ix1 = (ix0 + 1).min(self.nx - 1);
            let iy1 = (iy0 + 1).min(self.ny - 1);

            let wx1 = fx - ix0 as f32;
            let wy1 = fy - iy0 as f32;
            let wx0 = 1.0 - wx1;
            let wy0 = 1.0 - wy1;

            // Splat gradient to four corners (inverse bilinear)
            texture_grad[iy0 * self.nx + ix0] += grad * wx0 * wy0;
            texture_grad[iy0 * self.nx + ix1] += grad * wx1 * wy0;
            texture_grad[iy1 * self.nx + ix0] += grad * wx0 * wy1;
            texture_grad[iy1 * self.nx + ix1] += grad * wx1 * wy1;
        }

        Ok(Tensor::from_data(
            texture_grad,
            vec![self.ny, self.nx],
            &self.device,
        )?)
    }

    /// Clamp values to minimum stiffness
    fn clamp_values(&mut self) -> Result<()> {
        let data = self.values.to_cpu()?;
        let clamped: Vec<f32> = data.iter().map(|&v| v.max(self.min_stiffness)).collect();
        self.values = Tensor::from_data(clamped, vec![self.ny, self.nx], &self.device)?;
        Ok(())
    }

    /// Convert to StiffnessField for visualization
    pub fn to_field(&self, config: &MreConfig) -> Result<StiffnessField> {
        let data = self.values.to_cpu()?;

        // Convert from non-dim to kPa
        let values_kpa: Vec<f32> = data
            .iter()
            .map(|&v| config.nondim.dim_stiffness_kpa(v))
            .collect();

        Ok(StiffnessField {
            resolution: (self.nx, self.ny),
            values: values_kpa,
            bounds: (
                config.nondim.dim_length(self.x_min),
                config.nondim.dim_length(self.x_max),
                config.nondim.dim_length(self.y_min),
                config.nondim.dim_length(self.y_max),
            ),
        })
    }

    /// Get grid dimensions
    pub fn resolution(&self) -> (usize, usize) {
        (self.nx, self.ny)
    }

    /// Get raw values tensor
    pub fn values(&self) -> &Tensor {
        &self.values
    }

    /// Compute Total Variation (TV) regularization loss
    ///
    /// TV = sum(|mu(i+1,j) - mu(i,j)| + |mu(i,j+1) - mu(i,j)|)
    ///
    /// Encourages piecewise-constant stiffness maps (sharp edges)
    pub fn tv_loss(&self) -> Result<f32> {
        let data = self.values.to_cpu()?;
        let mut tv = 0.0f32;

        for j in 0..self.ny {
            for i in 0..self.nx {
                let idx = j * self.nx + i;
                let v = data[idx];

                // Horizontal difference
                if i + 1 < self.nx {
                    let v_right = data[idx + 1];
                    tv += (v_right - v).abs();
                }

                // Vertical difference
                if j + 1 < self.ny {
                    let v_below = data[idx + self.nx];
                    tv += (v_below - v).abs();
                }
            }
        }

        Ok(tv)
    }
}

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

    fn get_test_config() -> MreConfig {
        MreConfig::fast()
    }

    #[test]
    fn test_texture_creation() {
        let config = get_test_config();
        let device = Device::try_default().unwrap();
        let texture = StiffnessTexture::new(&config, &device).unwrap();

        assert_eq!(
            texture.resolution(),
            (config.stiffness_nx, config.stiffness_ny)
        );
    }

    #[test]
    fn test_uniform_sampling() {
        let config = get_test_config();
        let device = Device::try_default().unwrap();
        let texture = StiffnessTexture::new(&config, &device).unwrap();

        let x = Tensor::from_data(vec![0.25, 0.5, 0.75], vec![3], &device).unwrap();
        let y = Tensor::from_data(vec![0.25, 0.5, 0.75], vec![3], &device).unwrap();

        let mu = texture.sample(&x, &y).unwrap();
        let mu_data = mu.to_cpu().unwrap();

        // All should be 1.0 (initial uniform value)
        for val in &mu_data {
            assert!((*val - 1.0).abs() < 1e-5);
        }
    }

    #[test]
    fn test_bilinear_interpolation() {
        let config = MreConfig::fast().with_stiffness_resolution(3, 3);
        let device = Device::try_default().unwrap();

        // Create texture with known pattern
        let mut texture = StiffnessTexture::new(&config, &device).unwrap();

        // Set corners to different values
        // [0,0]=1, [0,2]=2, [2,0]=3, [2,2]=4
        let values = vec![1.0, 1.5, 2.0, 2.0, 2.5, 3.0, 3.0, 3.5, 4.0];
        texture.values = Tensor::from_data(values, vec![3, 3], &device).unwrap();

        // Sample at center (0.5, 0.5) should interpolate
        let x = Tensor::from_data(vec![0.5], vec![1], &device).unwrap();
        let y = Tensor::from_data(vec![0.5], vec![1], &device).unwrap();
        let mu = texture.sample(&x, &y).unwrap();
        let mu_val = mu.to_cpu().unwrap()[0];

        // Center value should be 2.5
        assert!((mu_val - 2.5).abs() < 1e-5);
    }

    #[test]
    fn test_gradient_computation() {
        let config = get_test_config();
        let device = Device::try_default().unwrap();
        let texture = StiffnessTexture::new(&config, &device).unwrap();

        let x = Tensor::from_data(vec![0.5], vec![1], &device).unwrap();
        let y = Tensor::from_data(vec![0.5], vec![1], &device).unwrap();

        let (mu, mu_x, mu_y) = texture.sample_with_gradients(&x, &y).unwrap();

        // For uniform texture, gradients should be ~0
        let mu_x_val = mu_x.to_cpu().unwrap()[0];
        let mu_y_val = mu_y.to_cpu().unwrap()[0];

        assert!(mu_x_val.abs() < 1e-5);
        assert!(mu_y_val.abs() < 1e-5);
    }
}