//! Tacotron 2 acoustic model
//!
//! Tacotron 2 is an autoregressive sequence-to-sequence model that generates
//! mel spectrograms from character/phoneme sequences using attention.

use crate::acoustic::{AcousticModel, MelSpectrogramConfig};
use crate::Result;
use rtx_nn::layers::{Module, linear::Linear, conv::{Conv1d, Conv1dConfig}};
use rtx_nn::layers::activation::{ReLU, Tanh, Sigmoid};
use rtx_nn::layers::dropout::Dropout;
use rtx_nn::layers::embedding::{Embedding, EmbeddingConfig};
use rtx_tensor::{Tensor, Device};
use serde::{Deserialize, Serialize};

/// Configuration for Tacotron2
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Tacotron2Config {
    /// Encoder embedding dimension
    pub encoder_dim: usize,

    /// Decoder LSTM dimension
    pub decoder_dim: usize,

    /// Attention dimension
    pub attention_dim: usize,

    /// Prenet dimension
    pub prenet_dim: usize,

    /// Postnet convolutional channels
    pub postnet_channels: usize,

    /// Mel spectrogram dimension
    pub mel_dim: usize,

    /// Maximum decoder steps (for inference)
    pub max_decoder_steps: usize,

    /// Character/phoneme vocabulary size
    pub vocab_size: usize,

    /// Dropout probability
    pub dropout: f32,

    /// Number of encoder conv layers
    pub encoder_n_convs: usize,

    /// Number of postnet conv layers
    pub postnet_n_convs: usize,

    /// Mel spectrogram configuration
    pub mel_config: MelSpectrogramConfig,
}

impl Default for Tacotron2Config {
    fn default() -> Self {
        Self {
            encoder_dim: 512,
            decoder_dim: 1024,
            attention_dim: 128,
            prenet_dim: 256,
            postnet_channels: 512,
            mel_dim: 80,
            max_decoder_steps: 1000,
            vocab_size: 100,
            dropout: 0.5,
            encoder_n_convs: 3,
            postnet_n_convs: 5,
            mel_config: MelSpectrogramConfig::default(),
        }
    }
}

impl Tacotron2Config {
    /// Validate configuration
    pub fn validate(&self) -> Result<()> {
        if self.encoder_dim == 0 {
            return Err(crate::TtsError::InvalidConfig("encoder_dim must be > 0".into()));
        }
        if self.decoder_dim == 0 {
            return Err(crate::TtsError::InvalidConfig("decoder_dim must be > 0".into()));
        }
        if self.mel_dim == 0 {
            return Err(crate::TtsError::InvalidConfig("mel_dim must be > 0".into()));
        }
        self.mel_config.validate()?;
        Ok(())
    }
}

/// Prenet for mel spectrogram processing
#[derive(Debug)]
struct PreNet {
    fc1: Linear,
    fc2: Linear,
    dropout: Dropout,
    activation: ReLU,
    training: bool,
}

impl PreNet {
    fn new(in_dim: usize, prenet_dim: usize, dropout: f32, device: &Device) -> Result<Self> {
        let fc1 = Linear::new(in_dim, prenet_dim, true, device)
            .map_err(|e| crate::TtsError::ModelError(format!("FC1 creation failed: {e}")))?;
        let fc2 = Linear::new(prenet_dim, prenet_dim, true, device)
            .map_err(|e| crate::TtsError::ModelError(format!("FC2 creation failed: {e}")))?;

        Ok(Self {
            fc1,
            fc2,
            dropout: Dropout::new(dropout),
            activation: ReLU::new(),
            training: true,
        })
    }

    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let h1 = self.fc1.forward(x)
            .map_err(|e| crate::TtsError::ModelError(format!("FC1 forward failed: {e}")))?;
        let h1 = self.activation.forward(&h1)
            .map_err(|e| crate::TtsError::ModelError(format!("ReLU failed: {e}")))?;
        let h1 = self.dropout.forward(&h1)
            .map_err(|e| crate::TtsError::ModelError(format!("Dropout failed: {e}")))?;

        let h2 = self.fc2.forward(&h1)
            .map_err(|e| crate::TtsError::ModelError(format!("FC2 forward failed: {e}")))?;
        let h2 = self.activation.forward(&h2)
            .map_err(|e| crate::TtsError::ModelError(format!("ReLU failed: {e}")))?;

