//! Unit tests for TFT LSTM Encoder
//!
//! Tests LSTM encoder initialization, forward pass, and gradient computation.

use anyhow::Result;
use candle_core::{Device, Tensor};
use ml::tft::lstm_encoder::LSTMEncoder;

#[test]
fn test_lstm_encoder_creation() -> Result<()> {
    let device = Device::Cpu;
    let input_size = 64;
    let hidden_size = 128;
    let num_layers = 2;

    let encoder = LSTMEncoder::new(input_size, hidden_size, num_layers, &device)?;

    // Verify encoder was created successfully
    assert_eq!(encoder.input_size(), input_size);
    assert_eq!(encoder.hidden_size(), hidden_size);
    assert_eq!(encoder.num_layers(), num_layers);

    Ok(())
}

#[test]
fn test_lstm_encoder_forward_pass() -> Result<()> {
    let device = Device::Cpu;
    let batch_size = 4;
    let seq_len = 10;
    let input_size = 64;
    let hidden_size = 128;
    let num_layers = 2;

    let encoder = LSTMEncoder::new(input_size, hidden_size, num_layers, &device)?;

    // Create input tensor [batch, seq_len, input_size]
    let input = Tensor::randn(0.0f32, 1.0, (batch_size, seq_len, input_size), &device)?;

    // Forward pass
    let output = encoder.forward(&input)?;

    // Verify output shape [batch, seq_len, hidden_size]
    let dims = output.dims();
    assert_eq!(dims.len(), 3);
    assert_eq!(dims[0], batch_size);
    assert_eq!(dims[1], seq_len);
    assert_eq!(dims[2], hidden_size);

    Ok(())
}

#[test]
fn test_lstm_encoder_hidden_state() -> Result<()> {
    let device = Device::Cpu;
    let batch_size = 2;
    let seq_len = 5;
    let input_size = 32;
    let hidden_size = 64;
    let num_layers = 1;

    let encoder = LSTMEncoder::new(input_size, hidden_size, num_layers, &device)?;
    let input = Tensor::randn(0.0f32, 1.0, (batch_size, seq_len, input_size), &device)?;

    let output = encoder.forward(&input)?;

    // Verify output has values (not all zeros)
    let output_sum = output.sum_all()?.to_scalar::<f32>()?;
    assert!(
        output_sum.abs() > 0.001,
        "Output should have non-zero values"
    );

    Ok(())
}

#[test]
fn test_lstm_encoder_batch_independence() -> Result<()> {
    let device = Device::Cpu;
    let seq_len = 8;
    let input_size = 48;
    let hidden_size = 96;
    let num_layers = 2;

    let encoder = LSTMEncoder::new(input_size, hidden_size, num_layers, &device)?;

    // Process batch_size=1
    let input1 = Tensor::randn(0.0f32, 1.0, (1, seq_len, input_size), &device)?;
    let output1 = encoder.forward(&input1)?;

    // Process batch_size=4 with same sequence
    let input4 = input1.repeat(&[4, 1, 1])?;
    let output4 = encoder.forward(&input4)?;

    // Verify shapes
    assert_eq!(output1.dim(0)?, 1);
    assert_eq!(output4.dim(0)?, 4);

    Ok(())
}

#[test]
fn test_lstm_encoder_sequence_length_invariance() -> Result<()> {
    let device = Device::Cpu;
    let batch_size = 2;
    let input_size = 64;
    let hidden_size = 128;
    let num_layers = 2;

    let encoder = LSTMEncoder::new(input_size, hidden_size, num_layers, &device)?;

    // Test different sequence lengths
    for seq_len in [5, 10, 20] {
        let input = Tensor::randn(0.0f32, 1.0, (batch_size, seq_len, input_size), &device)?;
        let output = encoder.forward(&input)?;

        assert_eq!(output.dim(0)?, batch_size);
        assert_eq!(output.dim(1)?, seq_len);
        assert_eq!(output.dim(2)?, hidden_size);
    }

    Ok(())
}

#[test]
fn test_lstm_encoder_deterministic() -> Result<()> {
    let device = Device::Cpu;
    let batch_size = 2;
    let seq_len = 8;
    let input_size = 32;
    let hidden_size = 64;
    let num_layers = 1;

    let encoder = LSTMEncoder::new(input_size, hidden_size, num_layers, &device)?;

    // Fixed input
    let input = Tensor::ones(
        (batch_size, seq_len, input_size),
        candle_core::DType::F32,
        &device,
    )?;

    // Two forward passes with same input
    let output1 = encoder.forward(&input)?;
    let output2 = encoder.forward(&input)?;

    // Outputs should be identical (deterministic)
    let diff = (&output1 - &output2)?
        .abs()?
        .sum_all()?
        .to_scalar::<f32>()?;
    assert!(diff < 1e-6, "Forward pass should be deterministic");

    Ok(())
}

#[test]
fn test_lstm_encoder_gradient_flow() -> Result<()> {
    let device = Device::Cpu;
    let batch_size = 2;
    let seq_len = 5;
    let input_size = 32;
    let hidden_size = 64;
    let num_layers = 1;

    let encoder = LSTMEncoder::new(input_size, hidden_size, num_layers, &device)?;
    let input = Tensor::randn(0.0f32, 1.0, (batch_size, seq_len, input_size), &device)?;

    let output = encoder.forward(&input)?;

    // Compute loss (sum for gradient check)
    let loss = output.sum_all()?;

    // Verify gradient can be computed
    loss.backward()?;

    // This test passes if backward() doesn't panic
    Ok(())
}

#[test]
fn test_lstm_encoder_multi_layer() -> Result<()> {
    let device = Device::Cpu;
    let batch_size = 2;
    let seq_len = 10;
    let input_size = 64;
    let hidden_size = 128;

    // Test different layer counts
    for num_layers in [1, 2, 3, 4] {
        let encoder = LSTMEncoder::new(input_size, hidden_size, num_layers, &device)?;
        let input = Tensor::randn(0.0f32, 1.0, (batch_size, seq_len, input_size), &device)?;
        let output = encoder.forward(&input)?;

        assert_eq!(output.dim(2)?, hidden_size);
    }

    Ok(())
}