//! Physics-Informed Denoising Diffusion Model (PIDDM).
//!
//! This module implements the core PIDDM model that combines diffusion
//! models with physics-informed constraints for PDE solving.

use rand::Rng;
use rand::prelude::*;
use thiserror::Error;

use rtx_backend::Backend;
use rtx_tensor::GenericTensor;

use crate::scheduler::{DDPMScheduler, NoiseScheduler};
use crate::unet::DiffusionUNet;

/// Errors that can occur in PIDDM operations.
#[derive(Error, Debug)]
pub enum PIDDMError {
    /// Invalid configuration
    #[error("Invalid configuration: {0}")]
    InvalidConfig(String),

    /// Shape mismatch
    #[error("Shape mismatch: expected {expected}, got {got}")]
    ShapeMismatch {
        /// The expected shape description
        expected: String,
        /// The actual shape received
        got: String,
    },

    /// Training error
    #[error("Training error: {0}")]
    TrainingError(String),

    /// Sampling error
    #[error("Sampling error: {0}")]
    SamplingError(String),
}

/// Configuration for PIDDM model.
#[derive(Debug, Clone)]
pub struct PIDDMConfig {
    /// Weight for physics loss (λ in combined loss = L_diffusion + λ * L_physics)
    pub physics_weight: f32,
    /// Whether to use physics loss during sampling (physics guidance)
    pub physics_guided_sampling: bool,
    /// Guidance strength during sampling (if physics_guided_sampling is true)
    pub guidance_strength: f32,
    /// Whether to clamp predictions to [-1, 1]
    pub clamp_predictions: bool,
}

impl Default for PIDDMConfig {
    fn default() -> Self {
        Self {
            physics_weight: 0.1,
            physics_guided_sampling: false,
            guidance_strength: 0.0,
            clamp_predictions: true,
        }
    }
}

impl PIDDMConfig {
    /// Create config optimized for PDE solving.
    pub fn for_pde() -> Self {
        Self {
            physics_weight: 0.5, // Higher physics weight for PDE applications
            physics_guided_sampling: true,
            guidance_strength: 0.1,
            clamp_predictions: false, // PDEs may have unbounded solutions
        }
    }
}

/// Physics-Informed Denoising Diffusion Model.
///
/// PIDDM combines diffusion models with physics constraints by:
/// 1. Adding physics residual loss during training
/// 2. Optionally using physics guidance during sampling
///
/// # Training Loss
///
/// The training loss combines diffusion loss and physics loss:
/// ```text
/// L = L_diffusion + λ * L_physics
///   = ||ε - ε_θ(x_t, t)||² + λ * ||L[x̂_0] - f||²
/// ```
///
/// where:
/// - ε is the noise added to the clean sample
/// - ε_θ is the predicted noise from the UNet
/// - x̂_0 is the predicted clean sample
/// - L[·] is the differential operator (e.g., Laplacian)
/// - f is the forcing term
pub struct PIDDM<B: Backend<FloatElem = f32>> {
    /// Noise scheduler
    scheduler: DDPMScheduler,
    /// UNet noise predictor
    unet: DiffusionUNet<B>,
    /// Configuration
    config: PIDDMConfig,
}

impl<B: Backend<FloatElem = f32>> PIDDM<B> {
    /// Create a new PIDDM model.
    pub fn new(scheduler: DDPMScheduler, unet: DiffusionUNet<B>, config: PIDDMConfig) -> Self {
        Self {
            scheduler,
            unet,
            config,
        }
    }

    /// Get reference to the scheduler.
    pub fn scheduler(&self) -> &DDPMScheduler {
        &self.scheduler
    }

    /// Get reference to the UNet.
    pub fn unet(&self) -> &DiffusionUNet<B> {
        &self.unet
    }

    /// Get reference to the configuration.
    pub fn config(&self) -> &PIDDMConfig {
        &self.config
    }