        self.dropout.forward(&h2)
            .map_err(|e| crate::TtsError::ModelError(format!("Dropout failed: {e}")))
    }

    fn set_training(&mut self, training: bool) {
        self.training = training;
    }
}

/// Postnet for mel spectrogram refinement
#[derive(Debug)]
struct PostNet {
    conv_layers: Vec<Conv1d>,
    batch_norm_layers: Vec<Linear>, // Simplified batch norm as linear
    dropout: Dropout,
    tanh: Tanh,
    training: bool,
}

impl PostNet {
    fn new(config: &Tacotron2Config, device: &Device) -> Result<Self> {
        let mut conv_layers = Vec::new();
        let mut batch_norm_layers = Vec::new();

        for i in 0..config.postnet_n_convs {
            let in_channels = if i == 0 { config.mel_dim } else { config.postnet_channels };
            let out_channels = if i == config.postnet_n_convs - 1 {
                config.mel_dim
            } else {
                config.postnet_channels
            };

            let conv_config = Conv1dConfig {
                in_channels,
                out_channels,
                kernel_size: 5,
                stride: 1,
                padding: 2,
                dilation: 1,
                groups: 1,
                bias: true,
            };

            conv_layers.push(Conv1d::with_config(conv_config, device)
                .map_err(|e| crate::TtsError::ModelError(format!("Conv1d creation failed: {e}")))?);

            // Simplified batch norm
            batch_norm_layers.push(Linear::new(out_channels, out_channels, true, device)
                .map_err(|e| crate::TtsError::ModelError(format!("Batch norm creation failed: {e}")))?);
        }

        Ok(Self {
            conv_layers,
            batch_norm_layers,
            dropout: Dropout::new(0.5),
            tanh: Tanh::new(),
            training: true,
        })
    }

    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let mut hidden = x.clone();

        for (i, (conv, bn)) in self.conv_layers.iter().zip(self.batch_norm_layers.iter()).enumerate() {
            let conv_out = conv.forward(&hidden)
                .map_err(|e| crate::TtsError::ModelError(format!("Conv forward failed: {e}")))?;

            // Transpose for batch norm
            let transposed = conv_out.transpose(1, 2)
                .map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;

            let bn_out = bn.forward(&transposed)
                .map_err(|e| crate::TtsError::ModelError(format!("Batch norm failed: {e}")))?;

            // Transpose back
            hidden = bn_out.transpose(1, 2)
                .map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;

            // Tanh activation except for last layer
            if i < self.conv_layers.len() - 1 {
                hidden = self.tanh.forward(&hidden)
                    .map_err(|e| crate::TtsError::ModelError(format!("Tanh failed: {e}")))?;
            }

            if self.training {
                hidden = self.dropout.forward(&hidden)
                    .map_err(|e| crate::TtsError::ModelError(format!("Dropout failed: {e}")))?;
            }
        }

        Ok(hidden)
    }

    fn set_training(&mut self, training: bool) {
        self.training = training;
        for conv in &mut self.conv_layers {
            conv.train(training);
        }
    }
}

/// Location-sensitive attention
#[derive(Debug)]
struct LocationAttention {
    query_layer: Linear,
    memory_layer: Linear,
    location_conv: Conv1d,
    location_layer: Linear,
    v: Linear,
    attention_dim: usize,
    device: Device,
}

impl LocationAttention {
    fn new(
        query_dim: usize,
        memory_dim: usize,
        attention_dim: usize,
        device: &Device,
    ) -> Result<Self> {
        let query_layer = Linear::new(query_dim, attention_dim, false, device)
            .map_err(|e| crate::TtsError::ModelError(format!("Query layer creation failed: {e}")))?;

        let memory_layer = Linear::new(memory_dim, attention_dim, false, device)
            .map_err(|e| crate::TtsError::ModelError(format!("Memory layer creation failed: {e}")))?;

        let location_conv_config = Conv1dConfig {
            in_channels: 2,
            out_channels: 32,
            kernel_size: 31,
            stride: 1,
            padding: 15,
            dilation: 1,
            groups: 1,
            bias: true,
        };

        let location_conv = Conv1d::with_config(location_conv_config, device)
            .map_err(|e| crate::TtsError::ModelError(format!("Location conv creation failed: {e}")))?;

        let location_layer = Linear::new(32, attention_dim, false, device)
            .map_err(|e| crate::TtsError::ModelError(format!("Location layer creation failed: {e}")))?;

        let v = Linear::new(attention_dim, 1, false, device)
            .map_err(|e| crate::TtsError::ModelError(format!("V layer creation failed: {e}")))?;

        Ok(Self {
            query_layer,
            memory_layer,
            location_conv,
            location_layer,
            v,
            attention_dim,
            device: device.clone(),
        })
    }

