//! Noise schedulers for diffusion models.
//!
//! This module implements DDPM and DDIM schedulers that control the noise
//! addition and removal process in diffusion models.

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

/// Configuration for noise schedulers.
#[derive(Debug, Clone)]
pub struct SchedulerConfig {
    /// Number of diffusion timesteps
    pub num_timesteps: usize,
    /// Starting value for beta schedule
    pub beta_start: f32,
    /// Ending value for beta schedule
    pub beta_end: f32,
    /// Type of beta schedule
    pub schedule_type: BetaSchedule,
}

impl Default for SchedulerConfig {
    fn default() -> Self {
        Self {
            num_timesteps: 1000,
            beta_start: 1e-4,
            beta_end: 0.02,
            schedule_type: BetaSchedule::Linear,
        }
    }
}

/// Beta schedule types for noise variance.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BetaSchedule {
    /// Linear schedule: β_t = β_start + t/T * (β_end - β_start)
    Linear,
    /// Cosine schedule (from "Improved Denoising Diffusion Probabilistic Models")
    Cosine,
    /// Scaled linear schedule
    ScaledLinear,
}

/// Trait for noise schedulers.
pub trait NoiseScheduler {
    /// Get the number of timesteps.
    fn num_timesteps(&self) -> usize;

    /// Get alpha_cumprod (ᾱ_t) at timestep t.
    fn alpha_cumprod(&self, t: usize) -> f32;

    /// Get beta (β_t) at timestep t.
    fn beta(&self, t: usize) -> f32;

    /// Add noise to clean data: x_t = √(ᾱ_t) * x_0 + √(1 - ᾱ_t) * ε
    fn add_noise<B: Backend<FloatElem = f32>, const D: usize>(
        &self,
        x0: &GenericTensor<B, D>,
        noise: &GenericTensor<B, D>,
        t: usize,
    ) -> GenericTensor<B, D> {
        let alpha_cumprod = self.alpha_cumprod(t);
        let sqrt_alpha_cumprod = alpha_cumprod.sqrt();
        let sqrt_one_minus_alpha_cumprod = (1.0 - alpha_cumprod).sqrt();

        // x_t = √(ᾱ_t) * x_0 + √(1 - ᾱ_t) * ε
        x0.mul_scalar(sqrt_alpha_cumprod)
            .add(&noise.mul_scalar(sqrt_one_minus_alpha_cumprod))
    }

    /// Predict x_0 from x_t and predicted noise.
    /// x_0 = (x_t - √(1 - ᾱ_t) * ε) / √(ᾱ_t)
    fn predict_x0<B: Backend<FloatElem = f32>, const D: usize>(
        &self,
        xt: &GenericTensor<B, D>,
        pred_noise: &GenericTensor<B, D>,
        t: usize,
    ) -> GenericTensor<B, D> {
        let alpha_cumprod = self.alpha_cumprod(t);
        let sqrt_alpha_cumprod = alpha_cumprod.sqrt();
        let sqrt_one_minus_alpha_cumprod = (1.0 - alpha_cumprod).sqrt();

        // x_0 = (x_t - √(1 - ᾱ_t) * ε) / √(ᾱ_t)
        xt.sub(&pred_noise.mul_scalar(sqrt_one_minus_alpha_cumprod))
            .mul_scalar(1.0 / sqrt_alpha_cumprod)
    }
}

/// DDPM (Denoising Diffusion Probabilistic Models) scheduler.
///
/// Implements the forward and reverse diffusion process from
/// Ho et al., "Denoising Diffusion Probabilistic Models" (2020).
#[derive(Debug, Clone)]
pub struct DDPMScheduler {
    /// Configuration
    config: SchedulerConfig,
    /// Beta values for each timestep
    betas: Vec<f32>,
    /// Alpha values (1 - beta)
    alphas: Vec<f32>,
    /// Cumulative product of alphas
    alphas_cumprod: Vec<f32>,
    /// √(ᾱ_t)
    sqrt_alphas_cumprod: Vec<f32>,
    /// √(1 - ᾱ_t)
    sqrt_one_minus_alphas_cumprod: Vec<f32>,
    /// Posterior variance σ²_t
    posterior_variance: Vec<f32>,
}

