# *****************************************************************************
#  Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of the NVIDIA CORPORATION nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# *****************************************************************************
import argparse
import json
import os
os.environ["LRU_CACHE_CAPACITY"] = "3"
import torch
import os.path
import sys
import time
from math import ceil, e, exp
import math

import numpy as np
import soundfile as sf

save_file_check_path = "save"

#=====START: ADDED FOR DISTRIBUTED======
from distributed import init_distributed, apply_gradient_allreduce, reduce_tensor
from torch.utils.data.distributed import DistributedSampler
from torch.distributed import broadcast
#=====END:   ADDED FOR DISTRIBUTED======

from torch.utils.data import DataLoader
from mel2samp import Mel2Samp
from tqdm import tqdm
import CookieTTS.utils.audio.stft as STFT

from mel2samp import load_wav_to_torch
from scipy import signal

class LossExplosion(Exception):
    """Custom Exception Class. If Loss Explosion, raise Error and automatically restart training script from previous best_val_model checkpoint."""
    pass

class StreamingMovingAverage:
    def __init__(self, window_size):
        self.window_size = window_size
        self.values = []
        self.sum = 0

    def process(self, value):
        self.values.append(value)
        self.sum += value
        if len(self.values) > self.window_size:
            self.sum -= self.values.pop(0)
        return float(self.sum) / len(self.values)

def load_checkpoint(checkpoint_paths, model, optimizer, scheduler, fp16_run, warm_start=False, nvidia_checkpoint=False):
    checkpoint_paths = [checkpoint_paths,] if type(checkpoint_paths) == str else checkpoint_paths

    model_dict = model.state_dict() # just initialized state_dict

    for checkpoint_path in checkpoint_paths:
        assert os.path.isfile(checkpoint_path), f'"{checkpoint_path}" does not exist'
        checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')

        if 'iteration' in checkpoint_dict:
            iteration = checkpoint_dict['iteration']

        if not warm_start:
            if 'optimizer' in checkpoint_dict.keys():
                optimizer.load_state_dict(checkpoint_dict['optimizer'])
            if fp16_run and 'amp' in checkpoint_dict.keys():
                amp.load_state_dict(checkpoint_dict['amp'])
            if scheduler and 'scheduler' in checkpoint_dict.keys():
                scheduler.load_state_dict(checkpoint_dict['scheduler'])

        checkpoint_model_dict = checkpoint_dict['model']
        if (str(type(checkpoint_model_dict)) != "<class 'collections.OrderedDict'>"):
            checkpoint_model_dict = checkpoint_model_dict.state_dict()

        if nvidia_checkpoint:# (optional) convert Nvidia/WaveGlow into my format
            checkpoint_model_dict = {k.replace(".in_layers",".WN.in_layers").replace(".res_skip_layers",".WN.res_skip_layers").replace(".start",".WN.start").replace(".end",".WN.end").replace(".conv.weight",".weight"): v for k, v in checkpoint_model_dict.items()}

        if warm_start:
            model_dict_missing = {k: v for k,v in checkpoint_model_dict.items() if k not in model_dict}
            if model_dict_missing:
                print(f"## {len(model_dict_missing.keys())} keys found that do not exist in the current model ##")
                print( '\n'.join([str(x) for x in model_dict_missing.keys()]) + '\n' )

            model_dict_mismatching = {k: v for k,v in checkpoint_model_dict.items() if k in model_dict and checkpoint_model_dict[k].shape != model_dict[k].shape}
            if model_dict_mismatching:
                print(f"## {len(model_dict_mismatching.keys())} keys found that do not match the shape in the current model ##")
                print( '\n'.join([str(x) for x in model_dict_mismatching.keys()]) + '\n' )

            pretrained_missing = {k: v for k,v in model_dict.items() if k not in checkpoint_model_dict}
            if pretrained_missing:
                print(f"## {len(pretrained_missing.keys())} keys in the current model that do not have a matching key in the checkpoint ##")
                print( '\n'.join([str(x) for x in pretrained_missing.keys()]) + '\n' )

            # Fiter out unneccessary keys
            filtered_dict = {k: v for k,v in checkpoint_model_dict.items() if k in model_dict and checkpoint_model_dict[k].shape == model_dict[k].shape}
            model_dict.update(filtered_dict)
        else:
            model_dict = {k.replace("invconv1x1","convinv").replace(".F.",".WN.").replace("WNs.","WN."): v for k, v in checkpoint_model_dict.items()}
        print(f"Loaded checkpoint path '{checkpoint_path}'")

    model.load_state_dict(model_dict)
    print(f"New state_dict loaded! (iteration {iteration})")
    return model, optimizer, iteration, scheduler

def save_checkpoint(model, optimizer, learning_rate, iteration, amp, scheduler, speaker_lookup, filepath):
    tqdm.write("Saving model and optimizer state at iteration {} to {}".format(
          iteration, filepath))
    #model_for_saving = WaveGlow(**waveglow_config).cuda()
    #model_for_saving.load_state_dict(model.state_dict())
    #state_dict = model_for_saving.state_dict()
    state_dict = model.state_dict()
    saving_dict = {'model': state_dict,
        'iteration': iteration,
        'optimizer': optimizer.state_dict(),
        'scheduler': scheduler.state_dict(),
        'learning_rate': learning_rate,
        'speaker_lookup': speaker_lookup,
        'waveglow_config': waveglow_config,
        }
    if amp: saving_dict['amp'] = amp.state_dict()
    torch.save(saving_dict, filepath)
    tqdm.write("Model Saved")

def save_weights(model, optimizer, learning_rate, iteration, filepath):
    print("Saving model weights without optimizer at iteration {} to {}".format(
          iteration, filepath))
    model_for_saving = WaveGlow(**waveglow_config).cuda()
    model_for_saving.load_state_dict(model.state_dict())
    #torch.save(model_for_saving.state_dict(), filepath)
    torch.save({'model': model_for_saving,
                'iteration': iteration}, filepath)
    print("Weights Saved")

import matplotlib
matplotlib.use("Agg")
import matplotlib.pylab as plt
def save_figure_to_numpy(fig):
    # save it to a numpy array.
    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    return data
def plot_spectrogram_to_numpy(spectrogram, range=None):
    fig, ax = plt.subplots(figsize=(16, 3))
    im = ax.imshow(spectrogram, cmap='inferno', aspect="auto", origin="lower",
                   interpolation='none')
    if range is not None:
        assert len(range) == 2, 'range params should be a 2 element List of [Min, Max].'
        assert range[1] > range[0], 'Max (element 1) must be greater than Min (element 0).'
        im.set_clim(range[0], range[1])
    plt.colorbar(im, ax=ax)
    plt.xlabel("Frames")
    plt.ylabel("Channels")
    plt.tight_layout()

    fig.canvas.draw()
    data = save_figure_to_numpy(fig)
    plt.close()
    return data