    fn forward(
        &self,
        query: &Tensor,
        memory: &Tensor,
        attention_weights_cat: &Tensor,
    ) -> Result<Tensor> {
        let processed_query = self.query_layer.forward(query)
            .map_err(|e| crate::TtsError::ModelError(format!("Query layer failed: {e}")))?;

        let processed_memory = self.memory_layer.forward(memory)
            .map_err(|e| crate::TtsError::ModelError(format!("Memory layer failed: {e}")))?;

        // Process location features
        let processed_location = self.location_conv.forward(attention_weights_cat)
            .map_err(|e| crate::TtsError::ModelError(format!("Location conv failed: {e}")))?;

        let processed_location = processed_location.transpose(1, 2)
            .map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;

        let processed_location = self.location_layer.forward(&processed_location)
            .map_err(|e| crate::TtsError::ModelError(format!("Location layer failed: {e}")))?;

        // Expand query to match memory sequence length
        let query_expanded = processed_query.unsqueeze(1)
            .map_err(|e| crate::TtsError::TensorError(format!("Unsqueeze failed: {e}")))?;

        // Compute alignment energies
        let energies = (&query_expanded + &processed_memory)
            .map_err(|e| crate::TtsError::TensorError(format!("Add failed: {e}")))?;
        let energies = (&energies + &processed_location)
            .map_err(|e| crate::TtsError::TensorError(format!("Add failed: {e}")))?;

        let energies = rtx_tensor::ops::tanh(&energies)
            .map_err(|e| crate::TtsError::TensorError(format!("Tanh failed: {e}")))?;

        let alignment = self.v.forward(&energies)
            .map_err(|e| crate::TtsError::ModelError(format!("V layer failed: {e}")))?;

        let alignment = alignment.squeeze(2)
            .map_err(|e| crate::TtsError::TensorError(format!("Squeeze failed: {e}")))?;

        // Softmax
        rtx_tensor::ops::softmax(&alignment, 1)
            .map_err(|e| crate::TtsError::TensorError(format!("Softmax failed: {e}")))
    }
}

/// Encoder for Tacotron2
#[derive(Debug)]
struct Encoder {
    embedding: Embedding,
    conv_layers: Vec<Conv1d>,
    dropout: Dropout,
    device: Device,
    training: bool,
}

impl Encoder {
    fn new(config: &Tacotron2Config, device: &Device) -> Result<Self> {
        let emb_config = EmbeddingConfig::new(config.vocab_size, config.encoder_dim);
        let embedding = Embedding::new(emb_config, device)
            .map_err(|e| crate::TtsError::ModelError(format!("Embedding creation failed: {e}")))?;

        let mut conv_layers = Vec::new();
        for _ in 0..config.encoder_n_convs {
            let conv_config = Conv1dConfig {
                in_channels: config.encoder_dim,
                out_channels: config.encoder_dim,
                kernel_size: 5,
                stride: 1,
                padding: 2,
                dilation: 1,
                groups: 1,
                bias: true,
            };

            conv_layers.push(Conv1d::with_config(conv_config, device)
                .map_err(|e| crate::TtsError::ModelError(format!("Conv1d creation failed: {e}")))?);
        }

        Ok(Self {
            embedding,
            conv_layers,
            dropout: Dropout::new(config.dropout),
            device: device.clone(),
            training: true,
        })
    }

    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let embedded = self.embedding.forward(x)
            .map_err(|e| crate::TtsError::ModelError(format!("Embedding failed: {e}")))?;

        // Transpose for conv: [batch, seq, dim] -> [batch, dim, seq]
        let mut hidden = embedded.transpose(1, 2)
            .map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;

        for conv in &self.conv_layers {
            let conv_out = conv.forward(&hidden)
                .map_err(|e| crate::TtsError::ModelError(format!("Conv forward failed: {e}")))?;

            hidden = rtx_tensor::ops::relu(&conv_out)
                .map_err(|e| crate::TtsError::TensorError(format!("ReLU failed: {e}")))?;

            if self.training {
                hidden = self.dropout.forward(&hidden)
                    .map_err(|e| crate::TtsError::ModelError(format!("Dropout failed: {e}")))?;
            }
        }