    /// Perform a training step.
    ///
    /// # Arguments
    /// * `x0` - Clean samples [batch, channels, height, width]
    /// * `physics_fn` - Function that computes physics residual from predicted x0
    /// * `device` - Device
    ///
    /// # Returns
    /// Tuple of (total_loss, diffusion_loss, physics_loss)
    pub fn training_step<F>(
        &self,
        x0: &GenericTensor<B, 4>,
        physics_fn: F,
        device: &B::Device,
    ) -> (f32, f32, f32)
    where
        F: Fn(&GenericTensor<B, 4>) -> GenericTensor<B, 4>,
    {
        let mut rng = thread_rng();

        // Sample random timestep
        let t = rng.gen_range(0..self.scheduler.num_timesteps());

        // Sample noise
        let noise = GenericTensor::<B, 4>::randn(x0.shape(), device);

        // Create noisy sample
        let xt = self.scheduler.add_noise(x0, &noise, t);

        // Predict noise
        let pred_noise = self.unet.forward(&xt, t, device);

        // Compute diffusion loss: ||ε - ε_θ||²
        let diff = noise.sub(&pred_noise);
        let diff_sq = diff.mul(&diff);
        let diffusion_loss = diff_sq.mean().to_vec()[0];

        // Predict x0 from noisy sample and predicted noise
        let pred_x0 = self.scheduler.predict_x0(&xt, &pred_noise, t);

        // Optionally clamp predictions
        let pred_x0 = if self.config.clamp_predictions {
            clamp(&pred_x0, -1.0, 1.0)
        } else {
            pred_x0
        };

        // Compute physics loss: ||L[x̂_0] - f||²
        let physics_residual = physics_fn(&pred_x0);
        let physics_sq = physics_residual.mul(&physics_residual);
        let physics_loss = physics_sq.mean().to_vec()[0];

        // Combined loss
        let total_loss = diffusion_loss + self.config.physics_weight * physics_loss;

        (total_loss, diffusion_loss, physics_loss)
    }

    /// Sample from the model using DDPM sampling.
    ///
    /// # Arguments
    /// * `shape` - Shape of samples to generate [batch, channels, height, width]
    /// * `device` - Device
    ///
    /// # Returns
    /// Generated samples
    pub fn sample(&self, shape: [usize; 4], device: &B::Device) -> GenericTensor<B, 4> {
        // Start from pure noise
        let mut x = GenericTensor::<B, 4>::randn(shape, device);

        // Reverse diffusion process
        for t in (0..self.scheduler.num_timesteps()).rev() {
            // Predict noise
            let pred_noise = self.unet.forward(&x, t, device);

            // Sample noise for stochastic step (not needed at t=0)
            let noise = if t > 0 {
                Some(GenericTensor::<B, 4>::randn(shape, device))
            } else {
                None
            };

            // Reverse step
            x = self.scheduler.step(&x, &pred_noise, t, noise.as_ref());

            // Optional: clamp intermediate results
            if self.config.clamp_predictions {
                x = clamp(&x, -1.0, 1.0);
            }
        }

        x
    }

    /// Sample with physics guidance.
    ///
    /// During sampling, we modify the predicted x0 to reduce physics residual:
    /// ```text
    /// x̂_0 = x̂_0 - γ * ∇_{x̂_0} ||L[x̂_0] - f||²
    /// ```
    ///
    /// # Arguments
    /// * `shape` - Shape of samples to generate
    /// * `physics_fn` - Function that computes physics residual
    /// * `physics_grad_fn` - Function that computes gradient of physics loss w.r.t. input
    /// * `device` - Device
    ///
    /// # Returns
    /// Generated samples satisfying physics constraints
    pub fn sample_with_guidance<F, G>(
        &self,
        shape: [usize; 4],
        physics_fn: F,
        physics_grad_fn: G,
        device: &B::Device,
    ) -> GenericTensor<B, 4>
    where
        F: Fn(&GenericTensor<B, 4>) -> GenericTensor<B, 4>,
        G: Fn(&GenericTensor<B, 4>) -> GenericTensor<B, 4>,
    {
        let mut x = GenericTensor::<B, 4>::randn(shape, device);

        for t in (0..self.scheduler.num_timesteps()).rev() {
            // Predict noise
            let pred_noise = self.unet.forward(&x, t, device);

            // Predict x0
            let mut pred_x0 = self.scheduler.predict_x0(&x, &pred_noise, t);

            // Apply physics guidance (gradient descent on physics loss)
            if self.config.physics_guided_sampling && self.config.guidance_strength > 0.0 {
                // Compute gradient of physics loss
                let grad = physics_grad_fn(&pred_x0);

                // Update x0 to reduce physics residual
                pred_x0 = pred_x0.sub(&grad.mul_scalar(self.config.guidance_strength));
            }

            // Clamp if configured
            if self.config.clamp_predictions {
                pred_x0 = clamp(&pred_x0, -1.0, 1.0);
            }

            // Reconstruct noise from corrected x0
            let alpha_cumprod = self.scheduler.alpha_cumprod(t);
            let corrected_noise = x
                .sub(&pred_x0.mul_scalar(alpha_cumprod.sqrt()))
                .mul_scalar(1.0 / (1.0 - alpha_cumprod).sqrt());

            // Sample noise for stochastic step
            let noise = if t > 0 {
                Some(GenericTensor::<B, 4>::randn(shape, device))
            } else {
                None
            };

            // Reverse step with corrected noise prediction
            x = self.scheduler.step(&x, &corrected_noise, t, noise.as_ref());
        }

        x
    }