impl DDPMScheduler {
    /// Create a new DDPM scheduler.
    pub fn new(num_timesteps: usize, beta_start: f32, beta_end: f32) -> Self {
        let config = SchedulerConfig {
            num_timesteps,
            beta_start,
            beta_end,
            schedule_type: BetaSchedule::Linear,
        };
        Self::from_config(config)
    }

    /// Create from configuration.
    pub fn from_config(config: SchedulerConfig) -> Self {
        let betas = Self::compute_betas(&config);
        let alphas: Vec<f32> = betas.iter().map(|b| 1.0 - b).collect();

        // Compute cumulative products
        let mut alphas_cumprod = Vec::with_capacity(config.num_timesteps);
        let mut cumprod = 1.0f32;
        for &alpha in &alphas {
            cumprod *= alpha;
            alphas_cumprod.push(cumprod);
        }

        // Precompute sqrt values
        let sqrt_alphas_cumprod: Vec<f32> = alphas_cumprod.iter().map(|a| a.sqrt()).collect();
        let sqrt_one_minus_alphas_cumprod: Vec<f32> =
            alphas_cumprod.iter().map(|a| (1.0 - a).sqrt()).collect();

        // Posterior variance: β̃_t = β_t * (1 - ᾱ_{t-1}) / (1 - ᾱ_t)
        let mut posterior_variance = Vec::with_capacity(config.num_timesteps);
        posterior_variance.push(betas[0]); // t=0
        for t in 1..config.num_timesteps {
            let var = betas[t] * (1.0 - alphas_cumprod[t - 1]) / (1.0 - alphas_cumprod[t]);
            posterior_variance.push(var.max(1e-20)); // Clip for numerical stability
        }

        Self {
            config,
            betas,
            alphas,
            alphas_cumprod,
            sqrt_alphas_cumprod,
            sqrt_one_minus_alphas_cumprod,
            posterior_variance,
        }
    }

    /// Compute beta schedule.
    fn compute_betas(config: &SchedulerConfig) -> Vec<f32> {
        let t = config.num_timesteps;

        match config.schedule_type {
            BetaSchedule::Linear => {
                // Linear interpolation from beta_start to beta_end
                (0..t)
                    .map(|i| {
                        config.beta_start
                            + (i as f32 / (t - 1) as f32) * (config.beta_end - config.beta_start)
                    })
                    .collect()
            }
            BetaSchedule::Cosine => {
                // Cosine schedule from "Improved DDPM"
                let s = 0.008f32; // Small offset to prevent β from being too small
                let max_beta = 0.999f32;

                let f = |t: f32| -> f32 {
                    let angle = (t + s) / (1.0 + s) * std::f32::consts::FRAC_PI_2;
                    angle.cos().powi(2)
                };

                let mut betas = Vec::with_capacity(t);
                for i in 0..t {
                    let t1 = i as f32 / t as f32;
                    let t2 = (i + 1) as f32 / t as f32;
                    let beta = 1.0 - f(t2) / f(t1);
                    betas.push(beta.min(max_beta));
                }
                betas
            }
            BetaSchedule::ScaledLinear => {
                // Scaled linear: β varies from sqrt(β_start) to sqrt(β_end), then squared
                let sqrt_start = config.beta_start.sqrt();
                let sqrt_end = config.beta_end.sqrt();
                (0..t)
                    .map(|i| {
                        let sqrt_beta =
                            sqrt_start + (i as f32 / (t - 1) as f32) * (sqrt_end - sqrt_start);
                        sqrt_beta * sqrt_beta
                    })
                    .collect()
            }
        }
    }

