""" 
test_conditioning.py
    Desc: In order to make sure that we can condition on phoneme encodings, need to ensure that
        the model can take in 2D Embeddings
"""

import torch 
from audio_diffusion_pytorch import DiffusionModel, UNetV0, VDiffusion, VSampler

import sys 
sys.path.append('.')

# Example build model function
def create_model():
    return DiffusionModel(
    net_t=UNetV0, # The model type used for diffusion 
    in_channels=2, # U-Net: number of input/output (audio) channels 
    channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer
    factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer 
    items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer
    attentions=[0, 0, 0, 0, 0, 1, 1, 1, 1], # U-Net: attention enabled/disabled at each layer 
    attention_heads=8, # U-Net: number of attention heads per attention block
    attention_features=64, # U-Net: number of attention features per attention block,
    diffusion_t=VDiffusion, # The diffusion method used 
    sampler_t=VSampler, # The diffusion sampler used 
)

def create_cond_model():
    return DiffusionModel(
    net_t=UNetV0, # The model type used for diffusion 
    in_channels=2, # U-Net: number of input/output (audio) channels 
    channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer
    factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer 
    items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer
    attentions=[0, 0, 0, 0, 0, 1, 1, 1, 1], # U-Net: attention enabled/disabled at each layer 
    attention_heads=8, # U-Net: number of attention heads per attention block
    attention_features=64, # U-Net: number of attention features per attention block,
    diffusion_t=VDiffusion, # The diffusion method used 
    sampler_t=VSampler, # The diffusion sampler used 
    embedding_features=768, # U-Net: embedding features
    cross_attentions=[0, 0, 0, 1, 1, 1, 1, 1, 1], # U-Net: cross-attention enabled/disabled at each layer 
)

def create_mel_model():
    return DiffusionModel(
    net_t=UNetV0, # The model type used for diffusion 
    in_channels=64, # U-Net: number of input/output (audio) channels 
    channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer
    factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer 
    items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer
    attentions=[0, 0, 0, 0, 0, 1, 1, 1, 1], # U-Net: attention enabled/disabled at each layer 
    attention_heads=8, # U-Net: number of attention heads per attention block
    attention_features=64, # U-Net: number of attention features per attention block,
    diffusion_t=VDiffusion, # The diffusion method used 
    sampler_t=VSampler, # The diffusion sampler used 
    embedding_features=768, # U-Net: embedding features
    cross_attentions=[0, 0, 0, 1, 1, 1, 1, 1, 1], # U-Net: cross-attention enabled/disabled at each layer 
)

# A file function for conditioning 2D Conditioning with random values
def test_2D_condition():
    model = create_model()
    cond_model = create_cond_model()

    # Pretend that input is waveform
    audio_wave = torch.randn(1, 2, 2**18) # [batch, in_channels, length]
    embedding = torch.randn(1, 1, 768) # [batch, num_embeddings, embedding_features] 

    y_nocond = model(audio_wave)
    y = cond_model(audio_wave, embedding=embedding)

    # Lets Try multiple Embeddings
    embedding = torch.randn(1, 20, 768) # [batch, num_embeddings, embedding_features] 
    y = cond_model(audio_wave, embedding=embedding)

    del model 
    del cond_model

    mel_model = create_mel_model()

    # Pretend that input is mel-spec
    audio_wave = torch.randn(1, 64, 2**18) # [batch, in_channels, length]
    embedding = torch.randn(1, 64, 768) # [batch, num_embeddings, embedding_features]; should be [batch, length, phone_embed_dim] 

    y = mel_model(audio_wave, embedding = embedding)

if __name__ == "__main__":
    # Note: The below function works as is 
    test_2D_condition()