def validate(model, loader_STFT, STFTs, logger, iteration, validation_files, speaker_lookup, sigma, output_directory, data_config, save_audio=True, max_length_s= 3, files_to_process=9):
    print("Validating... ", end="")
    from mel2samp import DTW
    val_sigma = sigma * 1.00
    model.eval()
    val_start_time = time.time()
    STFT_elapsed = samples_processed = 0
    with torch.no_grad():
        with torch.random.fork_rng(devices=[0,]):
            torch.random.manual_seed(0)# use same Z / random seed during validation so results are more consistent and comparable.

            with open(validation_files, encoding='utf-8') as f:
                audiopaths_and_melpaths = [line.strip().split('|') for line in f]

            if next(model.parameters()).type() == "torch.cuda.HalfTensor":
                model_type = "half"
            else:
                model_type = "float"

            timestr = time.strftime("%Y_%m_%d-%H_%M_%S")
            total_MAE = total_MSE = total = files_processed = 0
            input_mels = []
            gt_mels = []
            pred_mels = []
            MAE_specs = []
            for i, (audiopath, melpath, *remaining) in enumerate(audiopaths_and_melpaths):
                if files_processed >= files_to_process: # number of validation files to run.
                    break
                audio = load_wav_to_torch(audiopath)[0]/32768.0 # load audio from wav file to tensor
                if audio.shape[0] > (data_config['sampling_rate']*max_length_s):
                    continue # ignore audio over max_length_seconds

                gt_mel = loader_STFT.mel_spectrogram(audio.unsqueeze(0)).cuda()# [T] -> [1, T] -> [1, n_mel, T]
                if 'load_hidden_from_disk' in data_config.keys() and data_config['load_hidden_from_disk']:
                    mel = None
                    hidden_path = remaining[1]
                    model_input = np.load(hidden_path) # load tacotron hidden from file into numpy arr
                    model_input = torch.from_numpy(model_input).unsqueeze(0).cuda() # from numpy arr to tensor on GPU
                else:
                    if loader_STFT and data_config['load_mel_from_disk'] < 0.2:
                        mel = None
                        model_input = gt_mel.clone()
                    else:
                        mel = np.load(melpath) # load mel from file into numpy arr
                        mel = torch.from_numpy(mel).unsqueeze(0).cuda() # from numpy arr to tensor on GPU
                        assert mel[:, :gt_mel.shape[1], :].shape == gt_mel.shape, f'shapes of {mel[:, :gt_mel.shape[1], :].shape} and {gt_mel.shape} do not match.'
                        #if torch.nn.functional.mse_loss(mel[:, :gt_mel.shape[1], :], gt_mel) > 0.3:
                        #    continue # skip validation files that significantly vary from the target.
                        if model.has_logvar_channels:
                            if mel.shape[1] == model.n_mel_channels*2:
                                mel, logvar = mel.chunk(2, dim=1)# [1, n_mel*2, T] -> [1, n_mel, T], [1, n_mel, T]
                                mel = DTW(mel, gt_mel, scale_factor=8, range_=5)
                                model_input = torch.cat((mel, logvar), dim=1)
                            else:
                                raise Exception("Loaded mel from disk has wrong shape.")
                        else:
                            if mel.shape[1] == model.n_mel_channels*2:
                                mel = mel.chunk(2, dim=1)[0]# [1, n_mel*2, T] -> [1, n_mel, T]
                            #mel = DTW(mel, gt_mel, scale_factor=8, range_=5)
                            model_input = mel

                if hasattr(model, 'multispeaker') and model.multispeaker == True:
                    assert len(remaining), f"Speaker ID missing while multispeaker == True.\nLine: {i}\n'{'|'.join([autiopath, melpath])}'"
                    speaker_id = remaining[0]
                    speaker_id = torch.IntTensor([speaker_lookup[int(speaker_id)],])
                    speaker_id = speaker_id.cuda(non_blocking=True).long()
                else:
                    speaker_id = None

                if model_type == "half":
                    model_input = model_input.half() # for fp16 training

                audio_waveglow = model.infer(model_input, speaker_id, sigma=val_sigma).cpu().float()

                audio = audio.squeeze().unsqueeze(0) # crush extra dimensions and shape for STFT
                audio_waveglow = audio_waveglow.squeeze().unsqueeze(0).clamp(min=-0.999, max=0.999) # [1, T] crush extra dimensions and shape for STFT
                audio_waveglow[torch.isnan(audio_waveglow) | torch.isinf(audio_waveglow)] = 0.0 # and clamp any values over/under |1.0| (which should only exist very early in training)

                STFT_start_time = time.time()
                for j, STFT in enumerate(STFTs): # check Spectrogram Error with multiple window sizes
                    input_mels.append(mel)
                    mel_GT = STFT.mel_spectrogram(audio)# [1, T] -> [1, n_mel, T//hop_len]
                    gt_mels.append(mel_GT)

                    mel_waveglow = STFT.mel_spectrogram(audio_waveglow)[:,:,:mel_GT.shape[-1]]# [1, T] -> [1, n_mel, T//hop_len]
                    pred_mels.append(mel_waveglow)

                    MSE = (torch.nn.MSELoss()(mel_waveglow, mel_GT)).item() # get MSE (Mean Squared Error) between Ground Truth and WaveGlow inferred spectrograms.
                    MAE_spec = torch.nn.L1Loss(reduction='none')(mel_waveglow, mel_GT)
                    MAE = (MAE_spec.mean()).item() # get MAE (Mean Absolute Error) between Ground Truth and WaveGlow inferred spectrograms.
                    MAE_specs.append(MAE_spec)

                    total_MAE+=MAE
                    total_MSE+=MSE
                    total+=1
                STFT_elapsed += time.time()-STFT_start_time

                if save_audio:
                    audio_path = os.path.join(output_directory, "samples", str(iteration)+"-"+timestr, os.path.basename(audiopath)) # Write audio to checkpoint_directory/iteration/audiofilename.wav
                    os.makedirs(os.path.join(output_directory, "samples", str(iteration)+"-"+timestr), exist_ok=True)
                    sf.write(audio_path, audio_waveglow.squeeze().cpu().numpy(), data_config['sampling_rate'], "PCM_16") # save waveglow sample

                    audio_path = os.path.join(output_directory, "samples", "Ground Truth", os.path.basename(audiopath)) # Write audio to checkpoint_directory/iteration/audiofilename.wav
                    if not os.path.exists(audio_path):
                        os.makedirs(os.path.join(output_directory, "samples", "Ground Truth"), exist_ok=True)
                        sf.write(audio_path, audio.squeeze().cpu().numpy(), data_config['sampling_rate'], "PCM_16") # save ground truth
                files_processed+=1
                samples_processed+=audio_waveglow.shape[-1]

    if total:
        average_MSE = total_MSE/total
        average_MAE = total_MAE/total
        logger.add_scalar('val_MSE', average_MSE, iteration)
        logger.add_scalar('val_MAE', average_MAE, iteration)

        for idx, (gt_mel, pred_mel, input_mel, mae_mel) in enumerate(zip(gt_mels[-6:], pred_mels[-6:], input_mels[-6:], MAE_specs[-6:])):
            logger.add_image(f'mel_{idx}/pred',
                        plot_spectrogram_to_numpy(pred_mel[0].data.cpu().numpy(), range=[-11.5, 2.0]),
                        iteration, dataformats='HWC')
            if mae_mel is not None:
                logger.add_image(f'mel_{idx}/zmae',
                        plot_spectrogram_to_numpy(mae_mel[0].data.cpu().numpy(), range=[0.0, 2.5]),
                        iteration, dataformats='HWC')
            if iteration % 10000 == 0: # target doesn't change unless batch size or dataset changes so only needs to be plotted once in a while.
                logger.add_image(f'mel_{idx}/target',
                            plot_spectrogram_to_numpy(gt_mel[0].data.cpu().numpy(), range=[-11.5, 2.0]),
                            iteration, dataformats='HWC')
                if input_mel is not None:
                    logger.add_image(f'mel_{idx}/input',
                            plot_spectrogram_to_numpy(input_mel[0].data.cpu().numpy(), range=[-11.5, 2.0]),
                            iteration, dataformats='HWC')

        time_elapsed = time.time()-val_start_time
        time_elapsed_without_stft = time_elapsed-STFT_elapsed
        samples_per_second = samples_processed/time_elapsed_without_stft
        print(f"[Avg MSE: {average_MSE:.6f} MAE: {average_MAE:.6f}]",
            f"[{time_elapsed_without_stft:.3f}s]",
            f"[{time_elapsed:.3f}s_stft]",
            f"[{time_elapsed_without_stft/files_processed:.3f}s/file]",
            f"[{samples_per_second/data_config['sampling_rate']:.3f}rtf]",
            f"[{samples_per_second:,.0f}samples/s]"
          )
        logger.add_scalar('val_rtf', samples_per_second/data_config['sampling_rate'], iteration)
    else:
        average_MSE = 1e3
        average_MAE = 1e3
        print("Average MSE: N/A", "Average MAE: N/A")

    for convinv in model.convinv:
        if hasattr(convinv, 'W_inverse'):
            delattr(convinv, "W_inverse") # clear Inverse Weights.
    if hasattr(model, 'iso226'):
        delattr(model, 'iso226')
    mel = speaker_id = None # del GPU based tensors.
    torch.cuda.empty_cache() # clear cache for next training
    model.train()

    return average_MSE, average_MAE

def multiLR(model):
    model_parameters = []
    model_parameters.append( {"params": list(model.children())[0].parameters()} ) # ConvTranspose1d - Mel upscaling layer
    model_parameters.append( {"params": list(model.children())[2].parameters()} ) # Invertible1x1Conv - ConvInv

    # [0, 1]
    # [2:2+in_layers_len]

    in_layers_ = [ {"params": x.in_layers.parameters()} for x in list(model.children())[1] ]
    in_layers_offsets = [2, 2+len(in_layers_)]
    model_parameters.extend( in_layers_ ) # in layers

    res_skip_layers_ = [ {"params": x.res_skip_layers.parameters()} for x in list(model.children())[1] ]
    res_skip_layers_offsets = [in_layers_offsets[1], in_layers_offsets[1]+len(res_skip_layers_)]
    model_parameters.extend( res_skip_layers_ ) # residual and skip layers

    start_ = [ {"params": x.start.parameters()} for x in list(model.children())[1] ]
    start_offsets = [res_skip_layers_offsets[1], res_skip_layers_offsets[1]+len(start_)]
    model_parameters.extend( start_ ) # start layers

    end_ = [ {"params": x.end.parameters()} for x in list(model.children())[1] ]
    end_offsets = [start_offsets[1], start_offsets[1]+len(end_)]
    model_parameters.extend( end_ ) # end layers

    cond_layer_ = [ {"params": x.cond_layer.parameters()} for x in list(model.children())[1] ]
    cond_layer_offsets = [end_offsets[1], end_offsets[1]+len(cond_layer_)]
    model_parameters.extend( cond_layer_ ) # conditioning layers
    return model_parameters, (in_layers_offsets, res_skip_layers_offsets, start_offsets, end_offsets, cond_layer_offsets)