        // Transpose back: [batch, dim, seq] -> [batch, seq, dim]
        hidden.transpose(1, 2)
            .map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))
    }

    fn set_training(&mut self, training: bool) {
        self.training = training;
        for conv in &mut self.conv_layers {
            conv.train(training);
        }
    }
}

/// Tacotron 2 acoustic model
#[derive(Debug)]
pub struct Tacotron2 {
    config: Tacotron2Config,
    encoder: Encoder,
    prenet: PreNet,
    attention: LocationAttention,
    decoder_rnn_input: Linear,
    decoder_rnn_hidden: Linear,
    mel_projection: Linear,
    gate_projection: Linear,
    postnet: PostNet,
    device: Device,
    training: bool,
}

impl Tacotron2 {
    /// Create a new Tacotron2 model
    pub fn new(config: Tacotron2Config, device: &Device) -> Result<Self> {
        config.validate()?;

        let encoder = Encoder::new(&config, device)?;
        let prenet = PreNet::new(config.mel_dim, config.prenet_dim, config.dropout, device)?;

        let attention = LocationAttention::new(
            config.decoder_dim,
            config.encoder_dim,
            config.attention_dim,
            device,
        )?;

        let decoder_input_dim = config.prenet_dim + config.encoder_dim;
        let decoder_rnn_input = Linear::new(decoder_input_dim, config.decoder_dim, true, device)
            .map_err(|e| crate::TtsError::ModelError(format!("Decoder RNN input failed: {e}")))?;

        let decoder_rnn_hidden = Linear::new(config.decoder_dim, config.decoder_dim, true, device)
            .map_err(|e| crate::TtsError::ModelError(format!("Decoder RNN hidden failed: {e}")))?;

        let mel_projection = Linear::new(config.decoder_dim + config.encoder_dim, config.mel_dim, true, device)
            .map_err(|e| crate::TtsError::ModelError(format!("Mel projection failed: {e}")))?;

        let gate_projection = Linear::new(config.decoder_dim + config.encoder_dim, 1, true, device)
            .map_err(|e| crate::TtsError::ModelError(format!("Gate projection failed: {e}")))?;

        let postnet = PostNet::new(&config, device)?;

        Ok(Self {
            config,
            encoder,
            prenet,
            attention,
            decoder_rnn_input,
            decoder_rnn_hidden,
            mel_projection,
            gate_projection,
            postnet,
            device: device.clone(),
            training: true,
        })
    }

    /// Get configuration
    pub fn config(&self) -> &Tacotron2Config {
        &self.config
    }

    /// Decode one step
    fn decode_step(
        &self,
        decoder_input: &Tensor,
        decoder_hidden: &Tensor,
        encoder_outputs: &Tensor,
        attention_weights_cat: &Tensor,
    ) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
        // Prenet
        let prenet_out = self.prenet.forward(decoder_input)?;

        // Attention
        let attention_weights = self.attention.forward(
            decoder_hidden,
            encoder_outputs,
            attention_weights_cat,
        )?;

        // Apply attention to encoder outputs
        let attention_weights_expanded = attention_weights.unsqueeze(1)
            .map_err(|e| crate::TtsError::TensorError(format!("Unsqueeze failed: {e}")))?;

        let encoder_outputs_transposed = encoder_outputs.transpose(1, 2)
            .map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;

        let attention_context = rtx_tensor::ops::matmul(&attention_weights_expanded, &encoder_outputs_transposed)
            .map_err(|e| crate::TtsError::TensorError(format!("Matmul failed: {e}")))?;

        let attention_context = attention_context.squeeze(1)
            .map_err(|e| crate::TtsError::TensorError(format!("Squeeze failed: {e}")))?;

        // Concatenate prenet output and attention context
        let decoder_rnn_input = rtx_tensor::ops::cat(&[&prenet_out, &attention_context], 1)
            .map_err(|e| crate::TtsError::TensorError(format!("Cat failed: {e}")))?;

        // Decoder RNN step
        let rnn_input_proj = self.decoder_rnn_input.forward(&decoder_rnn_input)
            .map_err(|e| crate::TtsError::ModelError(format!("RNN input projection failed: {e}")))?;

