//! Diffusion UNet architecture.
//!
//! This module implements a simplified UNet architecture suitable for
//! diffusion models in the context of PDE solving.

use rtx_backend::Backend;
use rtx_nn::{GenericLinear, GenericModule};
use rtx_tensor::GenericTensor;

/// Configuration for the diffusion UNet.
#[derive(Debug, Clone)]
pub struct UNetConfig {
    /// Number of input channels
    pub in_channels: usize,
    /// Number of output channels (usually same as input)
    pub out_channels: usize,
    /// Base channel dimension
    pub base_channels: usize,
    /// Hidden dimension multiplier
    pub hidden_mult: usize,
    /// Dimension of time embedding
    pub time_embed_dim: usize,
}

impl Default for UNetConfig {
    fn default() -> Self {
        Self {
            in_channels: 1,
            out_channels: 1,
            base_channels: 64,
            hidden_mult: 2,
            time_embed_dim: 128,
        }
    }
}

impl UNetConfig {
    /// Create config for small model (for testing/debugging).
    pub fn small() -> Self {
        Self {
            in_channels: 1,
            out_channels: 1,
            base_channels: 32,
            hidden_mult: 2,
            time_embed_dim: 64,
        }
    }
}

/// Sinusoidal time embedding module.
///
/// Maps scalar timestep to a fixed-dimension embedding vector using
/// sinusoidal positional encoding.
pub struct TimeEmbedding<B: Backend<FloatElem = f32>> {
    /// Linear projection from embedding to hidden
    fc1: GenericLinear<B>,
    /// Linear projection from hidden to output
    fc2: GenericLinear<B>,
    /// Embedding dimension
    embed_dim: usize,
}

impl<B: Backend<FloatElem = f32>> TimeEmbedding<B> {
    /// Create a new time embedding module.
    pub fn new(embed_dim: usize, hidden_dim: usize, device: &B::Device) -> Self {
        let fc1 = GenericLinear::new(embed_dim, hidden_dim, true, device);
        let fc2 = GenericLinear::new(hidden_dim, hidden_dim, true, device);

        Self {
            fc1,
            fc2,
            embed_dim,
        }
    }

    /// Get sinusoidal embedding for timestep.
    pub fn get_embedding(&self, t: usize, device: &B::Device) -> GenericTensor<B, 2> {
        let half_dim = self.embed_dim / 2;
        let emb_scale = -(10000.0f32.ln()) / (half_dim as f32 - 1.0);

        let mut embedding = vec![0.0f32; self.embed_dim];

        for i in 0..half_dim {
            let freq = (i as f32 * emb_scale).exp();
            let angle = t as f32 * freq;
            embedding[i] = angle.sin();
            embedding[i + half_dim] = angle.cos();
        }

        GenericTensor::<B, 2>::from_slice(&embedding, [1, self.embed_dim], device)
    }

    /// Forward pass: embed timestep and project through MLP.
    pub fn forward(&self, t: usize, device: &B::Device) -> GenericTensor<B, 2> {
        let emb = self.get_embedding(t, device);

        // MLP with SiLU activation
        let h = GenericModule::forward(&self.fc1, &emb);
        let h = silu(&h);
        GenericModule::forward(&self.fc2, &h)
    }
}

/// SiLU (Swish) activation: x * sigmoid(x)
fn silu<B: Backend<FloatElem = f32>, const D: usize>(
    x: &GenericTensor<B, D>,
) -> GenericTensor<B, D> {
    // x * sigmoid(x) = x / (1 + exp(-x))
    let neg_x = x.mul_scalar(-1.0);
    let exp_neg_x = neg_x.exp();
    let one_plus_exp = exp_neg_x.add_scalar(1.0);
    x.div(&one_plus_exp)
}

/// Simplified Diffusion UNet for 2D fields.
///
/// This is a simplified architecture that treats the 2D field as a flattened
/// vector and processes it through fully-connected layers with time conditioning.
///
/// Architecture:
/// - Input: [batch, channels, height, width] -> flatten to [batch, h*w*c]
/// - Encoder: fc -> hidden -> fc
/// - Time conditioning added at hidden layer
/// - Decoder: fc -> output
/// - Output: reshape back to [batch, channels, height, width]
pub struct DiffusionUNet<B: Backend<FloatElem = f32>> {
    /// Configuration
    config: UNetConfig,
    /// Time embedding module
    time_embedding: TimeEmbedding<B>,
    /// Input projection
    encoder1: GenericLinear<B>,
    /// Hidden layer
    encoder2: GenericLinear<B>,
    /// Time projection to hidden
    time_proj: GenericLinear<B>,
    /// Decoder layer 1
    decoder1: GenericLinear<B>,
    /// Output projection
    decoder2: GenericLinear<B>,
    /// Expected spatial size (height * width)
    spatial_size: usize,
}

impl<B: Backend<FloatElem = f32>> DiffusionUNet<B> {
    /// Create a new diffusion UNet.
    ///
    /// # Arguments
    /// * `config` - Model configuration
    /// * `spatial_size` - Expected height * width of input
    /// * `device` - Device to create on
    pub fn new(config: UNetConfig, spatial_size: usize, device: &B::Device) -> Self {
        let input_dim = config.in_channels * spatial_size;
        let hidden_dim = config.base_channels * config.hidden_mult;
        let output_dim = config.out_channels * spatial_size;

        // Time embedding
        let time_embedding =
            TimeEmbedding::new(config.time_embed_dim, config.time_embed_dim, device);

        // Encoder
        let encoder1 = GenericLinear::new(input_dim, config.base_channels, true, device);
        let encoder2 = GenericLinear::new(config.base_channels, hidden_dim, true, device);

        // Time projection to hidden dimension
        let time_proj = GenericLinear::new(config.time_embed_dim, hidden_dim, true, device);

        // Decoder
        let decoder1 = GenericLinear::new(hidden_dim, config.base_channels, true, device);
        let decoder2 = GenericLinear::new(config.base_channels, output_dim, true, device);

        Self {
            config,
            time_embedding,
            encoder1,
            encoder2,
            time_proj,
            decoder1,
            decoder2,
            spatial_size,
        }
    }