    /// Compute the Laplacian of a 2D field using finite differences.
    ///
    /// This is a helper for common physics computations.
    /// ∇²u ≈ (u[i+1,j] + u[i-1,j] + u[i,j+1] + u[i,j-1] - 4*u[i,j]) / h²
    pub fn laplacian_2d(field: &GenericTensor<B, 4>, h: f32) -> GenericTensor<B, 4> {
        let shape = field.shape();
        let [batch, channels, height, width] = [shape[0], shape[1], shape[2], shape[3]];

        let field_vec = field.to_vec();
        let mut laplacian = vec![0.0f32; batch * channels * height * width];

        let h2 = h * h;

        for b in 0..batch {
            for c in 0..channels {
                for i in 1..height - 1 {
                    for j in 1..width - 1 {
                        let idx = |ii: usize, jj: usize| -> usize {
                            b * channels * height * width + c * height * width + ii * width + jj
                        };

                        let center = field_vec[idx(i, j)];
                        let left = field_vec[idx(i - 1, j)];
                        let right = field_vec[idx(i + 1, j)];
                        let down = field_vec[idx(i, j - 1)];
                        let up = field_vec[idx(i, j + 1)];

                        laplacian[idx(i, j)] = (left + right + down + up - 4.0 * center) / h2;
                    }
                }
            }
        }

        GenericTensor::<B, 4>::from_slice(&laplacian, shape, &field.device())
    }

    /// Create a Poisson residual function.
    ///
    /// For -∇²u = f, the residual is: R = -∇²u - f
    pub fn poisson_residual(
        u: &GenericTensor<B, 4>,
        f: &GenericTensor<B, 4>,
        h: f32,
    ) -> GenericTensor<B, 4> {
        let laplacian = Self::laplacian_2d(u, h);
        // Residual = -∇²u - f = -laplacian - f
        laplacian.mul_scalar(-1.0).sub(f)
    }
}