        let rnn_hidden_proj = self.decoder_rnn_hidden.forward(decoder_hidden)
            .map_err(|e| crate::TtsError::ModelError(format!("RNN hidden projection failed: {e}")))?;

        let decoder_hidden_new = (&rnn_input_proj + &rnn_hidden_proj)
            .map_err(|e| crate::TtsError::TensorError(format!("Add failed: {e}")))?;

        let decoder_hidden_new = rtx_tensor::ops::tanh(&decoder_hidden_new)
            .map_err(|e| crate::TtsError::TensorError(format!("Tanh failed: {e}")))?;

        // Projection to mel
        let projection_input = rtx_tensor::ops::cat(&[&decoder_hidden_new, &attention_context], 1)
            .map_err(|e| crate::TtsError::TensorError(format!("Cat failed: {e}")))?;

        let mel_output = self.mel_projection.forward(&projection_input)
            .map_err(|e| crate::TtsError::ModelError(format!("Mel projection failed: {e}")))?;

        let gate_output = self.gate_projection.forward(&projection_input)
            .map_err(|e| crate::TtsError::ModelError(format!("Gate projection failed: {e}")))?;

        Ok((mel_output, gate_output, decoder_hidden_new, attention_weights))
    }
}

impl AcousticModel for Tacotron2 {
    fn encode(&self, phonemes: &Tensor, _phoneme_lengths: Option<&Tensor>) -> Result<Tensor> {
        self.encoder.forward(phonemes)
    }

    fn decode(&self, hidden: &Tensor) -> Result<Tensor> {
        let shape = hidden.shape();
        let batch_size = shape.dims()[0];
        let seq_len = shape.dims()[1];

        // Initialize decoder state
        let decoder_hidden = Tensor::zeros(&[batch_size, self.config.decoder_dim], &self.device)
            .map_err(|e| crate::TtsError::TensorError(format!("Zeros failed: {e}")))?;

        let decoder_input = Tensor::zeros(&[batch_size, self.config.mel_dim], &self.device)
            .map_err(|e| crate::TtsError::TensorError(format!("Zeros failed: {e}")))?;

        let attention_weights_cat = Tensor::zeros(&[batch_size, 2, seq_len], &self.device)
            .map_err(|e| crate::TtsError::TensorError(format!("Zeros failed: {e}")))?;

        let mut mel_outputs = Vec::new();
        let mut current_hidden = decoder_hidden;
        let mut current_input = decoder_input;
        let mut current_attn = attention_weights_cat;

        // Autoregressive decoding
        for _ in 0..self.config.max_decoder_steps.min(200) {
            let (mel_out, gate_out, new_hidden, attn_weights) =
                self.decode_step(&current_input, &current_hidden, hidden, &current_attn)?;

            mel_outputs.push(mel_out.clone());
            current_hidden = new_hidden;
            current_input = mel_out;

            // Update attention weights cat
            let attn_expanded = attn_weights.unsqueeze(1)
                .map_err(|e| crate::TtsError::TensorError(format!("Unsqueeze failed: {e}")))?;

            current_attn = rtx_tensor::ops::cat(&[
                &current_attn.slice(1, 1, 2)
                    .map_err(|e| crate::TtsError::TensorError(format!("Slice failed: {e}")))?,
                &attn_expanded
            ], 1).map_err(|e| crate::TtsError::TensorError(format!("Cat failed: {e}")))?;

            // Check gate (stop condition)
            let gate_sigmoid = rtx_tensor::ops::sigmoid(&gate_out)
                .map_err(|e| crate::TtsError::TensorError(format!("Sigmoid failed: {e}")))?;

            let gate_val = gate_sigmoid.to_vec::<f32>()
                .map_err(|e| crate::TtsError::TensorError(format!("To vec failed: {e}")))?;

            if gate_val.iter().all(|&v| v > 0.5) {
                break;
            }
        }

        // Stack mel outputs: [batch, time, mel_dim]
        let mel_stacked = rtx_tensor::ops::stack(&mel_outputs.iter().collect::<Vec<_>>(), 1)
            .map_err(|e| crate::TtsError::TensorError(format!("Stack failed: {e}")))?;

        // Transpose to [batch, mel_dim, time]
        let mel_transposed = mel_stacked.transpose(1, 2)
            .map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;

        // Apply postnet
        let postnet_out = self.postnet.forward(&mel_transposed)?;

        // Add residual
        (&mel_transposed + &postnet_out)
            .map_err(|e| crate::TtsError::TensorError(format!("Add residual failed: {e}")))
    }