    /// Perform one reverse diffusion step (DDPM sampling).
    ///
    /// x_{t-1} = μ_θ(x_t, t) + σ_t * z, where z ~ N(0, 1)
    pub fn step<B: Backend<FloatElem = f32>, const D: usize>(
        &self,
        xt: &GenericTensor<B, D>,
        pred_noise: &GenericTensor<B, D>,
        t: usize,
        noise: Option<&GenericTensor<B, D>>,
    ) -> GenericTensor<B, D> {
        let alpha = self.alphas[t];
        let alpha_cumprod = self.alphas_cumprod[t];
        let beta = self.betas[t];

        // Predict x_0
        let pred_x0 = self.predict_x0(xt, pred_noise, t);

        // Compute mean for posterior: μ_θ(x_t, t) = √(ᾱ_{t-1}) * β_t / (1 - ᾱ_t) * x_0
        //                                        + √(α_t) * (1 - ᾱ_{t-1}) / (1 - ᾱ_t) * x_t
        let (coef_x0, coef_xt) = if t == 0 {
            (1.0, 0.0)
        } else {
            let alpha_cumprod_prev = self.alphas_cumprod[t - 1];
            let one_minus_alpha_cumprod = 1.0 - alpha_cumprod;
            let coef_x0 = alpha_cumprod_prev.sqrt() * beta / one_minus_alpha_cumprod;
            let coef_xt = alpha.sqrt() * (1.0 - alpha_cumprod_prev) / one_minus_alpha_cumprod;
            (coef_x0, coef_xt)
        };

        let mean = pred_x0.mul_scalar(coef_x0).add(&xt.mul_scalar(coef_xt));

        // Add noise if not at t=0
        if t > 0 {
            let std = self.posterior_variance[t].sqrt();
            if let Some(z) = noise {
                mean.add(&z.mul_scalar(std))
            } else {
                mean
            }
        } else {
            mean
        }
    }

    /// Get posterior variance at timestep t.
    pub fn posterior_variance(&self, t: usize) -> f32 {
        self.posterior_variance[t]
    }
}

impl NoiseScheduler for DDPMScheduler {
    fn num_timesteps(&self) -> usize {
        self.config.num_timesteps
    }

    fn alpha_cumprod(&self, t: usize) -> f32 {
        self.alphas_cumprod[t]
    }

    fn beta(&self, t: usize) -> f32 {
        self.betas[t]
    }
}

/// DDIM (Denoising Diffusion Implicit Models) scheduler.
///
/// Implements deterministic sampling from Song et al.,
/// "Denoising Diffusion Implicit Models" (2021).
#[derive(Debug, Clone)]
pub struct DDIMScheduler {
    /// Inner DDPM scheduler (shares noise schedule)
    inner: DDPMScheduler,
    /// DDIM eta parameter (0 = deterministic, 1 = DDPM)
    eta: f32,
    /// Timesteps to use for sampling (can be subset)
    timesteps: Vec<usize>,
}

impl DDIMScheduler {
    /// Create a new DDIM scheduler.
    ///
    /// # Arguments
    /// * `ddpm` - Base DDPM scheduler
    /// * `eta` - Stochasticity parameter (0 = deterministic, 1 = DDPM)
    /// * `num_inference_steps` - Number of steps for sampling (can be < training steps)
    pub fn new(ddpm: DDPMScheduler, eta: f32, num_inference_steps: usize) -> Self {
        let num_timesteps = ddpm.num_timesteps();

        // Create evenly spaced timesteps
        let step_ratio = num_timesteps / num_inference_steps;
        let timesteps: Vec<usize> = (0..num_inference_steps)
            .map(|i| (num_timesteps - 1 - i * step_ratio).min(num_timesteps - 1))
            .collect();

        Self {
            inner: ddpm,
            eta,
            timesteps,
        }
    }

    /// Get the timesteps for inference.
    pub fn timesteps(&self) -> &[usize] {
        &self.timesteps
    }

    /// Perform one DDIM sampling step.
    ///
    /// DDIM update:
    /// x_{t-1} = √(ᾱ_{t-1}) * x_0 + √(1 - ᾱ_{t-1} - σ²) * ε_θ + σ * z
    ///
    /// where σ² = η² * β̃_t (posterior variance)
    pub fn step<B: Backend<FloatElem = f32>, const D: usize>(
        &self,
        xt: &GenericTensor<B, D>,
        pred_noise: &GenericTensor<B, D>,
        t: usize,
        t_prev: usize,
        noise: Option<&GenericTensor<B, D>>,
    ) -> GenericTensor<B, D> {
        let alpha_cumprod = self.inner.alphas_cumprod[t];
        let alpha_cumprod_prev = if t_prev == 0 {
            1.0 // Convention: ᾱ_0 = 1
        } else {
            self.inner.alphas_cumprod[t_prev]
        };

        // Predict x_0
        let pred_x0 = self.inner.predict_x0(xt, pred_noise, t);

        // Compute σ for DDIM
        let sigma = if t_prev > 0 {
            let sigma_sq = self.eta * self.eta * self.inner.posterior_variance[t];
            sigma_sq.sqrt()
        } else {
            0.0
        };

        // Direction pointing to x_t
        let dir_xt = (1.0 - alpha_cumprod_prev - sigma * sigma).max(0.0).sqrt();

        // x_{t-1} = √(ᾱ_{t-1}) * x_0 + dir_xt * ε_θ + σ * z
        let x_prev = pred_x0
            .mul_scalar(alpha_cumprod_prev.sqrt())
            .add(&pred_noise.mul_scalar(dir_xt));

        if sigma > 0.0 {
            if let Some(z) = noise {
                x_prev.add(&z.mul_scalar(sigma))
            } else {
                x_prev
            }
        } else {
            x_prev
        }
    }
}