def train(num_gpus, rank, group_name, output_directory, epochs, learning_rate,
          sigma, loss_empthasis, iters_per_checkpoint, batch_size, seed, fp16_run,
          checkpoint_path, with_tensorboard, logdirname, datedlogdir, warm_start=False, optimizer='ADAM', start_zero=False):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    #=====START: ADDED FOR DISTRIBUTED======
    if num_gpus > 1:
        init_distributed(rank, num_gpus, group_name, **dist_config)
    #=====END:   ADDED FOR DISTRIBUTED======

    global WaveGlow
    global WaveGlowLoss

    ax = True # this is **really** bad coding practice :D
    if ax:
        from efficient_model_ax import WaveGlow
        from efficient_loss import WaveGlowLoss
    else:
        if waveglow_config["yoyo"]: # efficient_mode # TODO: Add to Config File
            from efficient_model import WaveGlow
            from efficient_loss import WaveGlowLoss
        else:
            from glow import WaveGlow, WaveGlowLoss

    criterion = WaveGlowLoss(sigma, loss_empthasis)
    model = WaveGlow(**waveglow_config).cuda()
    #=====START: ADDED FOR DISTRIBUTED======
    if num_gpus > 1:
        model = apply_gradient_allreduce(model)
    #=====END:   ADDED FOR DISTRIBUTED======
    STFTs = [STFT.TacotronSTFT(filter_length=window,
                                 hop_length=data_config['hop_length'],
                                 win_length=window,
                                 sampling_rate=data_config['sampling_rate'],
                                 n_mel_channels=160,
                                 mel_fmin=0, mel_fmax=16000) for window in data_config['validation_windows']]

    loader_STFT = STFT.TacotronSTFT(filter_length=data_config['filter_length'],
                                 hop_length=data_config['hop_length'],
                                 win_length=data_config['win_length'],
                                 sampling_rate=data_config['sampling_rate'],
                                 n_mel_channels=data_config['n_mel_channels'] if 'n_mel_channels' in data_config.keys() else 160,
                                 mel_fmin=data_config['mel_fmin'], mel_fmax=data_config['mel_fmax'])

    #optimizer = "Adam"
    optimizer = optimizer.lower()
    optimizer_fused = bool( 0 ) # use Apex fused optimizer, should be identical to normal but slightly faster and only works on RTX cards
    if optimizer_fused:
        from apex import optimizers as apexopt
        if optimizer == "adam":
            optimizer = apexopt.FusedAdam(model.parameters(), lr=learning_rate)
        elif optimizer == "lamb":
            optimizer = apexopt.FusedLAMB(model.parameters(), lr=learning_rate, max_grad_norm=200)
    else:
        if optimizer == "adam":
            optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
        elif optimizer == "lamb":
            from lamb import Lamb as optLAMB
            optimizer = optLAMB(model.parameters(), lr=learning_rate)
            #import torch_optimizer as optim
            #optimizer = optim.Lamb(model.parameters(), lr=learning_rate)
            #raise# PyTorch doesn't currently include LAMB optimizer.

    if fp16_run:
        global amp
        from apex import amp
        model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
    else:
        amp = None

    ## LEARNING RATE SCHEDULER
    if True:
        from torch.optim.lr_scheduler import ReduceLROnPlateau
        min_lr = 1e-8
        factor = 0.1**(1/5) # amount to scale the LR by on Validation Loss plateau
        scheduler = ReduceLROnPlateau(optimizer, 'min', factor=factor, patience=20, cooldown=2, min_lr=min_lr, verbose=True, threshold=0.0001, threshold_mode='abs')
        print("ReduceLROnPlateau used as Learning Rate Scheduler.")
    else: scheduler=False

    # Load checkpoint if one exists
    iteration = 0
    if checkpoint_path != "":
        model, optimizer, iteration, scheduler = load_checkpoint(checkpoint_path, model,
                                                      optimizer, scheduler, fp16_run, warm_start=warm_start)
        iteration += 1  # next iteration is iteration + 1
    if start_zero:
        iteration = 0

    trainset = Mel2Samp(**data_config, check_files=True)
    speaker_lookup = trainset.speaker_ids
    # =====START: ADDED FOR DISTRIBUTED======
    if num_gpus > 1:
        train_sampler = DistributedSampler(trainset, shuffle=True)
        shuffle = False
    else:
        train_sampler = None
        shuffle = True
    # =====END:   ADDED FOR DISTRIBUTED======
    train_loader = DataLoader(trainset, num_workers=3, shuffle=shuffle,
                              sampler=train_sampler,
                              batch_size=batch_size,
                              pin_memory=False,
                              drop_last=True)

    # Get shared output_directory ready
    if rank == 0:
        if not os.path.isdir(output_directory):
            os.makedirs(output_directory)
            os.chmod(output_directory, 0o775)
        print("output directory", output_directory)

    if with_tensorboard and rank == 0:
        from tensorboardX import SummaryWriter
        if datedlogdir:
            timestr = time.strftime("%Y_%m_%d-%H_%M_%S")
            log_directory = os.path.join(output_directory, logdirname, timestr)
        else:
            log_directory = os.path.join(output_directory, logdirname)
        logger = SummaryWriter(log_directory)

    moving_average = int(min(len(train_loader), 100)) # average loss over entire Epoch
    rolling_sum = StreamingMovingAverage(moving_average)
    start_time = time.time()
    start_time_iter = time.time()
    start_time_dekaiter = time.time()
    model.train()

    # best (averaged) training loss
    if os.path.exists(os.path.join(output_directory, "best_model")+".txt"):
        best_model_loss = float(str(open(os.path.join(output_directory, "best_model")+".txt", "r", encoding="utf-8").read()).split("\n")[0])
    else:
        best_model_loss = -6.20

    # best (validation) MSE on inferred spectrogram.
    if os.path.exists(os.path.join(output_directory, "best_val_model")+".txt"):
        best_MSE = float(str(open(os.path.join(output_directory, "best_val_model")+".txt", "r", encoding="utf-8").read()).split("\n")[0])
    else:
        best_MSE = 9e9

    epoch_offset = max(0, int(iteration / len(train_loader)))

    pytorch_total_params = sum(p.numel() for p in model.parameters())
    print("{:,} total parameters in model".format(pytorch_total_params))
    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print("{:,} trainable parameters.".format(pytorch_total_params))

    print(f"Segment Length: {data_config['segment_length']:,}\nBatch Size: {batch_size:,}\nNumber of GPUs: {num_gpus:,}\nSamples/Iter: {data_config['segment_length']*batch_size*num_gpus:,}")

    training = True
    while training:
        try:
            if rank == 0:
                epochs_iterator = tqdm(range(epoch_offset, epochs), initial=epoch_offset, total=epochs, smoothing=0.01, desc="Epoch", position=1, unit="epoch")
            else:
                epochs_iterator = range(epoch_offset, epochs)
            # ================ MAIN TRAINING LOOP! ===================
            for epoch in epochs_iterator:
                print(f"Epoch: {epoch}")
                if num_gpus > 1:
                    train_sampler.set_epoch(epoch)

                if rank == 0:
                    iters_iterator = tqdm(enumerate(train_loader), desc=" Iter", smoothing=0, total=len(train_loader), position=0, unit="iter", leave=True)
                else:
                    iters_iterator = enumerate(train_loader)
                for i, batch in iters_iterator:
                    # run external code every iter, allows the run to be adjusted without restarts
                    if (i==0 or iteration % param_interval == 0):
                        try:
                            with open("run_every_epoch.py") as f:
                                internal_text = str(f.read())
                                if len(internal_text) > 0:
                                    #code = compile(internal_text, "run_every_epoch.py", 'exec')
                                    ldict = {'iteration': iteration, 'seconds_elapsed': time.time()-start_time}
                                    exec(internal_text, globals(), ldict)
                                else:
                                    print("No Custom code found, continuing without changes.")
                        except Exception as ex:
                            print(f"Custom code FAILED to run!\n{ex}")
                        globals().update(ldict)
                        locals().update(ldict)
                        if show_live_params:
                            print(internal_text)
                    if not iteration % 50: # check actual learning rate every 20 iters (because I sometimes see learning_rate variable go out-of-sync with real LR)
                        learning_rate = optimizer.param_groups[0]['lr']
                    # Learning Rate Schedule
                    if custom_lr:
                        old_lr = learning_rate
                        if iteration < warmup_start:
                            learning_rate = warmup_start_lr
                        elif iteration < warmup_end:
                            learning_rate = (iteration-warmup_start)*((A_+C_)-warmup_start_lr)/(warmup_end-warmup_start) + warmup_start_lr # learning rate increases from warmup_start_lr to A_ linearly over (warmup_end-warmup_start) iterations.
                        else:
                            if iteration < decay_start:
                                learning_rate = A_ + C_
                            else:
                                iteration_adjusted = iteration - decay_start
                                learning_rate = (A_*(e**(-iteration_adjusted/B_))) + C_
                        assert learning_rate > -1e-8, "Negative Learning Rate."
                        if old_lr != learning_rate:
                            for param_group in optimizer.param_groups:
                                param_group['lr'] = learning_rate
                    else:
                        scheduler.patience = scheduler_patience
                        scheduler.cooldown = scheduler_cooldown
                        if override_scheduler_last_lr:
                            scheduler._last_lr = override_scheduler_last_lr
                        if override_scheduler_best:
                            scheduler.best = override_scheduler_best
                        if override_scheduler_last_lr or override_scheduler_best:
                            print("scheduler._last_lr =", scheduler._last_lr, "scheduler.best =", scheduler.best, "  |", end='')
                    model.zero_grad()
                    mel, audio, speaker_ids = batch
                    mel = torch.autograd.Variable(mel.cuda(non_blocking=True))
                    audio = torch.autograd.Variable(audio.cuda(non_blocking=True))
                    speaker_ids = speaker_ids.cuda(non_blocking=True).long().squeeze(1)
                    outputs = model(mel, audio, speaker_ids)

                    loss = criterion(outputs)
                    if num_gpus > 1:
                        reduced_loss = reduce_tensor(loss.data, num_gpus).item()
                    else:
                        reduced_loss = loss.item()

                    if fp16_run:
                        with amp.scale_loss(loss, optimizer) as scaled_loss:
                            scaled_loss.backward()
                    else:
                        loss.backward()

                    if (reduced_loss > LossExplosionThreshold) or (math.isnan(reduced_loss)):
                        model.zero_grad()
                        raise LossExplosion(f"\nLOSS EXPLOSION EXCEPTION ON RANK {rank}: Loss reached {reduced_loss} during iteration {iteration}.\n\n\n")

                    if use_grad_clip:
                        if fp16_run:
                            grad_norm = torch.nn.utils.clip_grad_norm_(
                                amp.master_params(optimizer), grad_clip_thresh)
                        else:
                            grad_norm = torch.nn.utils.clip_grad_norm_(
                                model.parameters(), grad_clip_thresh)
                        if type(grad_norm) == torch.Tensor:
                            grad_norm = grad_norm.item()
                        is_overflow = math.isinf(grad_norm) or math.isnan(grad_norm)
                    else: is_overflow = False; grad_norm=0.00001

                    optimizer.step()
                    if not is_overflow and rank == 0:
                        # get current Loss Scale of first optimizer
                        loss_scale = amp._amp_state.loss_scalers[0]._loss_scale if fp16_run else 32768

                        if with_tensorboard:
                            if (iteration % 100000 == 0):
                                # plot distribution of parameters
                                for tag, value in model.named_parameters():
                                    tag = tag.replace('.', '/')
                                    logger.add_histogram(tag, value.data.cpu().numpy(), iteration)
                            logger.add_scalar('training_loss', reduced_loss, iteration)
                            logger.add_scalar('training_loss_samples', reduced_loss, iteration*batch_size)
                            if (iteration % 20 == 0):
                                logger.add_scalar('learning.rate', learning_rate, iteration)
                            if (iteration % 10 == 0):
                                logger.add_scalar('duration', ((time.time() - start_time_dekaiter)/10), iteration)

                        average_loss = rolling_sum.process(reduced_loss)
                        if (iteration % 10 == 0):
                            tqdm.write("{} {}:  {:.3f} {:.3f}  {:.3f} {:08.3F} {:.8f}LR ({:.8f} Effective)  {:.2f}s/iter {:.4f}s/item".format(time.strftime("%H:%M:%S"), iteration, reduced_loss, average_loss, best_MSE, round(grad_norm,3), learning_rate, min((grad_clip_thresh/grad_norm)*learning_rate,learning_rate), (time.time() - start_time_dekaiter)/10, ((time.time() - start_time_dekaiter)/10)/(batch_size*num_gpus)))
                            start_time_dekaiter = time.time()
                        else:
                            tqdm.write("{} {}:  {:.3f} {:.3f}  {:.3f} {:08.3F} {:.8f}LR ({:.8f} Effective) {}LS".format(time.strftime("%H:%M:%S"), iteration, reduced_loss, average_loss, best_MSE, round(grad_norm,3), learning_rate, min((grad_clip_thresh/grad_norm)*learning_rate,learning_rate), loss_scale))
                        start_time_iter = time.time()

                    if rank == 0 and (len(rolling_sum.values) > moving_average-2):
                        if (average_loss+best_model_margin) < best_model_loss:
                            checkpoint_path = os.path.join(output_directory, "best_model")
                            try:
                                save_checkpoint(model, optimizer, learning_rate, iteration, amp, scheduler, speaker_lookup,
                                            checkpoint_path)
                            except KeyboardInterrupt: # Avoid corrupting the model.
                                save_checkpoint(model, optimizer, learning_rate, iteration, amp, scheduler, speaker_lookup,
                                            checkpoint_path)
                            text_file = open((f"{checkpoint_path}.txt"), "w", encoding="utf-8")
                            text_file.write(str(average_loss)+"\n"+str(iteration))
                            text_file.close()
                            best_model_loss = average_loss #Only save the model if X better than the current loss.
                    if rank == 0 and iteration > 0 and ((iteration % iters_per_checkpoint == 0) or (os.path.exists(save_file_check_path))):
                        checkpoint_path = f"{output_directory}/waveglow_{iteration}"
                        save_checkpoint(model, optimizer, learning_rate, iteration, amp, scheduler, speaker_lookup,
                                        checkpoint_path)
                        if (os.path.exists(save_file_check_path)):
                            os.remove(save_file_check_path)

                    if (iteration % validation_interval == 0):
                        if rank == 0:
                            MSE, MAE = validate(model, loader_STFT, STFTs, logger, iteration, data_config['validation_files'], speaker_lookup, sigma, output_directory, data_config)
                            if scheduler:
                                MSE = torch.tensor(MSE, device='cuda')
                                if num_gpus > 1:
                                    broadcast(MSE, 0)
                                scheduler.step(MSE.item())
                                if MSE < best_MSE:
                                    checkpoint_path = os.path.join(output_directory, "best_val_model")
                                    try:
                                        save_checkpoint(model, optimizer, learning_rate, iteration, amp, scheduler, speaker_lookup,
                                                    checkpoint_path)
                                    except KeyboardInterrupt: # Avoid corrupting the model.
                                        save_checkpoint(model, optimizer, learning_rate, iteration, amp, scheduler, speaker_lookup,
                                                    checkpoint_path)
                                    text_file = open((f"{checkpoint_path}.txt"), "w", encoding="utf-8")
                                    text_file.write(str(MSE.item())+"\n"+str(iteration))
                                    text_file.close()
                                    best_MSE = MSE.item() #Only save the model if X better than the current loss.
                        else:
                            if scheduler:
                                MSE = torch.zeros(1, device='cuda')
                                broadcast(MSE, 0)
                                scheduler.step(MSE.item())
                        learning_rate = optimizer.param_groups[0]['lr'] #check actual learning rate (because I sometimes see learning_rate variable go out-of-sync with real LR)
                    iteration += 1
            training = False # exit the While loop

        except LossExplosion as ex: # print Exception and continue from checkpoint. (turns out it takes < 4 seconds to restart like this, fucking awesome)
            print(ex) # print Loss
            checkpoint_path = os.path.join(output_directory, "best_model")
            assert os.path.exists(checkpoint_path), "best_val_model must exist for automatic restarts"

            # clearing VRAM for load checkpoint
            audio = mel = speaker_ids = loss = None
            torch.cuda.empty_cache()

            model.eval()
            model, optimizer, iteration, scheduler = load_checkpoint(checkpoint_path, model, optimizer, scheduler, fp16_run)
            learning_rate = optimizer.param_groups[0]['lr']
            epoch_offset = max(0, int(iteration / len(train_loader)))
            model.train()
            iteration += 1
            pass # and continue training.

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str,
                        help='JSON file for configuration')
    parser.add_argument('-r', '--rank', type=int, default=0,
                        help='rank of process for distributed')
    parser.add_argument('-g', '--group_name', type=str, default='',
                        help='name of group for distributed')
    parser.add_argument('--detect_anomaly', action='store_true',
                        help='detects NaN/Infs in autograd backward pass and gives additional debug info.')
    args = parser.parse_args()

    if args.detect_anomaly: # checks backprop for NaN/Infs and outputs very useful stack-trace. Runs slowly while enabled.
        torch.autograd.set_detect_anomaly(True)
        print("Autograd Anomaly Detection Enabled!\n(Code will run slower but backward pass will output useful info if crashing or NaN/inf values)")

    # Parse configs.  Globals nicer in this case
    with open(args.config) as f:
        data = f.read()
    config = json.loads(data)
    train_config = config["train_config"]
    global data_config
    data_config = config["data_config"]
    if "n_mel_channels" in config["waveglow_config"].keys():
        data_config["n_mel_channels"] = config["waveglow_config"]["n_mel_channels"]
    if not 'iso226_empthasis' in data_config.keys():
        data_config["iso226_empthasis"] = False
    global dist_config
    dist_config = config["dist_config"]
    global waveglow_config
    waveglow_config = {
        **config["waveglow_config"],
        'win_length': data_config['win_length'],
        'hop_length': data_config['hop_length'],
        'preempthasis': data_config['preempthasis'],
        'n_mel_channels': data_config["n_mel_channels"],
        'sampling_rate': data_config["sampling_rate"],
        'use_logvar_channels': data_config["use_logvar_channels"],
        'iso226_empthasis': data_config["iso226_empthasis"]
    }
    print(waveglow_config)

    num_gpus = torch.cuda.device_count()
    if num_gpus > 1:
        if args.group_name == '':
            print("WARNING: Multiple GPUs detected but no distributed group set")
            print("Only running 1 GPU.  Use distributed.py for multiple GPUs")
            num_gpus = 1

    if num_gpus == 1 and args.rank != 0:
        raise Exception("Doing single GPU training on rank > 0")

    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = True
    train(num_gpus, args.rank, args.group_name, **train_config)


