//! Tests for Classifier-Free Guidance in diffusion models
//!
//! TDD: Tests define behavior for CFG implementation

use rtx_diffuse::{CFGConfig, ClassifierFreeGuidance, DiffusionScheduler, GuidanceScale, Result};
use rtx_tensor::{Device, Tensor};

// Helper function to replace approx::assert_abs_diff_eq
fn assert_abs_diff_eq(a: f32, b: f32, epsilon: f32) {
    assert!(
        (a - b).abs() < epsilon,
        "Expected |{} - {}| < {}, but got {}",
        a,
        b,
        epsilon,
        (a - b).abs()
    );
}

#[test]
fn test_cfg_creation() {
    let device = Device::cpu();

    // Create CFG with standard configuration
    let config = CFGConfig {
        guidance_scale: 7.5,
        unconditional_guidance_scale: 1.0,
        guidance_rescale: 0.0,
        dynamic_thresholding: false,
        thresholding_percentile: 0.995,
    };

    let cfg = ClassifierFreeGuidance::new(config, &device);
    assert!(cfg.is_ok());

    let cfg = cfg.unwrap();
    assert_eq!(cfg.guidance_scale(), 7.5);
    assert!(!cfg.uses_dynamic_thresholding());
}

#[test]
fn test_cfg_conditional_unconditional_mixing() {
    let device = Device::cpu();

    let config = CFGConfig {
        guidance_scale: 7.5,
        unconditional_guidance_scale: 1.0,
        guidance_rescale: 0.0,
        dynamic_thresholding: false,
        thresholding_percentile: 0.995,
    };

    let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();

    // Create conditional and unconditional predictions
    let shape = vec![1, 3, 64, 64];
    let conditional = Tensor::ones(&shape, &device).unwrap();
    let unconditional = Tensor::zeros(&shape, &device).unwrap();

    // Apply CFG: output = unconditional + scale * (conditional - unconditional)
    let guided = cfg.apply(&conditional, &unconditional).unwrap();

    // Check the formula: guided = 0 + 7.5 * (1 - 0) = 7.5
    let guided_data = guided.to_vec().unwrap();
    assert_abs_diff_eq(guided_data[0], 7.5, 1e-5);
}

#[test]
fn test_cfg_with_batch() {
    let device = Device::cpu();

    let config = CFGConfig {
        guidance_scale: 5.0,
        unconditional_guidance_scale: 1.0,
        guidance_rescale: 0.0,
        dynamic_thresholding: false,
        thresholding_percentile: 0.995,
    };

    let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();

    // Batch of 4 samples
    let batch_size = 4;
    let shape = vec![batch_size, 3, 32, 32];

    let conditional = Tensor::randn(&shape, &device).unwrap();
    let unconditional = Tensor::randn(&shape, &device).unwrap();

    let guided = cfg.apply(&conditional, &unconditional).unwrap();

    // Output should have same shape
    assert_eq!(guided.shape().dims(), shape);
}

#[test]
fn test_cfg_guidance_rescale() {
    let device = Device::cpu();

    // Test guidance rescaling to prevent oversaturation
    let config = CFGConfig {
        guidance_scale: 10.0,
        unconditional_guidance_scale: 1.0,
        guidance_rescale: 0.7, // Rescale factor
        dynamic_thresholding: false,
        thresholding_percentile: 0.995,
    };

    let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();

    let shape = vec![1, 3, 64, 64];
    let conditional = Tensor::ones(&shape, &device)
        .unwrap()
        .mul_scalar(2.0)
        .unwrap();
    let unconditional = Tensor::ones(&shape, &device).unwrap();

    let guided = cfg.apply(&conditional, &unconditional).unwrap();

    // With rescaling, the values should be moderated
    let guided_data = guided.to_vec().unwrap();

    // Without rescale: 1 + 10 * (2 - 1) = 11
    // With rescale, it should be less
    assert!(
        guided_data[0] < 11.0,
        "Rescaling should reduce guidance magnitude"
    );
}

#[test]
fn test_cfg_dynamic_thresholding() {
    let device = Device::cpu();

    let config = CFGConfig {
        guidance_scale: 15.0, // High guidance
        unconditional_guidance_scale: 1.0,
        guidance_rescale: 0.0,
        dynamic_thresholding: true,
        thresholding_percentile: 0.995,
    };

    let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();

    // Create predictions that would overflow without thresholding
    let shape = vec![1, 3, 32, 32];
    let mut cond_data = vec![0.0f32; 3 * 32 * 32];
    let mut uncond_data = vec![0.0f32; 3 * 32 * 32];

    // Set some extreme values
    for i in 0..100 {
        cond_data[i] = 5.0;
        uncond_data[i] = -1.0;
    }

    let conditional = Tensor::from_data(cond_data, shape.clone(), &device).unwrap();
    let unconditional = Tensor::from_data(uncond_data, shape, &device).unwrap();

    let guided = cfg.apply(&conditional, &unconditional).unwrap();
    let guided_data = guided.to_vec().unwrap();

    // Dynamic thresholding should prevent extreme values
    let max_val = guided_data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
    let min_val = guided_data.iter().fold(f32::INFINITY, |a, &b| a.min(b));

    // Values should be clamped to reasonable range
    assert!(
        max_val < 100.0,
        "Dynamic thresholding should prevent extreme positive values"
    );
    assert!(
        min_val > -100.0,
        "Dynamic thresholding should prevent extreme negative values"
    );
}