    /// Forward pass for noise prediction.
    ///
    /// # Arguments
    /// * `x` - Noisy input [batch, channels, height, width]
    /// * `t` - Timestep (scalar)
    /// * `device` - Device
    ///
    /// # Returns
    /// Predicted noise with same shape as input
    pub fn forward(
        &self,
        x: &GenericTensor<B, 4>,
        t: usize,
        device: &B::Device,
    ) -> GenericTensor<B, 4> {
        let shape = x.shape();
        let batch = shape[0];
        let channels = shape[1];
        let height = shape[2];
        let width = shape[3];

        // Flatten: [batch, c, h, w] -> [batch, c*h*w]
        let x_flat = x.reshape([batch, channels * height * width]);

        // Get time embedding [1, time_embed_dim]
        let time_emb = self.time_embedding.forward(t, device);

        // Expand time embedding to batch size [batch, time_embed_dim]
        let time_emb = expand_batch(&time_emb, batch);

        // Encoder
        let h = GenericModule::forward(&self.encoder1, &x_flat);
        let h = silu(&h);
        let h = GenericModule::forward(&self.encoder2, &h);
        let h = silu(&h);

        // Add time conditioning
        let time_proj = GenericModule::forward(&self.time_proj, &time_emb);
        let time_proj = silu(&time_proj);
        let h = h.add(&time_proj);

        // Decoder
        let h = GenericModule::forward(&self.decoder1, &h);
        let h = silu(&h);
        let output = GenericModule::forward(&self.decoder2, &h);

        // Reshape back: [batch, c*h*w] -> [batch, c, h, w]
        output.reshape([batch, channels, height, width])
    }
}

/// Expand a [1, dim] tensor to [batch, dim] by repeating.
fn expand_batch<B: Backend<FloatElem = f32>>(
    x: &GenericTensor<B, 2>,
    batch: usize,
) -> GenericTensor<B, 2> {
    let shape = x.shape();
    let dim = shape[1];

    let x_vec = x.to_vec();
    let mut expanded = Vec::with_capacity(batch * dim);
    for _ in 0..batch {
        expanded.extend_from_slice(&x_vec);
    }

    GenericTensor::<B, 2>::from_slice(&expanded, [batch, dim], &x.device())
}

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

    #[test]
    fn test_time_embedding() {
        let device = CpuDevice::default();
        let time_emb = TimeEmbedding::<CpuBackend>::new(64, 128, &device);

        let emb = time_emb.forward(100, &device);
        let shape = emb.shape();

        assert_eq!(shape[0], 1);
        assert_eq!(shape[1], 128);
    }

    #[test]
    fn test_silu_activation() {
        let device = CpuDevice::default();
        let x = GenericTensor::<CpuBackend, 2>::from_slice(&[0.0, 1.0, -1.0, 2.0], [2, 2], &device);

        let y = silu(&x);
        let y_vec = y.to_vec();

        // silu(0) = 0
        assert!((y_vec[0] - 0.0).abs() < 1e-5);
        // silu(1) ≈ 0.731
        assert!((y_vec[1] - 0.731).abs() < 0.01);
        // silu(-1) ≈ -0.269
        assert!((y_vec[2] + 0.269).abs() < 0.01);
    }

    #[test]
    fn test_unet_creation() {
        let device = CpuDevice::default();
        let config = UNetConfig::small();
        let spatial_size = 64; // 8x8
        let _unet = DiffusionUNet::<CpuBackend>::new(config, spatial_size, &device);
    }

    #[test]
    fn test_unet_forward() {
        let device = CpuDevice::default();
        let config = UNetConfig {
            in_channels: 1,
            out_channels: 1,
            base_channels: 16,
            hidden_mult: 2,
            time_embed_dim: 32,
        };

        let height = 8;
        let width = 8;
        let spatial_size = height * width;

        let unet = DiffusionUNet::<CpuBackend>::new(config, spatial_size, &device);

        // Create test input [1, 1, 8, 8]
        let input = GenericTensor::<CpuBackend, 4>::rand([1, 1, height, width], &device);

        let output = unet.forward(&input, 50, &device);
        let output_shape = output.shape();

        assert_eq!(output_shape[0], 1); // batch
        assert_eq!(output_shape[1], 1); // channels
        assert_eq!(output_shape[2], height); // height
        assert_eq!(output_shape[3], width); // width
    }

    #[test]
    fn test_expand_batch() {
        let device = CpuDevice::default();
        let x = GenericTensor::<CpuBackend, 2>::from_slice(&[1.0, 2.0, 3.0], [1, 3], &device);

        let expanded = expand_batch(&x, 4);
        let shape = expanded.shape();

        assert_eq!(shape[0], 4);
        assert_eq!(shape[1], 3);

        let data = expanded.to_vec();
        // All rows should have same values
        for i in 0..4 {
            assert_eq!(data[i * 3], 1.0);
            assert_eq!(data[i * 3 + 1], 2.0);
            assert_eq!(data[i * 3 + 2], 3.0);
        }
    }
}