# -*- coding: utf-8 -*-
"""
Created on Thu Feb  6 11:19:57 2020

@author: keelin
"""

from pathlib import Path

import numpy as np
from keras.models import load_model
from opencxr.algorithms.base_algorithm import BaseAlgorithm
from opencxr.algorithms.imagesorter.preprocess import preprocess_img
import wget
import os

"""
An algorithm to sort 2D images into the following categories:
 - CXR (PA)
 - CXR (AP)
 - CXR (lateral)
 - Non-CXR
In addition, the following information is provided for each image
 - Rotation (in degrees, based on expectation of an upright CXR image)
 - Lateral flip (for a lateral CXR image, based on expectation of spine to right of lateral image)
 - Inversion (based on expectation of CXR image with dark intensitities for air filled regions)
"""

class ImageSorterAlgorithm(BaseAlgorithm):
    def name(self):
        return "ImageSorterAlgorithm"

    def __init__(self):
        """
        load the model
        """
        path_to_model_file = (
            Path(__file__).parent.parent / "model_weights" / "image_sorter.hdf5"
        )
        path_to_model_resolved = str(path_to_model_file.resolve())

        file_url = "https://github.com/DIAGNijmegen/opencxr/raw/master/opencxr/algorithms/model_weights/image_sorter.hdf5"
        # if the file does not exist (it's not included in whl file) then download it from github
        if not os.path.isfile(path_to_model_resolved):
            print("First use of imagesorter model, downloading the weights......")
            os.makedirs(os.path.dirname(path_to_model_resolved), exist_ok=True)
            wget.download(file_url, path_to_model_resolved)
            if not os.path.isfile(path_to_model_resolved):
                print("Failed to download file from", file_url)
                print(
                    "Please check the URL is valid and the following location is writeable",
                    path_to_model_resolved,
                )
                return
        if (
            not os.stat(path_to_model_resolved).st_size > 150000000
        ):  # model file needs to be larger than 150MB or it was not pulled correctly
            print(
                "ImageSorter model file seems incorrect (size ",
                os.stat(path_to_model_resolved).st_size,
                ") downloading the weights......",
                path_to_model_resolved,
            )
            os.remove(
                path_to_model_resolved
            )  # first delete the file, because wget will not overwrite
            wget.download(file_url, path_to_model_resolved)

        self.model = load_model(path_to_model_resolved)

    def run(self, image):
        """
        Run the image sorter algorithm
        :param image: The input image in x,y axis order (as from utils file_io read_image())
        :return: # a dict something like:
        {'Type': 'PA',
           'Rotation': '0',
           'Inversion': 'No',
           'Lateral_Flip': 'No',
           'Type_Probs_PA_AP_lateral_notCXR': [0.99999976, 2.5101654e-08, 2.4382584e-07, 1.0590604e-08],
           'Rotation_Probs_0_90_180_270': [0.9999999, 2.7740466e-08, 2.2800064e-08, 3.7591672e-08],
           'Inversion_Probs_No_Yes': [0.9999968589511354, 3.1410489e-06],
           'Lateral_Flip_Probs_No_Yes': [0.9999986753330177, 1.324667e-06]}
        The first four keys give classifications for Type, Rotation, Inversion, Lateral_Flip.
        The second four keys provide probabilities of all possible classes for users that might need this

        Possible values for the first four keys are as follows:
        Type:  ['PA', 'AP', 'lateral', 'notCXR']
        Rotation: ['0', '90', '180', '270']
        Inversion: ['No', 'Yes']
        Lateral_Flip: ['No', 'Yes']
        """
        # first preprocessing:
        image = preprocess_img(image)

        # now get ready and call the model to do prediction
        image = image / 255.0
        image = np.stack((image, image, image), axis=2)
        pred = self.model.predict(np.expand_dims(image, 0))

        # binarize the model outputs
        im_type = np.argmax(pred[0])
        im_rot = np.argmax(pred[1])
        im_inv = 1 if pred[2] > 0.5 else 0
        im_flip = 1 if pred[3] > 0.5 else 0

        # set up return labels
        type_labels = ["PA", "AP", "lateral", "notCXR"]
        rotation_labels = ["0", "90", "180", "270"]
        inversion_labels = ["No", "Yes"]
        lateral_flip_labels = ["No", "Yes"]

        # return dict
        return {
            "Type": type_labels[im_type],
            "Rotation": rotation_labels[im_rot],
            "Inversion": inversion_labels[im_inv],
            "Lateral_Flip": lateral_flip_labels[im_flip],
            "Type_Probs_PA_AP_lateral_notCXR": list(pred[0][0]),
            "Rotation_Probs_0_90_180_270": list(pred[1][0]),
            "Inversion_Probs_No_Yes": [1 - pred[2][0][0], pred[2][0][0]],
            "Lateral_Flip_Probs_No_Yes": [1 - pred[3][0][0], pred[3][0][0]],
        }


import subprocess as sb
import sys
import os
import numpy as np
import pandas as pd
from scipy.spatial.distance import cdist
import torch
from wmw.playlist import Playlist
from wmw.lstm_generator import LSTMGenerator
from dotenv import load_dotenv, find_dotenv

# Spotify API
sb.call([sys.executable, "-m", "pip", "install", 'spotipy'])
import spotipy
import spotipy.util as util
from spotipy.oauth2 import SpotifyClientCredentials
from spotipy.client import SpotifyException

