//! DDIM (Denoising Diffusion Implicit Models) sampler.
//!
//! Provides deterministic, fast sampling from a trained diffusion model
//! using a small number of steps (e.g., 10) instead of the full T=1000.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::{gpu_add, gpu_affine, gpu_sub, GpuTensor};

use super::denoiser::Denoiser;
use super::noise::NoiseScheduler;

/// DDIM sampler for fast, deterministic inference.
#[derive(Debug)]
pub struct DDIMSampler {
    /// Number of DDIM steps (much less than training timesteps).
    num_steps: usize,
    /// Total training timesteps.
    num_timesteps: usize,
    /// DDIM eta parameter: 0.0 = deterministic, 1.0 = DDPM.
    eta: f32,
}

impl DDIMSampler {
    pub fn new(num_steps: usize, num_timesteps: usize, eta: f32) -> Self {
        Self {
            num_steps: num_steps.max(1),
            num_timesteps,
            eta,
        }
    }

    /// Compute the uniformly spaced timestep subsequence for DDIM.
    fn timestep_sequence(&self) -> Vec<usize> {
        let mut seq = Vec::with_capacity(self.num_steps);
        for i in 0..self.num_steps {
            let t = ((self.num_timesteps - 1) as f64
                * (1.0 - i as f64 / self.num_steps.max(1) as f64))
                as usize;
            seq.push(t.min(self.num_timesteps - 1));
        }
        seq
    }

    /// Generate samples from pure noise using DDIM.
    pub fn sample(
        &self,
        denoiser: &Denoiser,
        scheduler: &NoiseScheduler,
        num_samples: usize,
        data_dim: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        // Start from pure noise
        let mut x_t = GpuTensor::randn(&[num_samples, data_dim], 1.0, stream)?;

        let timesteps = self.timestep_sequence();

        for (i, &t) in timesteps.iter().enumerate() {
            let t_vec = vec![t as u32; num_samples];

            // Predict noise
            let eps_pred = denoiser.forward(&x_t, &t_vec)?;

            let alpha_bar_t: f32 = scheduler.get_alpha_bar(t)?;

            let alpha_bar_prev: f32 = if i + 1 < timesteps.len() {
                let t_prev = timesteps.get(i + 1).copied().unwrap_or(0);
                scheduler.get_alpha_bar(t_prev)?
            } else {
                1.0_f32
            };

            let sqrt_ab = alpha_bar_t.sqrt();
            let sqrt_one_minus_ab = (1.0_f32 - alpha_bar_t).max(0.0_f32).sqrt();

            // Predicted clean sample:
            // x0 = (x_t - sqrt(1-ab)*eps) / sqrt(ab)
            let scale = 1.0_f64 / sqrt_ab.max(1e-8) as f64;
            let noise_sub = gpu_affine(&eps_pred, sqrt_one_minus_ab as f64, 0.0)?;
            let x0_pred = gpu_sub(&x_t, &noise_sub)?;
            let x0_pred = gpu_affine(&x0_pred, scale, 0.0)?;

            let sqrt_one_minus_ab_prev = (1.0_f32 - alpha_bar_prev).max(0.0_f32).sqrt();

            // Compute sigma
            let sigma: f32 = if self.eta > 0.0_f32 {
                let ratio =
                    (1.0_f32 - alpha_bar_prev) / (1.0_f32 - alpha_bar_t).max(1e-8_f32);
                let beta_t = 1.0_f32 - alpha_bar_t / alpha_bar_prev.max(1e-8_f32);
                self.eta * (ratio * beta_t).max(0.0_f32).sqrt()
            } else {
                0.0_f32
            };

            let dir_coeff = (sqrt_one_minus_ab_prev.powi(2) - sigma.powi(2))
                .max(0.0_f32)
                .sqrt();

            let sqrt_ab_prev = alpha_bar_prev.sqrt();
            let term1 = gpu_affine(&x0_pred, sqrt_ab_prev as f64, 0.0)?;
            let term2 = gpu_affine(&eps_pred, dir_coeff as f64, 0.0)?;
            x_t = gpu_add(&term1, &term2)?;

            // Add stochastic noise if eta > 0
            if sigma > 1e-8_f32 {
                let noise = GpuTensor::randn(&[num_samples, data_dim], 1.0, stream)?;
                let noise_term = gpu_affine(&noise, sigma as f64, 0.0)?;
                x_t = gpu_add(&x_t, &noise_term)?;
            }
        }

        Ok(x_t)
    }
}

#[cfg(test)]
#[allow(clippy::assertions_on_result_states)]
mod tests {
    use super::*;
    use super::super::config::{DiffusionConfig, NoiseSchedule};
    use cudarc::driver::CudaContext;

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA required");
        ctx.new_stream().expect("Failed to create stream")
    }

    fn make_test_components(
        stream: &Arc<CudaStream>,
    ) -> (Denoiser, NoiseScheduler, DDIMSampler) {
        let config = DiffusionConfig {
            num_timesteps: 100,
            sampling_steps: 5,
            seq_len: 16,
            feature_dim: 1,
            hidden_dim: 32,
            num_layers: 1,
            time_embed_dim: 16,
            ..Default::default()
        };

        let denoiser = Denoiser::new(
            config.data_dim(),
            config.hidden_dim,
            config.num_layers,
            config.time_embed_dim,
            stream,
        )
        .unwrap();

        let scheduler =
            NoiseScheduler::new(config.num_timesteps, &config.schedule, stream).unwrap();

        let sampler = DDIMSampler::new(config.sampling_steps, config.num_timesteps, 0.0);

        (denoiser, scheduler, sampler)
    }

    #[test]
    fn test_timestep_sequence_decreasing() {
        let sampler = DDIMSampler::new(5, 100, 0.0);
        let seq = sampler.timestep_sequence();
        assert_eq!(seq.len(), 5);
        for i in 1..seq.len() {
            assert!(
                seq.get(i).copied().unwrap_or(0) <= seq.get(i - 1).copied().unwrap_or(0),
                "Timesteps should be decreasing: {:?}",
                seq
            );
        }
    }

    #[test]
    fn test_ddim_sample_shape() {
        let stream = test_stream();
        let (denoiser, scheduler, sampler) = make_test_components(&stream);
        let samples = sampler
            .sample(&denoiser, &scheduler, 8, 16, &stream)
            .unwrap();
        assert_eq!(samples.shape, vec![8, 16]);
    }

    #[test]
    fn test_ddim_sample_finite() {
        let stream = test_stream();
        let (denoiser, scheduler, sampler) = make_test_components(&stream);
        let samples = sampler
            .sample(&denoiser, &scheduler, 4, 16, &stream)
            .unwrap();
        let vals = samples.to_vec().unwrap();
        for v in &vals {
            assert!(v.is_finite(), "Sample contains non-finite value: {v}");
        }
    }

    #[test]
    fn test_ddim_single_step() {
        let stream = test_stream();
        let denoiser = Denoiser::new(8, 16, 1, 8, &stream).unwrap();
        let scheduler = NoiseScheduler::new(10, &NoiseSchedule::Linear, &stream).unwrap();
        let sampler = DDIMSampler::new(1, 10, 0.0);
        let result = sampler.sample(&denoiser, &scheduler, 2, 8, &stream);
        assert!(result.is_ok());
    }
}