impl NoiseScheduler for DDIMScheduler {
    fn num_timesteps(&self) -> usize {
        self.inner.num_timesteps()
    }

    fn alpha_cumprod(&self, t: usize) -> f32 {
        self.inner.alpha_cumprod(t)
    }

    fn beta(&self, t: usize) -> f32 {
        self.inner.beta(t)
    }
}

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

    #[test]
    fn test_ddpm_scheduler_creation() {
        let scheduler = DDPMScheduler::new(1000, 1e-4, 0.02);

        assert_eq!(scheduler.num_timesteps(), 1000);
        assert!(scheduler.alpha_cumprod(0) > 0.99); // First alpha should be close to 1
        assert!(scheduler.alpha_cumprod(999) < 0.1); // Last alpha should be small
    }

    #[test]
    fn test_beta_schedule_monotonic() {
        let scheduler = DDPMScheduler::new(100, 1e-4, 0.02);

        // Betas should be monotonically increasing (for linear schedule)
        for t in 1..100 {
            assert!(
                scheduler.beta(t) >= scheduler.beta(t - 1),
                "Beta should be monotonic at t={}",
                t
            );
        }
    }

    #[test]
    fn test_alpha_cumprod_decreasing() {
        let scheduler = DDPMScheduler::new(100, 1e-4, 0.02);

        // Alpha cumprod should be monotonically decreasing
        for t in 1..100 {
            assert!(
                scheduler.alpha_cumprod(t) < scheduler.alpha_cumprod(t - 1),
                "Alpha cumprod should decrease at t={}",
                t
            );
        }
    }

    #[test]
    fn test_add_noise_and_predict_x0() {
        let scheduler = DDPMScheduler::new(100, 1e-4, 0.02);
        let device = CpuDevice::default();

        // Create test tensors
        let x0 = GenericTensor::<CpuBackend, 2>::from_slice(&[1.0, 2.0, 3.0, 4.0], [2, 2], &device);
        let noise =
            GenericTensor::<CpuBackend, 2>::from_slice(&[0.1, 0.2, 0.3, 0.4], [2, 2], &device);

        let t = 50;

        // Add noise
        let xt = scheduler.add_noise(&x0, &noise, t);

        // Predict x0 back
        let pred_x0 = scheduler.predict_x0(&xt, &noise, t);

        // Should recover x0 approximately (using known noise)
        let x0_vec = x0.to_vec();
        let pred_vec = pred_x0.to_vec();

        for (orig, pred) in x0_vec.iter().zip(pred_vec.iter()) {
            assert!(
                (orig - pred).abs() < 1e-5,
                "x0 recovery failed: {} vs {}",
                orig,
                pred
            );
        }
    }

    #[test]
    fn test_cosine_schedule() {
        let config = SchedulerConfig {
            num_timesteps: 100,
            beta_start: 1e-4,
            beta_end: 0.02,
            schedule_type: BetaSchedule::Cosine,
        };
        let scheduler = DDPMScheduler::from_config(config);

        // Cosine schedule should have smaller betas at start
        assert!(scheduler.beta(0) < 0.01);
        // And should never exceed max_beta
        for t in 0..100 {
            assert!(scheduler.beta(t) < 1.0);
        }
    }

    #[test]
    fn test_ddim_scheduler() {
        let ddpm = DDPMScheduler::new(1000, 1e-4, 0.02);
        let ddim = DDIMScheduler::new(ddpm, 0.0, 50); // 50 steps, deterministic

        assert_eq!(ddim.timesteps().len(), 50);
        // Timesteps should be evenly spaced from high to low
        assert!(ddim.timesteps()[0] > ddim.timesteps()[1]);
    }
}