# Joblib
sb.call([sys.executable, "-m", "pip", "install", 'joblib'])
import joblib

# find .env automagically by walking up directories until it's found
dotenv_path = find_dotenv()
load_dotenv(dotenv_path)

# ENV variables
MODEL_DIR = os.environ.get("MODEL_DIR")
MODEL_CHOICE = os.environ.get("MODEL_CHOICE")
SCALER_CHOICE = os.environ.get("SCALER_CHOICE")
PCA_CHOICE = os.environ.get("PCA_CHOICE")

# TODO: Contain recommendations pool to Best Of and add
# Deep Work as pool for intro tracks
class GenericPlaylist(Playlist):
    def __init__(self, track_pool, spotify_auth, username, playlist_id, model_type="LSTM"):
        """ Initiates pool of historic tracks, spotify api authentication and
            model of choice.
        """
        self.recommended_track_ids = pd.DataFrame()
        self.playlist_id = playlist_id
        self.track_pool = track_pool
        self.spotify_auth = spotify_auth
        self.username = username

        # Feature set
        self.feature_list = ['danceability', 'energy',
                             'loudness', 'speechiness',
                             'acousticness', 'instrumentalness',
                             'liveness', 'valence', 'tempo']

        # Setup feature and model pipeline
        self.std_scaler = joblib.load(os.path.join(MODEL_DIR, SCALER_CHOICE))
        self.dim_red = joblib.load(os.path.join(MODEL_DIR, PCA_CHOICE))
        self.model = self.setup_model(MODEL_DIR, MODEL_CHOICE)

        # Initiate pipeline
        self.intro_track = self.get_first_track()
        self.new_playlist = self.generate_playlist(self.model, self.intro_track)

    def setup_model(self, dir, filename):
        """Setup model in eval mode."""
        try:
            model = LSTMGenerator(9, 12, 2, 9, 1)
            model.load_state_dict(torch.load(os.path.join(dir, filename)))
            model.eval()
        except Exception as exc:
            print(f'{exc} Failed to import model.')
        return model

    def get_first_track(self):
        """Get first track based on recommendations."""
        # Sample an intro song from track pool
        song = self.track_pool.sample(1).copy()

        # Gather a recommendation based on the intro track using spotify
        song_res = self.spotify_auth.recommendations(seed_tracks=song['id'].values, limit=1)

        # Gather track freatures from spotify result
        for r in song_res['tracks']:
            track = {'id': r['id'], 'artists': ([i['name'] for i in r['artists']],), 'name': r['name']}
            track_features = self.spotify_auth.audio_features(r['id'])[0]
            track.update(track_features)
            self.intro_track = pd.DataFrame(track, index=[0])

        # Prepare features
        self.intro_track[self.feature_list] = self.std_scaler.transform(self.intro_track[self.feature_list])

        return self.intro_track

    def harmonic_match(self, key, mode):
        """Given a key and mode, return compatible keys according to the harmonic wheel."""

        # Harmonic Mixing Wheel: Pitch Class
        # 1A 0 - A flat minor: 8 | 1B 0 - B major: 11
        # 2A 1 - E flat minor: 3 | 2B 1 - F-sharp major: 6
        # 3A 2 - B-flat minor: 10 | 3B 2 - D-flat major: 1
        # 4A 3 - F minor: 5 | 4B 3 - A-flat major: 8
        # 5A 4 - C minor: 0 | 5B 4 - E-flat major: 3
        # 6A 5 - G minor: 7 | 6B 5 - B-flat major: 10
        # 7A 6 - D minor: 2 | 7B 6 - F major: 5
        # 8A 7 - A minor: 9 | 8B 7 - C major: 0
        # 9A 8 - E minor: 4 | 9B 8 - G major: 7
        # 10A 9 - B minor: 11 | 10B 9 - D major: 2
        # 11A 10 - F sharp minor: 6 | 11B 10 - A major: 9
        # 12A 11 - D flat minor: 1 | 12B 11 - E major: 4

        # Harmonic keys mapped to corresponding pitch classes
        pitch_to_harmonic_keys = {0: [4, 7], 1: [11, 2], 2: [6, 9],
                                  3: [1, 4], 4: [8, 11], 5: [3, 6],
                                  6: [10, 1], 7: [5, 8], 8: [0, 3],
                                  9: [7, 10], 10: [2, 5], 11: [9, 0]}

        # Extract values and keys
        dv = np.array(list(pitch_to_harmonic_keys.values()))
        dk = np.array(list(pitch_to_harmonic_keys.keys()))

        # Harmonic key code corresponding song pitch class
        harm_key = dv[np.where(dk == key)][0][mode]

        # Harmonic key codes
        harmonic_keys = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])

        # Get compatible key codes
        comp_keycodes = np.take(harmonic_keys,
                                [harm_key - 1, harm_key, harm_key + 1],
                                mode='wrap')

        # Compatible keys
        comp_keys = [np.where(dv[:, mode] == i)[0][0].tolist() for i in comp_keycodes]

        # Compatible up/down key
        inner_outer_key = np.array([np.where(dv[:, int(not bool(mode))] == harm_key)[0][0]])
        comp_keys = np.concatenate([comp_keys, inner_outer_key])
        return comp_keys, inner_outer_key

    def get_recommendations(self, sample_size=15, recommend_size=15):
        """Obtain a dataframe of recommended tracks for a specific track position."""
        recommendations = pd.DataFrame()
        sample_df = self.track_pool.sample(sample_size).copy()

        # Iterate full catalog of WMW songs
        for _, row in sample_df.iterrows():
            song_search = row['track_name'].partition('-')[0] + ' ' + row['artist_name']
            try:
                song_res = self.spotify_auth.search(song_search, limit=1)['tracks']['items'][0]
                results = self.spotify_auth.recommendations(seed_tracks=[song_res['id']], limit=10)
                for r in results['tracks']:
                    track = {'id': r['id'], 'artists': [i['name'] for i in r['artists']], 'name': r['name']}
                    track_features = self.spotify_auth.audio_features(r['id'])[0]
                    track.update(track_features)
                    final_track = pd.DataFrame(track)
                    recommendations = recommendations.append(final_track, ignore_index=True)
            except Exception as exc:
                print(f'{exc}: sample song not searchable.')
        recommendations[self.feature_list] = self.std_scaler.transform(recommendations[self.feature_list])
        return recommendations

    def pick_optimal_track(self, candidates, target):
        """Select the track with the minimum Euclidean distance between the candidate tracks."""
        candidates_reduced = self.dim_red.transform(candidates[self.feature_list])
        target_reduced = self.dim_red.transform(target[0])
        next_track_id = np.argmin(cdist(target_reduced, candidates_reduced))
        next_track = candidates.iloc[next_track_id]
        return next_track

    def generate_playlist(self, model, intro_track, playlist_len=15):
        """Generate playlist."""

        # Prepare prediction list
        predicted = intro_track

        # Prepare initial input
        input = torch.FloatTensor(intro_track[self.feature_list].values)

        print("Track #1")
        print("Selected:", predicted['name'].values[0], '-', ', '.join(predicted['artists'].values[0]))

        hidden_state = model.init_hidden()

        for p in range(2, playlist_len + 1):
            print("Track #%s" % p)

            # Important stuff about the last track
            current_track = predicted.iloc[-1]
            current_key = current_track['key']
            current_mode = current_track['mode']

            # Generate output feature set of next song
            output, hidden_state = model(input, hidden_state)

            output = output.detach().numpy()

            # Get mode and key from last song and generate compatible keys and modes
            keys, outer_inner_key = self.harmonic_match(current_key, current_mode)

            # Get recommended tracks for current track position
            recommendations = self.get_recommendations()

            # Filter for compatible tracks according to key and mode (harmonic wheel)
            next_tracks_curr_mode = recommendations[
                (recommendations['key'].isin(keys[:3])) & (recommendations['mode'] == current_mode)
            ]

            next_tracks_change_mode = recommendations[
                (recommendations['key'] == keys[-1]) & (recommendations['mode'] == abs(int(not current_mode)))
            ]

            candidate_tracks = pd.concat([next_tracks_curr_mode, next_tracks_change_mode]).reset_index(drop=True)

            # Ensure no duplicates exist in the playlist
            candidate_tracks = candidate_tracks[~candidate_tracks['id'].isin(predicted['id'])]

            # Pick optimal track
            next_track = self.pick_optimal_track(candidate_tracks, output)

            print("Selected:", next_track['name'], '-', next_track['artists'])

            # Set new input vector for next song
            input = torch.FloatTensor([next_track[self.feature_list]])

            # Append next song to playlist
            predicted = predicted.append(next_track, ignore_index=True)

            print('-' * 20)

        return predicted

    def post_playlist(self):
        """Posts playlist.

        Raises:
            SpotifyException: When query fails to gather playlist metadata.
        """
        try:
            self.spotify_auth.user_playlist_replace_tracks(
                self.username,
                self.playlist_id,
                self.new_playlist['id'].values)
        except SpotifyException as e:
            print(f"{e}: Failed to post playlist to Spotify.")


"""
Contains the Font class and the RenderError thrown when a font cannot be rendered correctly.
"""
import os
from math import pi, sqrt
import pickle
import numpy as np
import tensorflow as tf
from PIL import Image, ImageDraw, ImageFont, ImageStat, ImageOps
import typecat.config as config
import typecat.font2img as f2i

from pkg_resources import resource_string

#mean stddev min max
MEAN = 0
STDDEV = 1
MIN = 2
MAX = 3

class RenderError(Exception):
    """ Thrown when a font cannot be rendered correctly and is displayed as boxes. """
    pass