    fn forward(&self, phonemes: &Tensor, phoneme_lengths: Option<&Tensor>) -> Result<Tensor> {
        let encoded = self.encode(phonemes, phoneme_lengths)?;
        self.decode(&encoded)
    }

    fn mel_config(&self) -> &MelSpectrogramConfig {
        &self.config.mel_config
    }

    fn set_training(&mut self, training: bool) {
        self.training = training;
        self.encoder.set_training(training);
        self.prenet.set_training(training);
        self.postnet.set_training(training);
    }

    fn is_training(&self) -> bool {
        self.training
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rtx_tensor::Device;

    fn get_device() -> Device {
        Device::cuda(0).unwrap_or(Device::default())
    }

    #[test]
    fn test_config_default() {
        let config = Tacotron2Config::default();
        assert_eq!(config.encoder_dim, 512);
        assert_eq!(config.decoder_dim, 1024);
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_config_validation() {
        let mut config = Tacotron2Config::default();
        config.encoder_dim = 0;
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_prenet_creation() {
        let device = get_device();
        let prenet = PreNet::new(80, 256, 0.5, &device);
        assert!(prenet.is_ok());
    }

    #[test]
    fn test_prenet_forward() {
        let device = get_device();
        let prenet = PreNet::new(80, 256, 0.5, &device).unwrap();
        let input = Tensor::randn(&[2, 80], &device).unwrap();
        let output = prenet.forward(&input);
        assert!(output.is_ok());
    }

    #[test]
    fn test_postnet_creation() {
        let device = get_device();
        let config = Tacotron2Config::default();
        let postnet = PostNet::new(&config, &device);
        assert!(postnet.is_ok());
    }

    #[test]
    fn test_encoder_creation() {
        let device = get_device();
        let config = Tacotron2Config::default();
        let encoder = Encoder::new(&config, &device);
        assert!(encoder.is_ok());
    }

    #[test]
    fn test_encoder_forward() {
        let device = get_device();
        let config = Tacotron2Config::default();
        let encoder = Encoder::new(&config, &device).unwrap();

        let batch_size = 2;
        let seq_len = 10;
        let input = Tensor::randint(0, config.vocab_size as i64, &[batch_size, seq_len], &device).unwrap();

        let output = encoder.forward(&input);
        assert!(output.is_ok());

        let output = output.unwrap();
        let shape = output.shape();
        assert_eq!(shape.dims()[0], batch_size);
        assert_eq!(shape.dims()[1], seq_len);
        assert_eq!(shape.dims()[2], config.encoder_dim);
    }

    #[test]
    fn test_tacotron2_creation() {
        let device = get_device();
        let config = Tacotron2Config::default();
        let model = Tacotron2::new(config, &device);
        assert!(model.is_ok());
    }

    #[test]
    fn test_tacotron2_encode() {
        let device = get_device();
        let config = Tacotron2Config::default();
        let model = Tacotron2::new(config.clone(), &device).unwrap();

        let batch_size = 2;
        let seq_len = 10;
        let phonemes = Tensor::randint(0, config.vocab_size as i64, &[batch_size, seq_len], &device).unwrap();

        let encoded = model.encode(&phonemes, None);
        assert!(encoded.is_ok());

        let encoded = encoded.unwrap();
        let shape = encoded.shape();
        assert_eq!(shape.dims()[0], batch_size);
        assert_eq!(shape.dims()[1], seq_len);
        assert_eq!(shape.dims()[2], config.encoder_dim);
    }

    #[test]
    fn test_tacotron2_training_mode() {
        let device = get_device();
        let config = Tacotron2Config::default();
        let mut model = Tacotron2::new(config, &device).unwrap();

        assert!(model.is_training());
        model.set_training(false);
        assert!(!model.is_training());
    }

    #[test]
    fn test_config_serialization() {
        let config = Tacotron2Config::default();
        let json = serde_json::to_string(&config).unwrap();
        let deserialized: Tacotron2Config = serde_json::from_str(&json).unwrap();
        assert_eq!(config.encoder_dim, deserialized.encoder_dim);
    }

    #[test]
    fn test_location_attention_creation() {
        let device = get_device();
        let attention = LocationAttention::new(1024, 512, 128, &device);
        assert!(attention.is_ok());
    }
}