/// Clamp tensor values to [min, max].
fn clamp<B: Backend<FloatElem = f32>, const D: usize>(
    x: &GenericTensor<B, D>,
    min: f32,
    max: f32,
) -> GenericTensor<B, D> {
    let x_vec = x.to_vec();
    let clamped: Vec<f32> = x_vec.iter().map(|&v| v.clamp(min, max)).collect();
    GenericTensor::<B, D>::from_slice(&clamped, x.shape(), &x.device())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::unet::UNetConfig;
    use rtx_backend_cpu::{CpuBackend, CpuDevice};

    fn create_test_model() -> PIDDM<CpuBackend> {
        let device = CpuDevice::default();

        let scheduler = DDPMScheduler::new(100, 1e-4, 0.02); // Small for testing
        let unet_config = UNetConfig {
            in_channels: 1,
            out_channels: 1,
            base_channels: 16,
            hidden_mult: 2,
            time_embed_dim: 32,
        };
        let spatial_size = 64; // 8x8
        let unet = DiffusionUNet::new(unet_config, spatial_size, &device);
        let config = PIDDMConfig::default();

        PIDDM::new(scheduler, unet, config)
    }

    #[test]
    fn test_piddm_creation() {
        let model = create_test_model();
        assert_eq!(model.scheduler().num_timesteps(), 100);
    }

    #[test]
    fn test_laplacian_2d() {
        let device = CpuDevice::default();

        // Create a simple quadratic field: u(x,y) = x² + y²
        // For this field, ∇²u = 2 + 2 = 4 (constant)
        let n = 8;
        let h = 1.0 / (n - 1) as f32;

        let mut data = vec![0.0f32; n * n];
        for i in 0..n {
            for j in 0..n {
                let x = i as f32 * h;
                let y = j as f32 * h;
                data[i * n + j] = x * x + y * y;
            }
        }

        let field = GenericTensor::<CpuBackend, 4>::from_slice(&data, [1, 1, n, n], &device);

        let laplacian = PIDDM::<CpuBackend>::laplacian_2d(&field, h);
        let laplacian_vec = laplacian.to_vec();

        // Interior points should have Laplacian ≈ 4
        let center = n / 2;
        let center_val = laplacian_vec[center * n + center];
        assert!(
            (center_val - 4.0).abs() < 0.5,
            "Laplacian should be close to 4, got {}",
            center_val
        );
    }

    #[test]
    fn test_training_step() {
        let model = create_test_model();
        let device = CpuDevice::default();

        // Create dummy clean sample
        let x0 = GenericTensor::<CpuBackend, 4>::rand([1, 1, 8, 8], &device);

        // Dummy physics function (just returns zeros - no physics constraint)
        let physics_fn = |_x: &GenericTensor<CpuBackend, 4>| {
            GenericTensor::<CpuBackend, 4>::zeros([1, 1, 8, 8], &device)
        };

        let (total, diffusion, physics) = model.training_step(&x0, physics_fn, &device);

        // Losses should be non-negative
        assert!(total >= 0.0, "Total loss should be non-negative");
        assert!(diffusion >= 0.0, "Diffusion loss should be non-negative");
        assert!(physics >= 0.0, "Physics loss should be non-negative");

        // Physics loss should be 0 (since residual function returns zeros)
        assert!(
            physics < 1e-10,
            "Physics loss should be ~0 for zero residual"
        );
    }

    #[test]
    fn test_sample_shape() {
        let model = create_test_model();
        let device = CpuDevice::default();

        // Note: This is a simplified test - full sampling is slow
        // We just test that the forward pass works with the expected shapes
        let shape = [1, 1, 8, 8];

        // Test single step
        let x = GenericTensor::<CpuBackend, 4>::randn(shape, &device);
        let pred_noise = model.unet().forward(&x, 50, &device);

        let pred_shape = pred_noise.shape();
        assert_eq!(pred_shape, shape);
    }

    #[test]
    fn test_poisson_residual() {
        let device = CpuDevice::default();

        let n = 8;
        let h = 1.0 / (n - 1) as f32;

        // Create u = sin(πx)sin(πy), which satisfies -∇²u = 2π²sin(πx)sin(πy)
        let pi = std::f32::consts::PI;
        let mut u_data = vec![0.0f32; n * n];
        let mut f_data = vec![0.0f32; n * n];

        for i in 0..n {
            for j in 0..n {
                let x = i as f32 * h;
                let y = j as f32 * h;
                u_data[i * n + j] = (pi * x).sin() * (pi * y).sin();
                f_data[i * n + j] = 2.0 * pi * pi * (pi * x).sin() * (pi * y).sin();
            }
        }

        let u = GenericTensor::<CpuBackend, 4>::from_slice(&u_data, [1, 1, n, n], &device);
        let f = GenericTensor::<CpuBackend, 4>::from_slice(&f_data, [1, 1, n, n], &device);

        let residual = PIDDM::<CpuBackend>::poisson_residual(&u, &f, h);
        let residual_vec = residual.to_vec();

        // Interior residual should be small (discretization error)
        let center = n / 2;
        let center_residual = residual_vec[center * n + center];
        assert!(
            center_residual.abs() < 1.0,
            "Poisson residual should be small at center, got {}",
            center_residual
        );
    }
}