class Font(object):
    """
    Stores cached information about fonts in memory.
    Attributes:
        name
        family
        path
        pilfont: PIL font object
        size: current font size
        category: serif, sans, display, etc.
        languages
        thicknesses: the thickness values at all points in the typeface
        thickness: average thickness
        thickness_variation: standard dev. of thicknesses
        slant
        width
        height
        display_categories: list of strings and values to display,
                            first idx in tuple is name and second is dict value
        ratio: height/width
        style: bold, thin, italic etc.
        ascent: space between lowercase height and uppercase height
        descent: length of tail on 'p', 'q', and others
        features: str of special characters supported
        TODO curve quantification
    """

    ALPHABET = "abcdefghijklmnopqrstuvwxyz"
    SYMMETRIC = "HITl"
    MONOSPACE = "Monospace"
    SERIF = "Serif"
    SANS = "Sans Serif"
    DISPLAY = "Display"
    HANDWRITING = "Handwriting"
    BLOCK = "Block"

    CATEGORIES = {
        "SERIF": 0,
        "HANDWRITING": 1,
        "DISPLAY": 2,
        "MONOSPACE": 3,
        "SANS": 4,
        0: "SERIF",
        1: "HANDWRITING",
        2: "DISPLAY",
        3: "MONOSPACE",
        4: "SANS",
    }

    DELTA_RAD = pi/18
    compare = {
        "slant": -1,
        "thickness": -1,
        "width": -1,
        "height": -1,
        "ascent": -1,
        "descent": -1,
    }

    FIVE_CLASS_MODEL = resource_string(__name__, 'models/five_class_graph.pb')

    search_str = ""
    search_categories = []
    fonts = dict()
    scale_values = dict()

    graph = None

    @staticmethod
    def setup_graph():
        """ Loads the tensorflow graph from the trained .pb file """
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(Font.FIVE_CLASS_MODEL)
        Font.graph = tf.import_graph_def(graph_def, name='')

    def __init__(self, arg1, arg2=None):
        # option 1: a path to a font to go find the details yourself
        if isinstance(arg1, str):
            self.path = arg1
            if arg2 is not None:
                self.size = arg2
            else:
                self.size = 50

            # We don't find category until the end of loading but some functions just do some
            # preliminary stuff with them before
            self.category = None
            self.pilfont = None

            self.open_path()
            self.extract_pil()
            self.extract_width()
            self.extract_thickness()
            self.extract_slant()

            print("Loaded {} from {}".format(self.name, self.path))

    def open_path(self):
        """ Opens a font from a given system path and checks to see if it renders correctly """
        self.pilfont = ImageFont.truetype(self.path, size=self.size)
        a_img = f2i.single_pil("A", self.pilfont, fore=1, back=0)[0]
        b_img = f2i.single_pil("B", self.pilfont, fore=1, back=0)[0]
        if a_img == b_img:
            print("Could not render {} properly.".format(self.name))
            raise RenderError("Could not render {} properly.".format(self.name))

    def extract_pil(self):
        """ Extracts all data already collected by PIL fonts """
        self.name = "{} {}".format(self.pilfont.font.family, self.pilfont.font.style)
        self.ascent = self.pilfont.font.ascent
        self.descent = self.pilfont.font.descent

    def training_img(self):
        """ Generates an image to pass through the neural net in the style of fontjoy """
        temp_size = self.size
        self.set_size(60)
        glyphs = ['Laseg', 'dhum', 'Hloiv']
        current_pos = [25, 25]
        spacing = 60
        img = Image.new("L", (226, 226), 0)
        for glyph in glyphs:
            temp_img = Image.new(mode="L", size=(226, 226), color=0)
            temp_draw = ImageDraw.Draw(temp_img)
            temp_draw.text((0, 0), glyph, font=self.pilfont, fill=255)
            bbox = temp_img.getbbox()
            if bbox is None:
                self.set_size(temp_size)
                return temp_img
            temp_crop = temp_img.crop(bbox)
            img.paste(temp_crop, tuple(current_pos))
            current_pos[1] = current_pos[1] + spacing
        self.set_size(temp_size)
        return ImageOps.invert(img)

    def extract_width(self):
        """ Infers width by averaging individual character widths. """
        totalwidth = 0
        totalheight = 0
        fullabc = Font.ALPHABET + Font.ALPHABET.upper()
        for ch_img in list(fullabc):
            img = f2i.single_pil(ch_img, self.pilfont, fore=1, back=0)[0]
            bbox = img.getbbox()
            totalwidth += bbox[2] - bbox[0]
            totalheight += bbox[3] - bbox[1]
        self.height = totalheight / len(fullabc)
        self.width = totalwidth / len(fullabc)
        self.ratio = self.height / self.width

    def extract_thickness(self):
        """
        Calculates thickness by averaging dark bits
        """
        slstr = Font.ALPHABET + Font.ALPHABET.upper()
        img = f2i.single_pil(slstr, self.pilfont)[0]
        self.thickness = ImageStat.Stat(img).mean[0]

    def getsize(self, *args):
        """
        Doesn't get fontsize! wrapper for pilfont.getsize which gets the length it would take to
        render a string
        """
        return self.pilfont.getsize(*args)

    @staticmethod
    def pil2tensor(img):
        img = img.convert("RGB")
        img = np.array(img.resize((299, 299))).reshape(1, 299, 299, 3)
        #return np.expand_dims(np_img, axis=0)
        return img

    @staticmethod
    def extract_category():
        """
        Runs the font file through Inception & the retrained layer to find its category
        This is static because it's much faster to generate one graph and run a bunch than
        regenerate for every font
        """

        if Font.graph is None:
            Font.setup_graph()

        with tf.Session(graph=Font.graph) as sess:
            softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
            for font in Font.fonts.values():
                img = font.training_img()
                predictions = sess.run(softmax_tensor,
                                       {'Mul:0': Font.pil2tensor(img)})
                predictions = np.squeeze(predictions)
                max_index = max(enumerate(predictions), key=lambda p: p[1])[0]
                font.category = Font.CATEGORIES[max_index]

    def extract_slant(self):
        """ Compute slant by getting mean slope of characters """
        meanslant = 0
        # should one use symmetrics or the whole alphabet? Who knows
        # also TODO, when we get linear reps of each letter use that instead
        slstr = Font.ALPHABET + Font.ALPHABET.upper()
        for glyph in list(Font.SYMMETRIC):
            x_points = []
            y_points = []
            img = f2i.single_pil(glyph, self.pilfont)[0]
            pixels = img.load()
            for y_val in range(img.size[1]):
                point = 0
                number = 0
                for x_val in range(img.size[0]):
                    if pixels[x_val, y_val] == 0:
                        point += x_val
                        number += 1
                if number != 0:
                    # THIS IS INTENTIONAL, DON'T CHANGE IT
                    # We want to get the slant off of the normal so we reverse x
                    # and y
                    x_points.append(y_val)
                    y_points.append(point/number)
            slant = np.polyfit(x_points, y_points, 1)[0]
            meanslant += slant
        self.slant = -meanslant / len(slstr)

    def set_size(self, size):
        """ set current font size """
        self.size = size
        self.pilfont = ImageFont.truetype(self.path, self.size)

    def save(self):
        """ pickles to the cache directory and returns the name of the file """
        pil = self.pilfont
        self.pilfont = None
        pickle.dump(self, open(config.CACHE_LOCATION + "/" +
                               str(self.name) + ".pickle", "wb"))
        self.pilfont = pil
        return self.name.replace(" ", "_")

    def __getstate__(self):
        return self.__dict__

    def __setstate__(self, dic):
        self.__dict__ = dic
        self.open_path()

    def __str__(self):
        return "Font {} at path {}".format(self.name, self.path)

    @staticmethod
    def extract_name(path):
        """ Extracts just the name from a font to check if it's loaded """
        pilfont = ImageFont.truetype(path)
        family = pilfont.font.family
        style = pilfont.font.style
        return "{} {}".format(family, style)

    @staticmethod
    def scale(feature, value):
        """ Scales a value to its standard dev. across the font set. """
        xprime = ((value - Font.scale_values[feature][MIN]) /
                  (Font.scale_values[feature][MAX] - Font.scale_values[feature][MIN])) * 10 - 5
        return xprime

    @staticmethod
    def scale_features():
        """
        Calculates the stddev, mean, min, and max of each feature.
        """
        for feature in Font.compare:
            population = []
            for value in Font.fonts.values():
                population.append(value.__dict__[feature])
            maximum = max(population)
            minimum = min(population)
            mean = np.mean(population)
            #TODO STddev scaling
            stddev = np.std(population)
            Font.scale_values[feature] = (mean, stddev, max(population), min(population))

    def dist(self):
        """
        Finds the Euclidean distance between the features of this font and the desired features
        such as thickness, category, name, etc.
        """
        total = 0
        if self.category in Font.search_categories:
            total += 100
        if Font.search_str != "" and Font.search_str.lower() in self.name.lower():
            total += 100
        for feature, value in Font.compare.items():
            if value == -1:
                continue
            if isinstance(value, (float, int)):
                total += (value - Font.scale(feature, self.__dict__[feature])) ** 2
            else:
                total += 0 if value == self.__dict__[feature] else 1
        return sqrt(total)

    def __lt__(self, other):
        return self.dist() < other.dist()

    def __gt__(self, other):
        return self.dist() > other.dist()

    def __eq__(self, other):
        return self.dist() == other.dist()

    def __ne__(self, other):
        return self.dist() != other.dist()


import string
from json import dumps
import numpy as np
import numba as nb
from worms import util
from worms.util import jit
from worms.vis import format_atom
from worms.filters.clash import _chain_bounds
import numba.types as nt
import homog
from worms.util import hash_str_to_int

def BBlock(entry, pdbfile, filehash, pose, ss, null_base_names, **kw):

   json = dumps(entry)
   chains = util.get_chain_bounds(pose)
   ss = np.frombuffer(ss.encode(), dtype="i1")
   ncac = util.get_bb_coords(pose)
   cb = util.get_cb_coords(pose)
   stubs = _ncac_to_stubs(ncac)
   com = np.mean(cb, axis=0)
   rg = np.sqrt(np.sum((cb - com)**2) / len(cb))

   assert len(pose) == len(ncac)
   assert len(pose) == len(stubs)
   assert len(pose) == len(ss)
   conn = _make_connections_array(entry["connections"], chains)
   if len(conn) is 0:
      print("bad conn info!", pdbfile)
      assert 0
      return None, pdbfile  # new, missing
   if ncac.shape[-1] is 4:
      ncac = ncac.astype(np.float64)
   elif ncac.shape[-1] is 3:
      tmp = np.ones((ncac.shape[0], 3, 4), dtype=np.float64)
      tmp[..., :3] = ncac
      ncac = tmp
   else:
      assert 0, "bad ncac"
   assert cb.shape == (len(pose), 4)

   if entry["base"] in null_base_names: basehash = 0
   else: basehash = hash_str_to_int(entry["base"])

   def npfb(s):
      if isinstance(s, list):
         s = "[" + ",".join(s) + "]"
      return np.frombuffer(s.encode(), dtype="i1")

   validated = entry["validated"]
   if validated in ("na", "NA"):
      validated = False

   bblock = _BBlock(
      json=npfb(json),
      connections=conn,
      file=npfb(entry["file"]),
      filehash=filehash,
      components=npfb(str(entry["components"])),
      protocol=npfb(entry["protocol"]),
      name=npfb(entry["name"]),
      classes=npfb(",".join(entry["class"])),
      validated=validated,
      _type=npfb(entry["type"]),
      base=npfb(entry["base"]),
      basehash=basehash,
      ncac=np.ascontiguousarray(ncac),
      cb=np.ascontiguousarray(cb),
      chains=np.array(chains, dtype="i4"),
      ss=ss,
      stubs=np.ascontiguousarray(stubs.astype("f8")),
      com=com,
      rg=rg,
   )

   return bblock

