import os
import numpy as np
import torch
import argparse

from hparams import create_hparams
from model import lcm
from train import load_model
from torch.utils.data import DataLoader
from reader import TextMelIDLoader, TextMelIDCollate, id2sp
from inference_utils import plot_data

parser = argparse.ArgumentParser()
parser.add_argument('-c', '--checkpoint_path', type=str,
                        help='directory to save checkpoints')
parser.add_argument('--hparams', type=str,
                        required=False, help='comma separated name=value pairs')
args = parser.parse_args()

checkpoint_path=args.checkpoint_path

hparams = create_hparams(args.hparams)

model = load_model(hparams)
model.load_state_dict(torch.load(checkpoint_path)['state_dict'], strict=False)
_ = model.eval()


def gen_embedding(speaker):

    training_list = hparams.training_list

    train_set_A = TextMelIDLoader(training_list, hparams.mel_mean_std, hparams.speaker_A,
            hparams.speaker_B, 
            shuffle=False,pids=[speaker])
            
    collate_fn = TextMelIDCollate(lcm(hparams.n_frames_per_step_encoder,
                            hparams.n_frames_per_step_decoder))

    train_loader_A = DataLoader(train_set_A, num_workers=1, shuffle=False,
                                sampler=None,
                                batch_size=1, pin_memory=False,
                                drop_last=True, collate_fn=collate_fn)

    with torch.no_grad():

        speaker_embeddings = []

        for i,batch in enumerate(train_loader_A):
            #print i
            x, y = model.parse_batch(batch)
            text_input_padded, mel_padded, text_lengths, mel_lengths, speaker_id = x
            speaker_id, speaker_embedding = model.speaker_encoder.inference(mel_padded)

            speaker_embedding = speaker_embedding.data.cpu().numpy()
            speaker_embeddings.append(speaker_embedding)

        speaker_embeddings = np.vstack(speaker_embeddings)
        
    print(speaker_embeddings.shape)
    if not os.path.exists('outdir/embeddings'):
        os.makedirs('outdir/embeddings')
    
    np.save('outdir/embeddings/%s.npy'%speaker, speaker_embeddings)
    plot_data([speaker_embeddings], 
        'outdir/embeddings/%s.pdf'%speaker)


print('Generating embedding of %s ...'%hparams.speaker_A)
gen_embedding(hparams.speaker_A)

print('Generating embedding of %s ...'%hparams.speaker_B)
gen_embedding(hparams.speaker_B)

[((13, 9, 13, 34), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((22, 10, 22, 38), 'hparams.create_hparams', 'create_hparams', ({(22, 25, 22, 37): 'args.hparams'}, {}), '(args.hparams)', False, 'from hparams import create_hparams\n'), ((24, 8, 24, 27), 'train.load_model', 'load_model', ({(24, 19, 24, 26): 'hparams'}, {}), '(hparams)', False, 'from train import load_model\n'), ((33, 18, 35, 41), 'reader.TextMelIDLoader', 'TextMelIDLoader', (), '', False, 'from reader import TextMelIDLoader, TextMelIDCollate, id2sp\n'), ((40, 21, 43, 70), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((64, 4, 64, 67), 'numpy.save', 'np.save', ({(64, 12, 64, 46): "('outdir/embeddings/%s.npy' % speaker)", (64, 48, 64, 66): 'speaker_embeddings'}, {}), "('outdir/embeddings/%s.npy' % speaker, speaker_embeddings)", True, 'import numpy as np\n'), ((65, 4, 66, 43), 'inference_utils.plot_data', 'plot_data', ({(65, 14, 65, 34): '[speaker_embeddings]', (66, 8, 66, 42): "('outdir/embeddings/%s.pdf' % speaker)"}, {}), "([speaker_embeddings], 'outdir/embeddings/%s.pdf' % speaker)", False, 'from inference_utils import plot_data\n'), ((25, 22, 25, 49), 'torch.load', 'torch.load', ({(25, 33, 25, 48): 'checkpoint_path'}, {}), '(checkpoint_path)', False, 'import torch\n'), ((37, 34, 38, 62), 'model.lcm', 'lcm', ({(37, 38, 37, 71): 'hparams.n_frames_per_step_encoder', (38, 28, 38, 61): 'hparams.n_frames_per_step_decoder'}, {}), '(hparams.n_frames_per_step_encoder, hparams.n_frames_per_step_decoder)', False, 'from model import lcm\n'), ((45, 9, 45, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((58, 29, 58, 58), 'numpy.vstack', 'np.vstack', ({(58, 39, 58, 57): 'speaker_embeddings'}, {}), '(speaker_embeddings)', True, 'import numpy as np\n'), ((61, 11, 61, 46), 'os.path.exists', 'os.path.exists', ({(61, 26, 61, 45): '"""outdir/embeddings"""'}, {}), "('outdir/embeddings')", False, 'import os\n'), ((62, 8, 62, 40), 'os.makedirs', 'os.makedirs', ({(62, 20, 62, 39): '"""outdir/embeddings"""'}, {}), "('outdir/embeddings')", False, 'import os\n')]