#[test]
fn test_cfg_variable_guidance_scale() {
    let device = Device::cpu();

    // Test with time-varying guidance scale
    let cfg = ClassifierFreeGuidance::with_schedule(
        GuidanceScale::Linear {
            start: 1.0,
            end: 10.0,
        },
        &device,
    )
    .unwrap();

    let shape = vec![1, 3, 32, 32];
    let conditional = Tensor::ones(&shape, &device).unwrap();
    let unconditional = Tensor::zeros(&shape, &device).unwrap();

    // Early timestep (t=0.0) - should use start scale
    let guided_early = cfg
        .apply_with_timestep(&conditional, &unconditional, 0.0)
        .unwrap();
    let early_data = guided_early.to_vec().unwrap();
    assert_abs_diff_eq(early_data[0], 1.0, 1e-5); // scale = 1.0

    // Late timestep (t=1.0) - should use end scale
    let guided_late = cfg
        .apply_with_timestep(&conditional, &unconditional, 1.0)
        .unwrap();
    let late_data = guided_late.to_vec().unwrap();
    assert_abs_diff_eq(late_data[0], 10.0, 1e-5); // scale = 10.0
}

#[test]
fn test_cfg_with_multiple_conditions() {
    let device = Device::cpu();

    // Test CFG with multiple conditioning signals (e.g., text + image)
    let config = CFGConfig {
        guidance_scale: 7.5,
        unconditional_guidance_scale: 1.0,
        guidance_rescale: 0.0,
        dynamic_thresholding: false,
        thresholding_percentile: 0.995,
    };

    let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();

    let shape = vec![1, 3, 32, 32];

    // Multiple conditional predictions (e.g., from different prompts)
    let cond1 = Tensor::ones(&shape, &device).unwrap();
    let cond2 = Tensor::ones(&shape, &device)
        .unwrap()
        .mul_scalar(0.5)
        .unwrap();
    let unconditional = Tensor::zeros(&shape, &device).unwrap();

    // Combine conditions (average in this case)
    let combined = cfg
        .apply_multi(&[&cond1, &cond2], &unconditional, &[0.7, 0.3])
        .unwrap();

    // Check weighted combination
    let combined_data = combined.to_vec().unwrap();
    // Expected: 0 + 7.5 * ((0.7 * 1 + 0.3 * 0.5) - 0) = 7.5 * 0.85 = 6.375
    assert_abs_diff_eq(combined_data[0], 6.375, 1e-4);
}

#[test]
fn test_cfg_negative_prompting() {
    let device = Device::cpu();

    let config = CFGConfig {
        guidance_scale: 7.5,
        unconditional_guidance_scale: 1.0,
        guidance_rescale: 0.0,
        dynamic_thresholding: false,
        thresholding_percentile: 0.995,
    };

    let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();

    let shape = vec![1, 3, 32, 32];

    // Positive condition
    let positive = Tensor::ones(&shape, &device).unwrap();
    // Negative condition (what to avoid)
    let negative = Tensor::ones(&shape, &device)
        .unwrap()
        .mul_scalar(-0.5)
        .unwrap();

    // Apply with negative prompting
    let guided = cfg.apply_with_negative(&positive, &negative).unwrap();

    // Should push away from negative
    let guided_data = guided.to_vec().unwrap();
    assert!(
        guided_data[0] > 1.0,
        "Should amplify away from negative prompt"
    );
}

#[test]
fn test_cfg_gradient_scaling() {
    let device = Device::cpu();

    // Test gradient scaling for training stability
    let config = CFGConfig {
        guidance_scale: 7.5,
        unconditional_guidance_scale: 1.0,
        guidance_rescale: 0.0,
        dynamic_thresholding: false,
        thresholding_percentile: 0.995,
    };

    let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();

    // Create tensors that require gradients
    let shape = vec![2, 3, 32, 32];
    let mut conditional = Tensor::randn(&shape, &device).unwrap();
    conditional.set_requires_grad(true);
    let mut unconditional = Tensor::randn(&shape, &device).unwrap();
    unconditional.set_requires_grad(true);

    let guided = cfg.apply(&conditional, &unconditional).unwrap();

    // Check that gradients can flow
    assert!(guided.requires_grad());
}

#[test]
fn test_cfg_per_channel_guidance() {
    let device = Device::cpu();

    // Test different guidance scales per channel (RGB)
    let cfg = ClassifierFreeGuidance::with_per_channel_scales(
        vec![5.0, 7.5, 10.0], // Different scales for R, G, B
        &device,
    )
    .unwrap();

    let shape = vec![1, 3, 32, 32];
    let conditional = Tensor::ones(&shape, &device).unwrap();
    let unconditional = Tensor::zeros(&shape, &device).unwrap();

    let guided = cfg.apply(&conditional, &unconditional).unwrap();
    let guided_data = guided.to_vec().unwrap();

    // Check different scales applied to different channels
    let r_val = guided_data[0]; // Red channel, first pixel
    let g_val = guided_data[32 * 32]; // Green channel, first pixel
    let b_val = guided_data[2 * 32 * 32]; // Blue channel, first pixel

    assert_abs_diff_eq(r_val, 5.0, 1e-5);
    assert_abs_diff_eq(g_val, 7.5, 1e-5);
    assert_abs_diff_eq(b_val, 10.0, 1e-5);
}