@nb.jitclass(
    (
        ("json", nt.int8[:]),
        ("connections", nt.int32[:, :]),
        ("file", nt.int8[:]),
        ("filehash", nt.int64),
        ("components", nt.int8[:]),
        ("protocol", nt.int8[:]),
        ("name", nt.int8[:]),
        ("classes", nt.int8[:]),
        ("validated", nt.boolean),
        ("_type", nt.int8[:]),
        ("base", nt.int8[:]),
        ("basehash", nt.int64),
        ("ncac", nt.float64[:, :, :]),
        ("cb", nt.float64[:, :]),
        ("chains", nt.int32[:, :]),
        ("ss", nt.int8[:]),
        ("stubs", nt.float64[:, :, :]),
        ("com", nt.float64[:]),
        ("rg", nt.float64),
    )
)  # yapf: disable
class _BBlock:
   def __init__(
         self,
         json,
         connections,
         file,
         filehash,
         components,
         protocol,
         name,
         classes,
         validated,
         _type,
         base,
         basehash,
         ncac,
         cb,
         chains,
         ss,
         stubs,
         com,
         rg,
   ):
      self.json = json
      self.connections = connections
      self.file = file
      self.filehash = filehash
      self.components = components
      self.protocol = protocol
      self.name = name
      self.classes = classes
      self.validated = validated
      self._type = _type
      self.base = base
      self.basehash = basehash
      self.ncac = ncac
      self.cb = cb
      self.chains = chains
      self.ss = ss
      self.stubs = stubs
      self.com = com
      self.rg = rg
      assert np.isnan(np.sum(self.ncac)) == False
      assert np.isnan(np.sum(self.cb)) == False
      assert np.isnan(np.sum(self.stubs)) == False
      assert np.isnan(np.sum(self.ss)) == False
      assert np.isnan(np.sum(self.chains)) == False

   @property
   def n_connections(self):
      return len(self.connections)

   def conn_dirn(self, i):
      return self.connections[i, 0]

   def conn_resids(self, i):
      return self.connections[i, 2:self.connections[i, 1]]

   @property
   def _state(self):
      # MUST stay same order as args to __init__!!!!!
      return (
         self.json,
         self.connections,
         self.file,
         self.filehash,
         self.components,
         self.protocol,
         self.name,
         self.classes,
         self.validated,
         self._type,
         self.base,
         self.basehash,
         self.ncac,
         self.cb,
         self.chains,
         self.ss,
         self.stubs,
         self.com,
         self.rg,
      )

def bblock_dump_pdb(
      out,
      bblock,
      dirn=(2, 2),
      splice=(-1, -1),
      join="splice",
      pos=np.eye(4),
      chain=0,
      anum=1,
      rnum=1,
):
   close = False
   if isinstance(out, str):
      out = open(out, "w")
      close = True

   chains0 = _chain_bounds(dirn, splice, bblock.chains, trim=0)
   if dirn[0] == 2 and dirn[1] == 2:
      chains = chains0
   else:
      sponly = _chain_bounds(dirn, splice, bblock.chains, trim=0, spliced_only=1)
      # chains will have insplice at first ops, outsplice at last pos
      # either could be none
      chains = list()
      chains.append(sponly[0] if dirn[0] < 2 else None)
      for c in chains0:
         if np.all(sponly[0] == c) or np.all(sponly[-1] == c):
            continue
         chains.append(c)
      if len(sponly) > 1 or chains[0] is None:
         chains.append(sponly[-1] if dirn[1] < 2 else None)

   aname = [" N  ", " CA ", " C  "]
   for ic, lbub in enumerate(chains):
      if lbub is None:
         continue
      for i in range(*lbub):
         for j in (0, 1, 2):
            xyz = pos @ bblock.ncac[i, j]
            out.write(
               format_atom(
                  atomi=anum,
                  atomn=aname[j],
                  resn="GLY",
                  chain=string.ascii_uppercase[chain],
                  resi=rnum,
                  x=xyz[0],
                  y=xyz[1],
                  z=xyz[2],
                  occ=1.0,
               ))
            anum += 1
         rnum += 1
      if join == "bb":
         continue
      if join == "splice" and ic + 1 == len(chains) and dirn[1] < 2:
         continue
      chain += 1
   if join == "bb":
      chain += 1
   if close:
      out.close()
   return chain, anum, rnum

def _ncac_to_stubs(ncac):
   """
        Vector const & center,
        Vector const & a,
        Vector const & b,
        Vector const & c
    )
    {
        Vector e1( a - b);
        e1.normalize();

        Vector e3( cross( e1, c - b ) );
        e3.normalize();

        Vector e2( cross( e3,e1) );
        M.col_x( e1 ).col_y( e2 ).col_z( e3 );
        v = center;
    """
   assert ncac.shape[1:] == (3, 4)
   stubs = np.zeros((len(ncac), 4, 4), dtype=np.float64)
   ca2n = (ncac[:, 0] - ncac[:, 1])[..., :3]
   ca2c = (ncac[:, 2] - ncac[:, 1])[..., :3]
   # tgt1 = ca2n + ca2c  # thought this might make
   # tgt2 = ca2n - ca2c  # n/c coords match better
   tgt1 = ca2n  # rosetta style
   tgt2 = ca2c  # seems better
   a = tgt1
   a /= np.linalg.norm(a, axis=-1)[:, None]
   c = np.cross(a, tgt2)
   c /= np.linalg.norm(c, axis=-1)[:, None]
   b = np.cross(c, a)
   assert np.allclose(np.sum(a * b, axis=-1), 0)
   assert np.allclose(np.sum(b * c, axis=-1), 0)
   assert np.allclose(np.sum(c * a, axis=-1), 0)
   assert np.allclose(np.linalg.norm(a, axis=-1), 1)
   assert np.allclose(np.linalg.norm(b, axis=-1), 1)
   assert np.allclose(np.linalg.norm(c, axis=-1), 1)
   stubs[:, :3, 0] = a
   stubs[:, :3, 1] = b
   stubs[:, :3, 2] = c
   stubs[:, :3, 3] = ncac[:, 1, :3]
   stubs[:, 3, 3] = 1
   return stubs

def bb_splice_res(bb, dirn):
   r = []
   for iconn in range(bb.n_connections):
      if bb.conn_dirn(iconn) == dirn:
         r.append(bb.conn_resids(iconn))
   return np.concatenate(r)

def bb_splice_res_N(bb):
   return splice_res(bb, 0)

def bb_splice_res_C(bb):
   return splice_res(bb, 1)

class BBlockWrap:
   def __init__(self, _bblock):
      self._bblock = _bblock

   def __setstate__(self, state):
      self._bblock = _BBlock(*state)

   def __getstate__(self):
      return self._bblock._state

@jit
def chain_of_ires(bb, ires):
   chain = np.empty_like(ires)
   for i, ir in enumerate(ires):
      if ir < 0:
         chain[i] = -1
      else:
         for c in range(len(bb.chains)):
            if bb.chains[c, 0] <= ir < bb.chains[c, 1]:
               chain[i] = c
   return chain

def _make_connections_array(entries, chain_bounds):
   try:
      reslists = [_get_connection_residues(e, chain_bounds) for e in entries]
   except Exception as e:
      print("make_connections_array failed on", entries, "error was:", e)
      return np.zeros((0, 0))

   order = np.argsort([x[0] for x in reslists])
   mx = max(len(x) for x in reslists)
   conn = np.zeros((len(reslists), mx + 2), "i4") - 1
   for i, iord in enumerate(order):
      conn[i, 0] = entries[iord]["direction"] == "C"
      conn[i, 1] = len(reslists[iord]) + 2
      conn[i, 2:conn[i, 1]] = reslists[iord]
   return conn

def _get_connection_residues(entry, chain_bounds):
   """should return sorted list of resi positions"""
   chain_bounds[-1][-1]
   r, c, d = entry["residues"], int(entry["chain"]), entry["direction"]
   if isinstance(r, str) and r.startswith("["):
      r = eval(r)
   if isinstance(r, list):
      try:
         return sorted(int(i) for i in r)
      except ValueError:
         assert len(r) is 1
         r = r[0]
   if r.count(","):
      c2, r = r.split(",")
      assert int(c2) == c
   b, e = r.split(":")
   if b == "-":
      b = 0
   if e == "-":
      e = -1
   nres = chain_bounds[c - 1][1] - chain_bounds[c - 1][0]
   b = int(b) if b else 0
   e = int(e) if e else nres
   if e < 0:
      e += nres
   return np.array(range(*chain_bounds[c - 1])[b:e], dtype="i4")

def bblock_components(bblock):
   return eval(bytes(bblock.components))

def bblock_str(bblock):
   return "\n".join([
      "jitclass BBlock(",
      "    file=" + str(bytes(bblock.file)),
      "    components=" + str(bblock_components(bblock)),
      "    protocol=" + str(bytes(bblock.protocol)),
      "    name=" + str(bytes(bblock.name)),
      "    classes=" + str(bytes(bblock.classes)),
      "    validated=" + str(bblock.validated),
      "    _type=" + str(bytes(bblock._type)),
      "    base=" + str(bytes(bblock.base)),
      "    ncac=array(shape=" + str(bblock.ncac.shape) + ", dtype=" + str(bblock.ncac.dtype) +
      ")",
      "    chains=" + str(bblock.chains),
      "    ss=array(shape=" + str(bblock.ss.shape) + ", dtype=" + str(bblock.ss.dtype) + ")",
      "    stubs=array(shape=" + str(bblock.stubs.shape) + ", dtype=" +
      str(bblock.connections.dtype) + ")",
      "    connectionsZ=array(shape=" + str(bblock.connections.shape) + ", dtype=" +
      str(bblock.connections.dtype) + ")",
      ")",
   ])


"""Wrapping PaccMann model."""
import os
import json
import zipfile
import numpy as np
import tensorflow as tf
from copy import deepcopy
from paccmann.models import paccmann_model_fn, MODEL_SPECIFICATION_FACTORY
from .smiles import process_smiles, get_smiles_language
from ..core import TextModel, get_model_file
from ...core import Task

