from cvae import CVAE
import torch
from typing import Sequence
import re

instruments = ['bass_acoustic', 'brass_acoustic', 'flute_acoustic', 'guitar_acoustic', 'keyboard_acoustic', 'mallet_acoustic', 'organ_acoustic', 'reed_acoustic', 'string_acoustic', 'synth_lead_acoustic', 'vocal_acoustic', 'bass_synthetic', 'brass_synthetic', 'flute_synthetic', 'guitar_synthetic', 'keyboard_synthetic', 'mallet_synthetic', 'organ_synthetic', 'reed_synthetic', 'string_synthetic', 'synth_lead_synthetic', 'vocal_synthetic', 'bass_electronic', 'brass_electronic', 'flute_electronic', 'guitar_electronic', 'keyboard_electronic', 'mallet_electronic', 'organ_electronic', 'reed_electronic', 'string_electronic', 'synth_lead_electronic', 'vocal_electronic']

model = CVAE.load_from_checkpoint(
    'epoch=17-step=650718.ckpt',
    io_channels=1,
    io_features=16000*4,
    latent_features=5,
    channels=[32, 64, 128, 256, 512],
    num_classes=len(instruments),
    learning_rate=1e-5
)

def format(text):
    text = text.split(' ')[-1]
    return text.replace(" ", "").lower()

def choice_to_tensor(choice: Sequence[str]) -> torch.Tensor:
    choice = '_'.join([format(i) for i in choice])
    return torch.tensor(instruments.index(choice))

def generate(choice: Sequence[str], params: Sequence[int]=None):
    noise = torch.tensor(params).unsqueeze(0).to('cuda') if params else torch.randn(1, 5).to('cuda')
    return model.sample(eps=noise, c = choice_to_tensor(choice).to('cuda')).cpu().numpy()[0]