MODEL_PARAMS_JSON = 'model_params.json'
MODEL_CHECKPOINT = 'model.ckpt-375000'
NUMBER_OF_GENES = 2128
SMILES_LENGTH = 155

class CachedGraphPaccMannPredictor(object):
    """
    A paccmann predictor that caches the graph avoiding reloading it.

    Inspired by:
    https://github.com/marcsto/rl/blob/master/src/fast_predict2.py
    """

    def __init__(
        self, estimator, input_fn,
        batch_size, checkpoint_path=None
    ):
        """
        Initialize CachedGraphPaccMannPredictor.

        Arguments:
            estimator (tf.estimator.Estimator): a PaccMann estimator.
            input_fn (Callable): an input function accepting a generator.
            batch_size (int): size of the batch supported by the estimator.
            checkpoint_path (str): path to the checkpoint to load.
                Defaults to loading the latest.
        """
        self.estimator = estimator
        self.first_run = True
        self.closed = False
        self.input_fn = input_fn
        self.batch_size = batch_size
        self.checkpoint_path = checkpoint_path

    def _create_generator(self):
        """
        Create a generator of examples.

        Returns:
            a generator of example processed in self.input_fn.
        """
        while not self.closed:
            for example in self.next_examples:
                yield example

    def predict(self, examples):
        """
        Predict on the given examples.

        Arguments:
            examples (Iterable): an iterable of examples.

        Returns:
            a np.array representing the logits [non-effective, effective].
        """
        # make sure we don't alter the input
        examples = deepcopy(examples)
        if isinstance(examples, np.ndarray) and len(examples.shape) > 1:
            examples = [row for row in examples]
        if not isinstance(examples, list):
            examples = [examples]
        number_of_examples = len(examples)
        iterations = number_of_examples // self.batch_size
        remainder = number_of_examples % self.batch_size
        # handle remainder in samples
        if remainder > 0:
            examples += [examples[-1]]*(self.batch_size - remainder)
        self.next_examples = examples
        if self.first_run:
            self.predictions = self.estimator.predict(
                input_fn=self.input_fn(self._create_generator()),
                yield_single_examples=False,
                checkpoint_path=self.checkpoint_path
            )
            self.first_run = False
        results = []
        for _ in range(iterations):
            prediction = next(self.predictions)
            results.extend(
                [value, 1 - value]
                for value in prediction['IC50']
            )
        # get rid of the optional remainder
        if remainder > 0:
            prediction = next(self.predictions)
            results.extend(
                [value, 1 - value]
                for value in prediction['IC50'][:remainder]
            )
        return np.array(results)

    def close(self):
        """
        Close the predictor.
        """
        self.closed = True
        try:
            next(self.predictions)
        except:
            pass

def generator_fn_to_dataset(
    generator_fn,
    number_of_genes, smiles_length,
    batch_size
):
    """
    Get a tf.data.Dataset from a generator function.

    Arguments:
        generator_fn (Callable): a generator function.
        number_of_genes (int): number of the gene considered.
        smiles_length (int): maximum SMILES length.
        batch_size (int): size of the batch.
    Returns:
        An input function.
    """
    return tf.data.Dataset.from_generator(
        generator_fn,
        output_types={
            'selected_genes_20': tf.float32,
            'smiles_atom_tokens': tf.int64
        },
        output_shapes={
            'selected_genes_20': number_of_genes,
            'smiles_atom_tokens': smiles_length
        }
    ).batch(batch_size)

def paccmann_smiles_input_fn(
    generator, cell_line,
    number_of_genes, smiles_length,
    batch_size
):
    """
    Input for SMILES data.

    Arguments:
        generator (generator): a generator providing SMILES.
        cell_line (np.array): an array containing cell line gene expression.
        number_of_genes (int): number of the gene considered.
        smiles_length (int): maximum SMILES length.
        batch_size (int): size of the batch.
    Returns:
        an input function accepting a SMILES generator.
    """
    def generator_fn():
        for smiles in generator:
            yield {
                'selected_genes_20': cell_line.astype(np.float32),
                'smiles_atom_tokens': np.array(process_smiles(smiles))
            }
    return lambda: generator_fn_to_dataset(
        generator_fn,
        number_of_genes, smiles_length,
        batch_size
    )

def paccmann_cell_line_input_fn(
    generator, smiles,
    number_of_genes, smiles_length,
    batch_size
):
    """
    Input for cell line data.

    Arguments:
        generator (generator): a generator providing cell line data.
        smiles (str): a SMILES representing a molecule.
        number_of_genes (int): number of the gene considered.
        smiles_length (int): maximum SMILES length.
        batch_size (int): size of the batch.
    Returns:
        an input function accepting a cell line generator.
    """
    def generator_fn():
        for cell_line in generator:
            yield {
                'selected_genes_20': cell_line.astype(np.float32),
                'smiles_atom_tokens': np.array(process_smiles(smiles))
            }
    return lambda: generator_fn_to_dataset(
        generator_fn,
        number_of_genes, smiles_length,
        batch_size
    )

class PaccMann(TextModel):
    """Multimodal classification of drug sensitivity."""

    def __init__(
        self, filename='paccmann.zip',
        origin='https://ibm.box.com/shared/static/dy2x4cen1dsrc738uewmv1iccdawlqwd.zip',
        model_type='mca',
        model_params_json='model_params.json',
        model_checkpoint='model.ckpt-375000',
        number_of_genes=2128,
        smiles_length=155,
        *args, **kwargs
    ):
        """
        Initalize the Model.

        Arguments:
            filename (str): model zip.
            origin (str): url where the model is stored.
            model_type (str): multimodal encoder type.
            model_params_json (str): name of the json containing the
                parameters.
            model_checkpoints (str): name of the checkpoint.
            number_of_genes (int): number of the gene considered.
            smiles_length (int): maximum SMILES length.
            args (Iterable): list of arguments.
            kwargs (dict): list of key-value arguments.
        """
        super().__init__(
            Task.CLASSIFICATION, filename, origin, *args, **kwargs
        )
        # store initalization parameters
        self.model_type = model_type
        self.model_params_json = model_params_json
        self.model_checkpoint = model_checkpoint
        self.number_of_genes = number_of_genes
        self.smiles_length = smiles_length
        # make sure the model is present
        self.model_dir = os.path.join(
            os.path.dirname(self.model_path),
            'paccmann'
        )
        if not os.path.exists(self.model_dir):
            with zipfile.ZipFile(self.model_path) as zip_fp:
                zip_fp.extractall(os.path.dirname(self.model_path))
        # parse the parameters
        self.params = None
        with open(os.path.join(self.model_dir, self.model_params_json)) as fp:
            self.params = json.load(fp)
        self.batch_size = self.params['eval_batch_size']
        self.checkpoint_path = os.path.join(self.model_dir, self.model_checkpoint)
        # initialize the estimator
        self.estimator = tf.estimator.Estimator(
            model_fn=(
                lambda features, labels, mode, params: paccmann_model_fn(
                    features, labels, mode, params,
                    model_specification_fn=(
                        MODEL_SPECIFICATION_FACTORY[self.model_type]
                    )
                )
            ),
            model_dir=self.model_dir,
            params=self.params
        )
        # create the SMILES language
        self.language = get_smiles_language()
        # placeholder for a predictor
        self.predictor = None

    def predict(self, sample, **kwargs):
        """
        Run the model for inference on a given sample and with the provided
        parameters.

        Arguments:
            sample (object): an input sample for the model.
            kwargs (dict): list of key-value arguments.

        Returns:
            a prediction for the model on the given sample.
        """
        return self.predictor.predict(sample)

class PaccMannSmiles(PaccMann):
    """Multimodal classification of drug sensitivity for a given cell line."""

    def __init__(
        self, cell_line,
        filename='paccmann.zip',
        origin='https://ibm.box.com/shared/static/dy2x4cen1dsrc738uewmv1iccdawlqwd.zip',
        model_type='mca',
        model_params_json='model_params.json',
        model_checkpoint='model.ckpt-375000',
        number_of_genes=2128,
        smiles_length=155,
        *args, **kwargs
    ):
        """
        Initalize the Model.

        Arguments:
            cell_line (np.array): an array containing cell line gene expression.
            filename (str): model zip.
            origin (str): url where the model is stored.
            model_type (str): multimodal encoder type.
            model_params_json (str): name of the json containing the
                parameters.
            model_checkpoints (str): name of the checkpoint.
            number_of_genes (int): number of the gene considered.
            smiles_length (int): maximum SMILES length.
            args (Iterable): list of arguments.
            kwargs (dict): list of key-value arguments.
        """
        self.cell_line = cell_line
        super().__init__(
            filename, origin, model_type, model_params_json,
            model_checkpoint, number_of_genes, smiles_length,
            *args, **kwargs
        )
        self.input_fn = lambda generator: paccmann_smiles_input_fn(
            generator, self.cell_line,
            self.number_of_genes, self.smiles_length,
            self.batch_size
        )
        self.predictor = CachedGraphPaccMannPredictor(
            self.estimator, self.input_fn,
            self.batch_size, self.checkpoint_path
        )

class PaccMannCellLine(PaccMann):
    """Multimodal classification of drug sensitivity for a given SMILES."""

    def __init__(
        self, smiles,
        filename='paccmann.zip',
        origin='https://ibm.box.com/shared/static/dy2x4cen1dsrc738uewmv1iccdawlqwd.zip',
        model_type='mca',
        model_params_json='model_params.json',
        model_checkpoint='model.ckpt-375000',
        number_of_genes=2128,
        smiles_length=155,
        *args, **kwargs
    ):
        """
        Initalize the Model.

        Arguments:
            smiles (str): a SMILES representing a molecule.
            filename (str): model zip.
            origin (str): url where the model is stored.
            model_type (str): multimodal encoder type.
            model_params_json (str): name of the json containing the
                parameters.
            model_checkpoints (str): name of the checkpoint.
            number_of_genes (int): number of the gene considered.
            smiles_length (int): maximum SMILES length.
            args (Iterable): list of arguments.
            kwargs (dict): list of key-value arguments.
        """
        self.smiles = smiles
        super().__init__(
            filename, origin, model_type, model_params_json,
            model_checkpoint, number_of_genes, smiles_length,
            *args, **kwargs
        )
        self.input_fn = lambda generator: paccmann_cell_line_input_fn(
            generator, self.smiles,
            self.number_of_genes, self.smiles_length,
            self.batch_size
        )
        self.predictor = CachedGraphPaccMannPredictor(
            self.estimator, self.input_fn,
            self.batch_size, self.checkpoint_path
        )