# *****************************************************************************
#  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

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, criterion, optimizer_d, scheduler, fp16_run, stage, 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()}

        if stage >= 2:
            if 'criterion' in checkpoint_dict and checkpoint_dict['hifigan_config']['stage'] == 2:
                criterion.load_state_dict(checkpoint_dict['criterion'])
            if 'optimizer_d' in checkpoint_dict:
                optimizer_d.load_state_dict(checkpoint_dict['optimizer_d'])
        print(f"Loaded checkpoint path '{checkpoint_path}'")

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

def save_checkpoint(model, optimizer, criterion, optimizer_d, learning_rate, iteration, amp, scheduler, speaker_lookup, filepath):
    tqdm.write("Saving model and optimizer state at iteration {} to {}".format(
          iteration, filepath))

    state_dict = model.state_dict()
    saving_dict = {'model': state_dict,
        'iteration': iteration,
        'scheduler': scheduler.state_dict(),
        'learning_rate': learning_rate,
        'speaker_lookup': speaker_lookup,
        'hifigan_config': hifigan_config,
        }
    if criterion is not None:
        saving_dict['criterion'] = criterion.state_dict()
    if optimizer_d is not None:
        saving_dict['optimizer_d'] = optimizer_d.state_dict()
    if optimizer is not None:
        saving_dict['optimizer'] = optimizer.state_dict()
    if amp:
        saving_dict['amp'] = amp.state_dict()
    torch.save(saving_dict, filepath)
    tqdm.write("Model 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, trainset, logger, iteration, validation_files, speaker_lookup, output_directory, data_config, save_audio=True, max_length_s= 5., files_to_process=9):
    print("Validating... ", end="")

    model.eval()
    val_start_time = time.time()
    total = 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]

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

            timestr = time.strftime("%Y_%m_%d-%H_%M_%S")
            files_processed = 0
            for i, (audiopath, melpath, *remaining) in enumerate(audiopaths_and_melpaths):
                if files_processed >= files_to_process: # number of validation files to run.
                    break
                noisy_audio, audio = trainset.get_from_path(audiopath, segment_length=720000) # 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

                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

                noisy_audio = noisy_audio.half() if model_type == "half" else noisy_audio # for fp16 training

                pred_audio = model(noisy_audio.cuda().unsqueeze(0)).cpu().float().clamp(min=-0.999, max=0.999).squeeze(0)

                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, pred_audio.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))
                    if not os.path.exists(audio_path):# save ground truth
                        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")

                    audio_path = os.path.join(output_directory, "samples", str(iteration)+"-"+timestr+"_NT", os.path.basename(audiopath))
                    if not os.path.exists(audio_path):# save noisy ground truth
                        os.makedirs(os.path.join(output_directory, "samples", str(iteration)+"-"+timestr+"_NT"), exist_ok=True)
                        sf.write(audio_path, noisy_audio.squeeze().cpu().numpy(), data_config['sampling_rate'], "PCM_16")
                files_processed+=1
                samples_processed+=pred_audio.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)
    else:
        average_MSE = 1e3
        average_MAE = 1e3
        print("Average MSE: N/A", "Average MAE: N/A")

    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)

    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 get_optimizer(model, optimizer: str, fp16_run=True, optimizer_fused=True, learning_rate=1e-4, max_grad_norm=200):
    optimizer = optimizer.lower()
    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=max_grad_norm)
    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)

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

def print_params(model, name='model'):
    pytorch_total_params = sum(p.numel() for p in model.parameters())
    print(f"{pytorch_total_params:,} total parameters in {name}")
    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"{pytorch_total_params:,} trainable parameters.")

def train(num_gpus, rank, group_name, stage, output_directory, epochs, learning_rate, sigma, 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======

    from model import HiFiGAN, HiFiGANLoss
    criterion = HiFiGANLoss(**hifigan_config).cuda()
    model = HiFiGAN(**hifigan_config).cuda()
    #=====START: ADDED FOR DISTRIBUTED======
    if num_gpus > 1:
        model = apply_gradient_allreduce(model)
        if stage >= 2:
            criterion = apply_gradient_allreduce(criterion)
    #=====END:   ADDED FOR DISTRIBUTED======

    criterion, optimizer_d = get_optimizer(criterion, optimizer, fp16_run, optimizer_fused=True) if stage >= 2 else (criterion, None)
    model, optimizer = get_optimizer(model, optimizer, fp16_run, optimizer_fused=True)

    ## 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, criterion, optimizer_d, iteration, scheduler = load_checkpoint(checkpoint_path, model,
                                                      optimizer, criterion, optimizer_d, scheduler, fp16_run, stage, 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), 200)) # 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 = 9e9

    # 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)))

    print_params(model, name='generator')

    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)
                    # 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."
                        for param_group in optimizer.param_groups:
                            param_group['lr'] = learning_rate
                        if optimizer_d is not None:
                            for param_group in optimizer_d.param_groups:
                                param_group['lr'] = learning_rate*d_lr_scale
                    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(f"scheduler._last_lr = {scheduler._last_lr} scheduler.best = {scheduler.best}  |", end='')
                    model.zero_grad()
                    noisy_audio, gt_audio, speaker_ids = batch
                    noisy_audio = torch.autograd.Variable(noisy_audio.cuda(non_blocking=True))
                    gt_audio = torch.autograd.Variable(gt_audio.cuda(non_blocking=True))
                    speaker_ids = speaker_ids.cuda(non_blocking=True).long().squeeze(1)
                    pred_audio = model(noisy_audio)#, speaker_ids)

                    metrics = criterion(pred_audio, gt_audio, amp, model, optimizer, optimizer_d, num_gpus, use_grad_clip, grad_clip_thresh)

                    if not metrics['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)
                            for key, value in metrics.items():
                                if key not in ['is_overflow',]:
                                    logger.add_scalar(key, value, iteration)
                            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)

                        logged_loss = metrics['g_train_loss'] if stage >= 2 else metrics['train_loss']
                        grad_norm = metrics['grad_norm']
                        average_loss = rolling_sum.process(logged_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, logged_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, logged_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, criterion, optimizer_d, learning_rate, iteration, amp, scheduler, speaker_lookup, checkpoint_path)
                            except KeyboardInterrupt: # Avoid corrupting the model.
                                save_checkpoint(model, optimizer, criterion, optimizer_d, 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, criterion, optimizer_d, 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, trainset, logger, iteration, data_config['validation_files'], speaker_lookup, 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, criterion, optimizer_d, 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()
                        else:
                            if scheduler:
                                MSE = torch.zeros(1, device='cuda')
                                broadcast(MSE, 0)
                                scheduler.step(MSE.item())
                    iteration += 1
            training = False # exit the training 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_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"]
    global dist_config
    dist_config = config["dist_config"]
    global hifigan_config
    hifigan_config = config["hifigan_config"]
    print(hifigan_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, hifigan_config['stage'], **train_config)


# create predictions from the other stuff
"""
Go from proposals + scores to relationships.

pred-cls: No bbox regression, obj dist is exactly known
sg-cls : No bbox regression
sg-det : Bbox regression

in all cases we'll return:
boxes, objs, rels, pred_scores

"""

import numpy as np
import torch
from lib.pytorch_misc import unravel_index
from lib.fpn.box_utils import bbox_overlaps
# from ad3 import factor_graph as fg
from time import time

def ass_embed_dist(ass_embed):
    num = ass_embed.size()[0]
    size = (num, num, ass_embed.size()[1])
    ass_embed_a = ass_embed.unsqueeze(dim=1).expand(*size)
    ass_embed_b = ass_embed_a.permute(1, 0, 2)
    diff_push = (ass_embed_a - ass_embed_b)
    diff_push = torch.pow(diff_push, 2).sum(dim=2)
    diff_push_t = (diff_push - diff_push.min()) / (diff_push.max() - diff_push.min())
    return  diff_push_t.max() - diff_push_t

def filter_dets(ids, im_inds, boxes, obj_scores, obj_classes, adj_mat_rel=None, adj_mat_obj=None,
                rel_inds=None, pred_scores=None,
                nongt_box = False, with_adj_mat=False, with_gt_adj_mat=False, gt_adj_mat=None,
                alpha=0.5, feat=None, ext_feat=False, mode=None, beta=0.5,
                ass_embed_obj=None, ass_embed_rel=None):
    """
    Filters detections....
    when testing model, the batch size always equal to one.
    :param boxes: [num_box, topk, 4] if bbox regression else [num_box, 4]
    :param obj_scores: [num_box] probabilities for the scores
    :param obj_classes: [num_box] class labels for the topk
    :param adj_mat: [num_box, 65] adj matrix
    :param rel_inds: [num_rel, 2] TENSOR consisting of (im_ind0, im_ind1)
    :param pred_scores: [topk, topk, num_rel, num_predicates]
    :param use_nms: True if use NMS to filter dets.
    :return: boxes, objs, rels, pred_scores

    """
    if boxes.dim() != 2:
        raise ValueError("Boxes needs to be [num_box, 4] but its {}".format(boxes.size()))

    num_box = boxes.size(0)
    if nongt_box == False:
        assert obj_scores.size(0) == num_box

    assert obj_classes.size() == obj_scores.size()
    num_rel = rel_inds.size(0)
    assert rel_inds.size(1) == 2
    assert pred_scores.size(0) == num_rel

    obj_scores0 = obj_scores.data[rel_inds[:,0]]
    obj_scores1 = obj_scores.data[rel_inds[:,1]]

    pred_scores_max, pred_classes_argmax = pred_scores.data[:,1:].max(1)
    pred_classes_argmax = pred_classes_argmax + 1

    rel_scores_argmaxed = pred_scores_max * obj_scores0 * obj_scores1
    rel_scores_vs, rel_scores_idx = torch.sort(rel_scores_argmaxed.view(-1), dim=0, descending=True)

    if with_adj_mat:
        '''
        rel_adj_mat = torch.zeros_like(adj_mat)
        rel_adj_mat[rel_inds[:,0],rel_inds[:,1]] = 1.0
        adj_mat = adj_mat * rel_adj_mat
        adj_mat = adj_mat[:,: adj_mat.size(0)].data
        im_inds = im_inds.data
        bool_adj_mat = im_inds[:,None] == im_inds[None,:]
        bool_adj_mat = bool_adj_mat
        dual_ind = torch.arange(adj_mat.size(0))
        dual_ind = dual_ind[:,None] == dual_ind[None,:]
        dual_ind = dual_ind.type_as(adj_mat)
        dual_ind = 1.0 - dual_ind
        bool_adj_mat = bool_adj_mat.type_as(adj_mat)
        adj_mat = adj_mat * bool_adj_mat
        adj_mat = adj_mat * dual_ind
        '''
        if adj_mat_rel is not None:
            adj_mat_rel = adj_mat_rel.data.contiguous()
        else:
            adj_mat_rel = None

        if adj_mat_obj is not None:
            adj_mat_obj = adj_mat_obj.data.contiguous()
        else:
            adj_mat_obj = None
        #adj_mat = adj_mat[rel_inds[:,0],rel_inds[:,1]]

        if with_gt_adj_mat:
            rel_scores_argmaxed = pred_scores_max * obj_scores0 * obj_scores1 * gt_adj_mat.data.contiguous()
            adj_mat_obj = gt_adj_mat.data.contiguous()
            adj_mat_rel = gt_adj_mat.data.contiguous()
        else:
            if mode != 'sgdet':
                if adj_mat_obj is None and adj_mat_rel is not None:
                    rel_scores_argmaxed = ((pred_scores_max * obj_scores0 * obj_scores1) ** (1.0 - alpha)) * \
                                          ((adj_mat_rel) ** (alpha))
                elif adj_mat_rel is None and adj_mat_obj is not None:
                    rel_scores_argmaxed = ((pred_scores_max * obj_scores0 * obj_scores1) ** (1.0 - alpha)) * \
                                          ((adj_mat_obj) ** (alpha))
                elif adj_mat_rel is not None and adj_mat_obj is not None:
                    rel_scores_argmaxed = ((pred_scores_max * obj_scores0 * obj_scores1) ** (1.0 - alpha)) * \
                                          (((adj_mat_rel) * (adj_mat_obj)) ** (alpha))
                else:
                    rel_scores_argmaxed = ((pred_scores_max * obj_scores0 * obj_scores1) )

            else:
                if adj_mat_obj is not None:
                    '''
                    rel_scores_argmaxed = ((pred_scores_max * obj_scores0 * obj_scores1) ** (1.0 - alpha)) * \
                                          ((adj_mat_rel * adj_mat_obj) ** (alpha))
                    '''
                    rel_scores_argmaxed = ((pred_scores_max * obj_scores0 * obj_scores1) ** (1.0 - alpha)) * \
                                      (((adj_mat_rel**(beta)) * (adj_mat_obj**((1.0 - beta)))) ** (alpha))

                else:
                    rel_scores_argmaxed = ((pred_scores_max * obj_scores0 * obj_scores1) ** (1.0 - alpha)) * \
                                          ((adj_mat_rel) ** (alpha))
            '''
            if ass_embed_obj is not None and ass_embed_rel is not None:
                ass_embed_obj_dist = (ass_embed_dist(ass_embed_obj)[rel_inds[:,0],rel_inds[:,1]]).data.contiguous()
                ass_embed_rel_dist = (ass_embed_dist(ass_embed_rel)[rel_inds[:,0],rel_inds[:,1]]).data.contiguous()
                rel_scores_argmaxed = rel_scores_argmaxed * ass_embed_obj_dist * ass_embed_rel_dist
            '''
        rel_scores_vs, rel_scores_idx = torch.sort(rel_scores_argmaxed.view(-1), dim=0, descending=True)

    rels = rel_inds[rel_scores_idx].cpu().numpy()

    pred_scores_sorted = pred_scores[rel_scores_idx].data.cpu().numpy()
    obj_scores_np = obj_scores.data.cpu().numpy()
    objs_np = obj_classes.data.cpu().numpy()
    boxes_out = boxes.data.cpu().numpy()

    if adj_mat_rel is not None:
        adj_mat_rel_np = adj_mat_rel[rel_scores_idx].cpu().numpy()
    else:
        adj_mat_rel_np = None

    if adj_mat_obj is not None:
        adj_mat_obj_np = adj_mat_obj[rel_scores_idx].cpu().numpy()
    else:
        adj_mat_obj_np = None

    if gt_adj_mat is not None:
        gt_adj_mat_np = gt_adj_mat[rel_scores_idx].data.cpu().numpy()
    else:
        gt_adj_mat_np = None
    if ext_feat:
        feat_np = feat[rel_scores_idx].data.cpu().numpy()
        return ids, boxes_out, objs_np, obj_scores_np, rels, pred_scores_sorted, \
               adj_mat_rel_np, adj_mat_obj_np, gt_adj_mat_np, feat_np
    return boxes_out, objs_np, obj_scores_np, rels, pred_scores_sorted, \
           adj_mat_rel_np, adj_mat_obj_np, gt_adj_mat_np

# def _get_similar_boxes(boxes, obj_classes_topk, nms_thresh=0.3):
#     """
#     Assuming bg is NOT A LABEL.
#     :param boxes: [num_box, topk, 4] if bbox regression else [num_box, 4]
#     :param obj_classes: [num_box, topk] class labels
#     :return: num_box, topk, num_box, topk array containing similarities.
#     """
#     topk = obj_classes_topk.size(1)
#     num_box = boxes.size(0)
#
#     box_flat = boxes.view(-1, 4) if boxes.dim() == 3 else boxes[:, None].expand(
#         num_box, topk, 4).contiguous().view(-1, 4)
#     jax = bbox_overlaps(box_flat, box_flat).data > nms_thresh
#     # Filter out things that are not gonna compete.
#     classes_eq = obj_classes_topk.data.view(-1)[:, None] == obj_classes_topk.data.view(-1)[None, :]
#     jax &= classes_eq
#     boxes_are_similar = jax.view(num_box, topk, num_box, topk)
#     return boxes_are_similar.cpu().numpy().astype(np.bool)


import os
import json
import numpy as np
import scipy.sparse as sp

from src.model.linear_svm import LinearSVM
from src.model.random_forest import RandomForest
from src.metric.uar import get_UAR, get_post_probability, get_late_fusion_UAR
from src.utils.io import load_proc_baseline_feature, save_UAR_results
from src.utils.io import save_post_probability, load_post_probability
from src.utils.io import save_cv_results
from src.utils.preprocess import upsample
from src.utils.preprocess import k_fold_cv

'''
BASELINE CLASSIFICATION (py) PROVIDED BY AVEC2018

features        | computed level
--------        | --------------
MFCCs           | frame level
eGeMAPS         | turn level
DeepSpectrum    | activations in ALEXNET
BoAW            | window size (2s)
FAUs            | session level
BoVW            | window size (11s)
'''

class BaseLine():
    """
    Baseline system in BD classification, based on SVM/RF using LLDs and fusion
    ---
    Attributes
    -----------
    model_name: str
        model for BaseLine() instance, SVM or RF
    feature_name: str
        feature for BaseLine() instance, MFCC/eGeMAPS/Deep/BoAW/FAU/BoVW
    test: bool
        whether to test BaseLine() or not
    ----------------------------------------------------------------------
    Functions
    -----------
    run(): public
        main function
    run_[MFCC,eGeMAPS,DeepSpectrum,BoAW,AU,BoVW](): public
        run classifier on specified feature (single modality)
    run_fusion(): public
        run late fusion on a pair of specified features
    """
    def __init__(self, model_name, feature_name, test=False):
        # para model: determine the model in baseline system
        # para name: determine the feature in baseline system
        self.model_name = model_name
        self.feature_name = feature_name
        self.test = test
        print("\nbaseline system initialized, model %s feature %s" % (self.model_name, self.feature_name))

    def run(self):
        """main function of BaseLine() instance
        """
        if self.feature_name == 'FUSE':
            feature_name_1 = ''
            feature_name_2 = ''
            self.run_fusion(feature_name_1, feature_name_2)
        elif self.feature_name == 'MFCC':
            self.run_MFCC()
        elif self.feature_name == 'eGeMAPS':
            self.run_eGeMAPS()
        elif self.feature_name == 'Deep':
            self.run_DeepSpectrum()
        elif self.feature_name == 'BoAW':
            self.run_BoAW()
        elif self.feature_name == 'AU':
            self.run_AU()
        elif self.feature_name == 'BoVW':
            self.run_BoVW()

    def run_MFCC(self):
        """run classifier on MFCC feature (single modality)
        """
        print("\nbuilding a classifier on MFCC features (both frame-level and session-level)")
        X_train, y_train, train_inst, X_dev, y_dev, dev_inst = load_proc_baseline_feature('MFCC', verbose=True)

        if self.model_name == 'RF_cv':
            y_train, y_dev = np.ravel(y_train), np.ravel(y_dev)
            train_inst, dev_inst = np.ravel(train_inst), np.ravel(dev_inst)

            X = np.vstack((X_train, X_dev))
            y = np.hstack((y_train, y_dev))
            inst = np.hstack((train_inst, dev_inst))
            assert len(X) == len(y) == len(inst)
            cv_ids = k_fold_cv(len(X))
            cv_res = []
            for (ids_train, ids_dev) in cv_ids:
                X_train = X[ids_train]
                y_train = y[ids_train]
                X_dev = X[ids_dev]
                y_dev = y[ids_dev]
                dev_inst = inst[ids_dev]
                RF_MFCC = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
                RF_MFCC.run()
                y_pred_train, y_pred_dev = RF_MFCC.evaluate()
                _, session_res = get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=True)
                cv_res.append(session_res)
            save_cv_results(cv_res, self.model_name, self.feature_name, 'baseline')

        print("\nupsampling training data to address class imbalance")
        X_train, y_train, train_inst = upsample(X_train, y_train, train_inst)
        print("\nobtaining sparse matrix for better classification")
        # X_train = sp.csr_matrix(np.vstack((X_train, X_dev)))
        # X_dev = sp.csr_matrix(X_dev)
        # y_train = np.hstack((y_train, y_dev))
        X_train, X_dev = sp.csr_matrix(X_train), sp.csr_matrix(X_dev)

        if self.model_name == 'SVM':
            SVM_MFCC = LinearSVM(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            SVM_MFCC.run()
            y_pred_train, y_pred_dev = SVM_MFCC.evaluate()
        elif self.model_name == 'RF':
            RF_MFCC = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            RF_MFCC.run()
            y_pred_train, y_pred_dev = RF_MFCC.evaluate()

        get_UAR(y_pred_train, y_train, train_inst, self.model_name, self.feature_name, 'baseline', baseline=True, train_set=True, test=self.test)
        get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=self.test)
        if not self.test:
            get_post_probability(y_pred_dev, y_dev, dev_inst, np.array([]), self.model_name, self.feature_name)

    def run_eGeMAPS(self):
        """run classifier on eGeMAPS feature (single modality)
        """
        print("\nbuilding a classifier on eGeMAPS features (both frame-level and session-level)")
        X_train, y_train, train_inst, X_dev, y_dev, dev_inst = load_proc_baseline_feature('eGeMAPS', verbose=True)

        if self.model_name == 'RF_cv':
            X = np.vstack((X_train, X_dev))
            y = np.hstack((y_train, y_dev))
            assert len(X) == len(y)
            cv_ids = k_fold_cv(len(X))
            cv_res = []
            for (ids_train, ids_dev) in cv_ids:
                X_train = X[ids_train]
                y_train = y[ids_train]
                X_dev = X[ids_dev]
                y_dev = y[ids_dev]
                RF_MFCC = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
                RF_MFCC.run()
                y_pred_train, y_pred_dev = RF_MFCC.evaluate()
                _, session_res = get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=True)
                cv_res.append(session_res)
            save_cv_results(cv_res, self.model_name, self.feature_name, 'baseline')

        print("\nupsampling training data to address class imbalance")
        X_train, y_train, train_inst = upsample(X_train, y_train, train_inst)
        print("\nobtaining sparse matrix for better classification")
        # X_train = sp.csr_matrix(np.vstack((X_train, X_dev)))
        # X_dev = sp.csr_matrix(X_dev)
        # y_train = np.hstack((y_train, y_dev))
        X_train, X_dev = sp.csr_matrix(X_train), sp.csr_matrix(X_dev)

        if self.model_name == 'SVM':
            SVM_eGeMAPS = LinearSVM(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            SVM_eGeMAPS.run()
            y_pred_train, y_pred_dev = SVM_eGeMAPS.evaluate()
        elif self.model_name == 'RF':
            RF_eGeMAPS = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            RF_eGeMAPS.run()
            y_pred_train, y_pred_dev = RF_eGeMAPS.evaluate()

        get_UAR(y_pred_train, y_train, train_inst, self.model_name, self.feature_name, 'baseline', baseline=True, train_set=True, test=self.test)
        get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=self.test)
        if not self.test:
            get_post_probability(y_pred_dev, y_dev, dev_inst, np.array([]), self.model_name, self.feature_name)

    def run_DeepSpectrum(self):
        """run classifier on DeepSpectrum feature (single modality)
        """
        print("\nbuilding a classifier on Deep features (both frame-level and session-level)")
        X_train, y_train, train_inst, X_dev, y_dev, dev_inst = load_proc_baseline_feature('Deep', verbose=True)

        if self.model_name == 'RF_cv':
            X = np.vstack((X_train, X_dev))
            y = np.hstack((y_train, y_dev))
            assert len(X) == len(y)
            cv_ids = k_fold_cv(len(X))
            cv_res = []
            for (ids_train, ids_dev) in cv_ids:
                X_train = X[ids_train]
                y_train = y[ids_train]
                X_dev = X[ids_dev]
                y_dev = y[ids_dev]
                RF_MFCC = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
                RF_MFCC.run()
                y_pred_train, y_pred_dev = RF_MFCC.evaluate()
                _, session_res = get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=True)
                cv_res.append(session_res)
            save_cv_results(cv_res, self.model_name, self.feature_name, 'baseline')

        print("\nupsampling training data to address class imbalance")
        X_train, y_train, train_inst = upsample(X_train, y_train, train_inst)
        print("\nobtaining sparse matrix for better classification")
        # X_train = sp.csr_matrix(np.vstack((X_train, X_dev)))
        # X_dev = sp.csr_matrix(X_dev)
        # y_train = np.hstack((y_train, y_dev))
        X_train, X_dev = sp.csr_matrix(X_train), sp.csr_matrix(X_dev)

        if self.model_name == 'SVM':
            SVM_Deep = LinearSVM(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            SVM_Deep.run()
            y_pred_train, y_pred_dev = SVM_Deep.evaluate()
        elif self.model_name == 'RF':
            RF_Deep = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            RF_Deep.run()
            y_pred_train, y_pred_dev = RF_Deep.evaluate()

        get_UAR(y_pred_train, y_train, train_inst, self.model_name, self.feature_name, 'baseline', baseline=True, train_set=True, test=self.test)
        get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=self.test)
        if not self.test:
            get_post_probability(y_pred_dev, y_dev, dev_inst, np.array([]), self.model_name, self.feature_name)

    def run_BoAW(self):
        """run classifier on BoAW feature (single modality)
        """
        print("\nbuilding a classifier on BoAW features (both frame-level and session-level)")
        X_train, y_train, train_inst, X_dev, y_dev, dev_inst = load_proc_baseline_feature('BoAW', verbose=True)

        if self.model_name == 'RF_cv':
            y_train, y_dev = np.ravel(y_train), np.ravel(y_dev)

            X = np.vstack((X_train, X_dev))
            y = np.hstack((y_train, y_dev))
            assert len(X) == len(y)
            cv_ids = k_fold_cv(len(X))
            cv_res = []
            for (ids_train, ids_dev) in cv_ids:
                X_train = X[ids_train]
                y_train = y[ids_train]
                X_dev = X[ids_dev]
                y_dev = y[ids_dev]
                RF_MFCC = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
                RF_MFCC.run()
                y_pred_train, y_pred_dev = RF_MFCC.evaluate()
                _, session_res = get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=True)
                cv_res.append(session_res)
            save_cv_results(cv_res, self.model_name, self.feature_name, 'baseline')

        print("\nupsampling training data to address class imbalance")
        X_train, y_train, train_inst = upsample(X_train, y_train, train_inst)
        print("\nobtaining sparse matrix for better classification")
        # X_train = sp.csr_matrix(np.vstack((X_train, X_dev)))
        # X_dev = sp.csr_matrix(X_dev)
        # y_train = np.hstack((y_train, y_dev))
        X_train, X_dev = sp.csr_matrix(X_train), sp.csr_matrix(X_dev)

        if self.model_name == 'SVM':
            SVM_BoAW = LinearSVM(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            SVM_BoAW.run()
            y_pred_train, y_pred_dev = SVM_BoAW.evaluate()
        elif self.model_name == 'RF':
            RF_BoAW = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            RF_BoAW.run()
            y_pred_train, y_pred_dev = RF_BoAW.evaluate()

        get_UAR(y_pred_train, y_train, train_inst, self.model_name, self.feature_name, 'baseline', baseline=True, train_set=True, test=self.test)
        get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=self.test)
        if not self.test:
            get_post_probability(y_pred_dev, y_dev, dev_inst, np.array([]), self.model_name, self.feature_name)

    def run_AU(self):
        """run classifier on AU feature (single modality)
        """
        print("\nbuilding a classifier on AU features (already session-level)")
        X_train, y_train, _, X_dev, y_dev, _ = load_proc_baseline_feature('AU', verbose=True)

        if self.model_name == 'RF_cv':
            X = np.vstack((X_train, X_dev))
            y = np.hstack((y_train, y_dev))
            assert len(X) == len(y)
            cv_ids = k_fold_cv(len(X))
            cv_res = []
            for (ids_train, ids_dev) in cv_ids:
                X_train = X[ids_train]
                y_train = y[ids_train]
                X_dev = X[ids_dev]
                y_dev = y[ids_dev]
                RF_MFCC = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
                RF_MFCC.run()
                y_pred_train, y_pred_dev = RF_MFCC.evaluate()
                _, session_res = get_UAR(y_pred_dev, y_dev, np.array([]), self.model_name, self.feature_name, 'baseline', baseline=True, test=True)
                cv_res.append(session_res)
            save_cv_results(cv_res, self.model_name, self.feature_name, 'baseline')

        print("\nupsampling training data to address class imbalance")
        X_train, y_train, _ = upsample(X_train, y_train, np.array([]))
        print("\nobtaining sparse matrix for better classification")
        # X_train = sp.csr_matrix(np.vstack((X_train, X_dev)))
        # X_dev = sp.csr_matrix(X_dev)
        # y_train = np.hstack((y_train, y_dev))
        X_train, X_dev = sp.csr_matrix(X_train), sp.csr_matrix(X_dev)

        if self.model_name == 'SVM':
            SVM_AU = LinearSVM(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            SVM_AU.run()
            y_pred_train, y_pred_dev = SVM_AU.evaluate()
            session_prob = SVM_AU.get_session_probability()
        elif self.model_name == 'RF':
            RF_AU = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            RF_AU.run()
            y_pred_train, y_pred_dev = RF_AU.evaluate()
            session_prob = RF_AU.get_session_probability()

        get_UAR(y_pred_train, y_train, np.array([]), self.model_name, self.feature_name, 'baseline', baseline=True, train_set=True, test=self.test)
        get_UAR(y_pred_dev, y_dev, np.array([]), self.model_name, self.feature_name, 'baseline', baseline=True, test=self.test)

    def run_BoVW(self):
        """run classifier on BoVW feature (single modality)
        """
        print("\nbuilding a classifier on BoVW features (both frame-level and session-level)")
        X_train, y_train, train_inst, X_dev, y_dev, dev_inst = load_proc_baseline_feature('BoVW', verbose=True)

        if self.model_name == 'RF_cv':
            X = np.vstack((X_train, X_dev))
            y = np.hstack((y_train, y_dev))
            assert len(X) == len(y)
            cv_ids = k_fold_cv(len(X))
            cv_res = []
            for (ids_train, ids_dev) in cv_ids:
                X_train = X[ids_train]
                y_train = y[ids_train]
                X_dev = X[ids_dev]
                y_dev = y[ids_dev]
                RF_MFCC = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
                RF_MFCC.run()
                y_pred_train, y_pred_dev = RF_MFCC.evaluate()
                _, session_res = get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=True)
                cv_res.append(session_res)
            save_cv_results(cv_res, self.model_name, self.feature_name, 'baseline')

        print("\nupsampling training data to address class imbalance")
        X_train, y_train, train_inst = upsample(X_train, y_train, train_inst)
        print("\nobtaining sparse matrix for better classification")
        # X_train = sp.csr_matrix(np.vstack((X_train, X_dev)))
        # X_dev = sp.csr_matrix(X_dev)
        # y_train = np.hstack((y_train, y_dev))
        X_train, X_dev = sp.csr_matrix(X_train), sp.csr_matrix(X_dev)

        if self.model_name == 'SVM':
            SVM_BoVW = LinearSVM(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            SVM_BoVW.run()
            y_pred_train, y_pred_dev = SVM_BoVW.evaluate()
        elif self.model_name == 'RF':
            RF_BoVW = RandomForest(self.feature_name, X_train, y_train, X_dev, y_dev, baseline=True, test=self.test)
            RF_BoVW.run()
            y_pred_train, y_pred_dev = RF_BoVW.evaluate()

        get_UAR(y_pred_train, y_train, train_inst, self.model_name, self.feature_name, 'baseline', baseline=True, train_set=True, test=self.test)
        get_UAR(y_pred_dev, y_dev, dev_inst, self.model_name, self.feature_name, 'baseline', baseline=True, test=self.test)
        if not self.test:
            get_post_probability(y_pred_dev, y_dev, dev_inst, np.array([]), self.model_name, self.feature_name)

    def run_fusion(self, feature_name_1, feature_name_2):
        """run late fusion on a pair of specified features
        """
        get_late_fusion_UAR(self.model_name, feature_name_1, feature_name_2, baseline=True)


# Copyright 2021 The TensorFlow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import defaultdict
from copy import copy
from pathlib import Path
from typing import DefaultDict, Dict, List, MutableMapping, MutableSequence
from typing import Optional, Sequence, Union
import numpy as np
from tabulate import tabulate
import tensorflow as tf
from tensorflow.keras.optimizers import Optimizer
from tensorflow.keras.metrics import Metric
from tensorflow.keras.losses import Loss
from tqdm.auto import tqdm

from tensorflow_similarity.classification_metrics import ClassificationMetric
from tensorflow_similarity.classification_metrics import make_classification_metric  # noqa
from tensorflow_similarity.distances import Distance
from tensorflow_similarity.distances import distance_canonicalizer
from tensorflow_similarity.training_metrics import DistanceMetric
from tensorflow_similarity.evaluators.evaluator import Evaluator
from tensorflow_similarity.indexer import Indexer
from tensorflow_similarity.losses import MetricLoss
from tensorflow_similarity.matchers import ClassificationMatch
from tensorflow_similarity.retrieval_metrics import RetrievalMetric
from tensorflow_similarity.stores import Store
from tensorflow_similarity.search import Search
from tensorflow_similarity.types import FloatTensor, Lookup, IntTensor, Tensor
from tensorflow_similarity.types import PandasDataFrame, CalibrationResults

@tf.keras.utils.register_keras_serializable(package="Similarity")
class SimilarityModel(tf.keras.Model):
    """Specialized Keras.Model which implement the core features needed for
    metric learning. In particular, `SimilarityModel()` supports indexing,
    searching and saving the embeddings predicted by the network.

    All Similarity models classes derive from this class to benefits from those
    core features.
    """

    # @property
    # def _index(self):
    #     if not hasattr(self, '_index'):
    #         ValueError("Index doesn't exist: index data before quering it")
    #     return self._index

    # @index.setter
    # def _index(self, index):
    #     self._index: Indexer = index

    def compile(self,
                optimizer: Union[Optimizer, str, Dict, List] = 'rmsprop',  # noqa
                distance: Union[Distance, str] = 'auto',
                loss: Union[Loss, MetricLoss, str, Dict, List] = None,
                metrics: Union[Metric, DistanceMetric, str, Dict, List] = None,
                embedding_output: int = None,
                kv_store: Union[Store, str] = 'memory',
                search: Union[Search, str] = 'nmslib',
                evaluator: Union[Evaluator, str] = 'memory',
                stat_buffer_size: int = 1000,
                loss_weights: List = None,
                weighted_metrics: List = None,
                run_eagerly: bool = False,
                steps_per_execution: int = 1,
                **kwargs):
        """Configures the model for training.

        Args:

            optimizer: String (name of optimizer) or optimizer instance. See
            [tf.keras.optimizers](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers).

            loss: String (name of objective function), objective function,
            any `tensorflow_similarity.loss.*` instance or a
            `tf.keras.losses.Loss` instance. See the [Losses
            documentation](../losses.md) for a list of metric learning
            specifics loss offered by TensorFlow Similairy and
            [tf.keras.losses](https://www.tensorflow.org/api_docs/python/tf/keras/losses)
            for the losses available directly in TensorFlow.

            metrics: List of metrics to be evaluated by the model during
            training and testing. Each of those can be a string,
            a function or a [tensorflow_similairty.metrics.*](../metrics.md)
            instance. Note that the metrics used for some type of
            metric-learning such as distance learning (e.g via triplet loss)
            have a different prototype than the metrics used in
            standard models and you can't use the `tf.keras.metrics` for those
            type of learning.

            Additionally many distance metrics are computed based of the
            [Indexer()](../indexer.md) performance. E.g Matching Top 1
            accuracy. For technical and performance reasons, indexing data at
            each training batch to compute those is impractical so
            those metrics are computed at epoch end via
            the [EvalCallback](../callbacks.md)

            See [Evaluation Metrics](../eval_metrics.md) for a list of
            available metrics.

            For multi-output models you can specify different metrics for
            different outputs by passing a dictionary, such as
            `metrics={'similarity': 'min_neg_gap', 'other': ['accuracy',
            'mse']}`.  You can also pass a list (len = len(outputs)) of lists
            of metrics such as `metrics=[['min_neg_gap'], ['accuracy', 'mse']]`
            or `metrics=['min_neg_gap', ['accuracy', 'mse']]`. For outputs
            which are not related to metrics learning, you can use any of the
            standard `tf.keras.metrics`.

            loss_weights: Optional list or dictionary specifying scalar
            coefficients (Python floats) to weight the loss contributions of
            different model outputs. The loss value that will be minimized
            by the model will then be the *weighted sum* of all individual
            losses, weighted by the `loss_weights` coefficients.
            If a list, it is expected to have a 1:1 mapping to the model's
            outputs. If a dict, it is expected to map output names (strings)
            to scalar coefficients.

            weighted_metrics: List of metrics to be evaluated and weighted by
            sample_weight or class_weight during training and testing.

            run_eagerly: Bool. Defaults to `False`. If `True`, this `Model`'s
            logic will not be wrapped in a `tf.function`. Recommended to leave
            this as `None` unless your `Model` cannot be run inside a
            `tf.function`.

            steps_per_execution: Int. Defaults to 1. The number of batches to
            run during each `tf.function` call. Running multiple batches
            inside a single `tf.function` call can greatly improve performance
            on TPUs or small models with a large Python overhead.
            At most, one full epoch will be run each execution. If a number
            larger than the size of the epoch is passed,  the execution will be
            truncated to the size of the epoch.
            Note that if `steps_per_execution` is set to `N`,
            `Callback.on_batch_begin` and `Callback.on_batch_end` methods will
            only be called every `N` batches (i.e. before/after each
            `tf.function` execution).

    Raises:
        ValueError: In case of invalid arguments for
            `optimizer`, `loss` or `metrics`.
    """
        # Fetching the distance used from the first loss if auto
        if distance == 'auto':
            if isinstance(loss, list):
                metric_loss = loss[0]
            else:
                metric_loss = loss

            try:
                distance = metric_loss.distance
            except:  # noqa
                msg = "distance='auto' only works if the first loss is a\
                       metric loss"

                raise ValueError(msg)
            print("Distance metric automatically set to", distance,
                  "use the distance arg to override.")
        else:
            distance = distance_canonicalizer(distance)

        # check if we we need to set the embedding head
        num_outputs = len(self.output_names)
        if embedding_output and embedding_output > num_outputs:
            raise Exception("Embedding_output value exceed number of model "
                            "outputs")

        if not embedding_output and num_outputs > 1:
            print("Embedding output set to be model output 0",
                  "Use the embedding_output arg to override this")
            embedding_output = 0

        # fetch embedding size as some ANN libs requires it for init
        if num_outputs > 1:
            self.embedding_size = self.outputs[embedding_output].shape[1]
        else:
            self.embedding_size = self.outputs[0].shape[1]

        # init index
        self._index = Indexer(embedding_size=self.embedding_size,
                              distance=distance,
                              search=search,
                              kv_store=kv_store,
                              evaluator=evaluator,
                              embedding_output=embedding_output,
                              stat_buffer_size=stat_buffer_size)

        # call underlying keras method
        super().compile(optimizer=optimizer,
                        loss=loss,
                        metrics=metrics,
                        loss_weights=loss_weights,
                        weighted_metrics=weighted_metrics,
                        run_eagerly=run_eagerly,
                        steps_per_execution=steps_per_execution,
                        **kwargs)

    def index(self,
              x: Tensor,
              y: IntTensor = None,
              data: Optional[Tensor] = None,
              build: bool = True,
              verbose: int = 1):
        """Index data.

        Args:
            x: Samples to index.

            y: class ids associated with the data if any. Defaults to None.

            store_data: store the data associated with the samples in the key
            value store. Defaults to True.

            build: Rebuild the index after indexing. This is needed to make the
            new samples searchable. Set it to false to save processing time
            when calling indexing repeatidly without the need to search between
            the indexing requests. Defaults to True.

            verbose: Output indexing progress info. Defaults to 1.
        """

        if not self._index:
            raise Exception('You need to compile the model with a valid'
                            'distance to be able to use the indexing')
        if verbose:
            print('[Indexing %d points]' % len(x))
            print('|-Computing embeddings')
        predictions = self.predict(x)
        self._index.batch_add(predictions=predictions,
                              labels=y,
                              data=data,
                              build=build,
                              verbose=verbose)

    def lookup(self,
               x: Tensor,
               k: int = 5,
               verbose: int = 1) -> List[List[Lookup]]:
        """Find the k closest matches in the index for a set of samples.

        Args:
            x: Samples to match.

            k: Number of nearest neighboors to lookup. Defaults to 5.

            verbose: display progress. Default to 1.

        Returns
            list of list of k nearest neighboors:
            List[List[Lookup]]
        """
        predictions = self.predict(x)
        return self._index.batch_lookup(predictions=predictions,
                                        k=k,
                                        verbose=verbose)

    def single_lookup(self,
                      x: Tensor,
                      k: int = 5) -> List[Lookup]:
        """Find the k closest matches in the index for a given sample.

        Args:
            x: Sample to match.

            k: Number of nearest neighboors to lookup. Defaults to 5.

        Returns
            list of the k nearest neigboors info:
            List[Lookup]
        """
        x = tf.expand_dims(x, axis=0)
        prediction = self.predict(x)
        return self._index.single_lookup(prediction=prediction, k=k)

    def index_summary(self):
        "Display index info summary."
        self._index.print_stats()

    def calibrate(
            self,
            x: FloatTensor,
            y: IntTensor,
            thresholds_targets: MutableMapping[str, float] = {},
            k: int = 1,
            calibration_metric: Union[str, ClassificationMetric] = "f1",
            matcher: Union[str, ClassificationMatch] = 'match_nearest',
            extra_metrics: MutableSequence[Union[str, ClassificationMetric]] = ['precision', 'recall'],  # noqa
            rounding: int = 2,
            verbose: int = 1) -> CalibrationResults:
        """Calibrate model thresholds using a test dataset.

        TODO: more detailed explaination.

        Args:

            x: examples to use for the calibration.

            y: labels associated with the calibration examples.

            thresholds_targets: Dict of performance targets to (if possible)
            meet with respect to the `calibration_metric`.

            calibration_metric:
            [ClassificationMetric()](classification_metrics/overview.md) used
            to evaluate the performance of the index.

            k: How many neighboors to use during the calibration.
            Defaults to 1.

            matcher: {'match_nearest', 'match_majority_vote'} or
            ClassificationMatch object. Defines the classification matching,
            e.g., match_nearest will count a True Positive if the query_label
            is equal to the label of the nearest neighbor and the distance is
            less than or equal to the distance threshold. Defaults to
            'match_nearest'.

            extra_metrics: List of additional
            `tf.similarity.classification_metrics.ClassificationMetric()`
            to compute and report. Defaults to ['precision', 'recall'].

            rounding: Metric rounding. Default to 2 digits.

            verbose: Be verbose and display calibration results.
            Defaults to 1.

        Returns:
            CalibrationResults containing the thresholds and cutpoints Dicts.
        """

        # predict
        predictions = self.predict(x)

        # calibrate
        return self._index.calibrate(predictions=predictions,
                                     target_labels=y,
                                     thresholds_targets=thresholds_targets,
                                     k=k,
                                     calibration_metric=calibration_metric,
                                     matcher=matcher,
                                     extra_metrics=extra_metrics,
                                     rounding=rounding,
                                     verbose=verbose)

    def match(self,
              x: FloatTensor,
              cutpoint='optimal',
              no_match_label=-1,
              k=1,
              matcher: Union[str, ClassificationMatch] = 'match_nearest',
              verbose=0):
        """Match a set of examples against the calibrated index

        For the match function to work, the index must be calibrated using
        calibrate().

        Args:
            x: Batch of examples to be matched against the index.

            cutpoint: Which calibration threshold to use.
            Defaults to 'optimal' which is the optimal F1 threshold computed
            using calibrate().

            no_match_label: Which label value to assign when there is no
            match. Defaults to -1.

            k: How many neighboors to use during the calibration.
            Defaults to 1.

            matcher: {'match_nearest', 'match_majority_vote'} or
            ClassificationMatch object. Defines the classification matching,
            e.g., match_nearest will count a True Positive if the query_label
            is equal to the label of the nearest neighbor and the distance is
            less than or equal to the distance threshold.

            verbose. Be verbose. Defaults to 0.

        Returns:
            List of class ids that matches for each supplied example

        Notes:
            This function matches all the cutpoints at once internally as there
            is little performance downside to do so and allows to do the
            evaluation in a single go.

        """
        # basic checks
        if not self._index.is_calibrated:
            raise ValueError('Uncalibrated model: run model.calibration()')

        # get predictions
        predictions = self.predict(x)

        # matching
        matches = self._index.match(predictions,
                                    no_match_label=no_match_label,
                                    k=k,
                                    matcher=matcher,
                                    verbose=verbose)

        # select which matches to return
        if cutpoint == 'all':  # returns all the cutpoints for eval purpose.
            return matches
        else:  # normal match behavior - returns a specific cut point
            return matches[cutpoint]

    def evaluate_retrieval(
            self,
            x: Tensor,
            y: IntTensor,
            retrieval_metrics: Sequence[RetrievalMetric],  # noqa
            verbose: int = 1) -> Dict[str, np.ndarray]:
        """Evaluate the quality of the index against a test dataset.

        Args:
            x: Examples to be matched against the index.

            y: Label associated with the examples supplied.

            retrieval_metrics: List of
            [RetrievalMetric()](retrieval_metrics/overview.md) to compute.

            verbose (int, optional): Display results if set to 1 otherwise
            results are returned silently. Defaults to 1.

        Returns:
            Dictionary of metric results where keys are the metric names and
            values are the metrics values.
        """
        # get embeddings
        if verbose:
            print("|-Computing embeddings")
        predictions = self.predict(x)

        if verbose:
            print("|-Computing retrieval metrics")

        results = self._index.evaluate_retrieval(
                predictions=predictions,
                target_labels=y,
                retrieval_metrics=retrieval_metrics,
                verbose=verbose,
        )

        if verbose:
            table = zip(results.keys(), results.values())
            headers = ['metric', 'Value']
            print('\n [Summary]\n')
            print(tabulate(table, headers=headers))

        return results

    def evaluate_classification(
            self,
            x: Tensor,
            y: IntTensor,
            k: int = 1,
            extra_metrics: MutableSequence[Union[str, ClassificationMetric]] = ['precision', 'recall'],  # noqa
            matcher: Union[str, ClassificationMatch] = 'match_nearest',
            verbose: int = 1
            ) -> DefaultDict[str, Dict[str, Union[str, np.ndarray]]]:
        """Evaluate model classification matching on a given evaluation dataset.

        Args:
            x: Examples to be matched against the index.

            y: Label associated with the examples supplied.

            k: How many neighbors to use to perform the evaluation.
            Defaults to 1.

            extra_metrics: List of additional
            `tf.similarity.classification_metrics.ClassificationMetric()` to
            compute and report. Defaults to ['precision', 'recall'].

            matcher: {'match_nearest', 'match_majority_vote'} or
            ClassificationMatch object. Defines the classification matching,
            e.g., match_nearest will count a True Positive if the query_label
            is equal to the label of the nearest neighbor and the distance is
            less than or equal to the distance threshold.

            verbose (int, optional): Display results if set to 1 otherwise
            results are returned silently. Defaults to 1.

        Returns:
            Dictionary of (distance_metrics.md)[evaluation metrics]
        """
        # There is some code duplication in this function but that is the best
        # solution to keep the end-user API clean and doing inferences once.

        if not self._index.is_calibrated:
            raise ValueError('Uncalibrated model: run model.calibration()')
        cal_metric = self._index.get_calibration_metric()

        # get embeddings
        if verbose:
            print("|-Computing embeddings")
        predictions = self.predict(x)

        results: DefaultDict[str, Dict[str, Union[str, np.ndarray]]] = (
                defaultdict(dict))

        if verbose:
            pb = tqdm(total=len(self._index.cutpoints),
                      desc='Evaluating cutpoints')

        for cp_name, cp_data in self._index.cutpoints.items():
            # create a metric that match at the requested k and threshold
            distance_threshold = float(cp_data['distance'])
            metric = make_classification_metric(cal_metric.name)
            metrics = copy(extra_metrics)
            metrics.append(metric)

            res: Dict[str, Union[str, np.ndarray]] = {}
            res.update(
                self._index.evaluate_classification(
                    predictions,
                    y,
                    [distance_threshold],
                    metrics=metrics,
                    matcher=matcher,
                    k=k
                )
            )
            res['distance'] = tf.constant([distance_threshold])
            res['name'] = cp_name
            results[cp_name] = res
            if verbose:
                pb.update()

        if verbose:
            pb.close()

        if verbose:
            headers = ['name', cal_metric.name]
            for i in results['optimal'].keys():
                if i not in headers:
                    headers.append(str(i))
            rows = []
            for data in results.values():
                rows.append([data[v] for v in headers])
            print('\n [Summary]\n')
            print(tabulate(rows, headers=headers))

        return results

    def reset_index(self):
        "Reinitialize the index"
        self._index.reset()

    def index_size(self) -> int:
        "Return the index size"
        return self._index.size()

    def load_index(self, filepath: str):
        """Load Index data from a checkpoint and initialize underlying
        structure with the reloaded data.

        Args:
            path: Directory where the checkpoint is located.
            verbose: Be verbose. Defaults to 1.
        """

        index_path = Path(filepath) / "index"
        self._index = Indexer.load(index_path)

    def save_index(self, filepath, compression=True):
        """Save the index to disk

        Args:
            path: directory where to save the index
            compression: Store index data compressed. Defaults to True.
        """
        index_path = Path(filepath) / "index"
        self._index.save(index_path, compression=compression)

    def save(self,
             filepath: str,
             save_index: bool = True,
             compression: bool = True,
             overwrite: bool = True,
             include_optimizer: bool = True,
             signatures=None,
             options=None,
             save_traces: bool = True):
        """Save the model and the index.

        Args:
            filepath: where to save the model.

            save_index: Save the index content. Defaults to True.

            compression: Compress index data. Defaults to True.

            overwrite: Overwrite previous model. Defaults to True.

            include_optimizer: Save optimizer state. Defaults to True.

            signatures: Signatures to save with the model. Defaults to None.

            options: A `tf.saved_model.SaveOptions` to save with the model.
            Defaults to None.

            save_traces (optional): When enabled, the SavedModel will
            store the function traces for each layer. This can be disabled,
            so that only the configs of each layer are stored.
            Defaults to True. Disabling this will decrease serialization
            time and reduce file size, but it requires that all
            custom layers/models implement a get_config() method.
        """

        # save trace doesn't exist prior to 2.4 -- asking for it but not
        # using it

        # call underlying keras method to save the mode graph and its weights
        tf.keras.models.save_model(self,
                                   filepath,
                                   overwrite=overwrite,
                                   include_optimizer=include_optimizer,
                                   signatures=signatures,
                                   options=options,
                                   save_traces=save_traces)
        if hasattr(self, '_index') and self._index and save_index:
            self.save_index(filepath, compression=compression)
        else:
            print('Index not saved as save_index=False')

    def to_data_frame(self, num_items: int = 0) -> PandasDataFrame:
        """Export data as pandas dataframe

        Args:
            num_items (int, optional): Num items to export to the dataframe.
            Defaults to 0 (unlimited).

        Returns:
            pd.DataFrame: a pandas dataframe.
        """
        return self._index.to_data_frame(num_items=num_items)

    # We don't need from_config as the index is reloaded separatly.
    # this is kept as a reminder that it was looked into and decided to split
    # the index reloading instead of overloading this method.
    # @classmethod
    # def from_config(cls, config):
    #     return super().from_config(**config)


import argparse
import cv2
import numpy as np
from keras.models import Sequential, load_model
from segmentImage import segmentImage
width=608
height=608
confThreshold = 0.5
nmsThreshold = 0.4

parser = argparse.ArgumentParser(description='Object Detection using YOLO in OPENCV')
parser.add_argument('--image', help='Path to image file.')
args = parser.parse_args()
outputFile = args.image[:-4] + '_yolo_out_py.jpg'
outputPlate = args.image[:-4] + '_plate.jpg'
classesFile = "classes.names"

classes = None
with open(classesFile, 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')

predConf="darknet-yolov3.cfg"
weights='lapi.weights'

model = load_model('cnn_classifier.h5')

predictor = cv2.dnn.readNetFromDarknet(predConf, weights)
predictor.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
predictor.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

cap = cv2.VideoCapture(args.image)

def drawPred(classId, conf, left, top, right, bottom):
    # Draw a bounding box.
    #    cv.rectangle(frame, (left, top), (right, bottom), (255, 178, 50), 3)
    cv2.rectangle(img, (left, top), (right, bottom), (0, 255, 0), 3)
    plate = img[top:bottom, left:right]
    #cv2.imshow("CROPPED",plate)
    predResults=segmentImage(plate)
    cv2.imwrite(outputPlate, plate.astype(np.uint8))
    print(predResults)
    return predResults
    label = '%.2f' % conf

    # Get the label for the class name and its confidence
    '''
    if classes:
        assert(classId < len(classes))
        label = '%s:%s' % (classes[classId], label)

    #Display the label at the top of the bounding box
    labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
    top = max(top, labelSize[1])
    cv2.rectangle(img, (left, top - round(1.5*labelSize[1])), (left + round(1.5*labelSize[0]), top + baseLine), (0, 0, 255), cv2.FILLED)
    #cv.rectangle(frame, (left, top - round(1.5*labelSize[1])), (left + round(1.5*labelSize[0]), top + baseLine),    (255, 255, 255), cv.FILLED)
    cv2.putText(img, label, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 2)
    '''
def process(img,res):
    frameHeight=img.shape[0]
    frameWidth=img.shape[1]
    classIds = []
    confidences = []
    boxes = []
    for out in res:
        #print("out.shape : ", out.shape)
        for detection in out:
            #if detection[4]>0.001:
            scores = detection[5:]
            classId = np.argmax(scores)
            #if scores[classId]>confThreshold:
            confidence = scores[classId]
            #print(detection)
            #if detection[4]>confThreshold:
                #print(detection[4], " - ", scores[classId], " - th : ", confThreshold)
                #print(detection)
            if confidence > confThreshold:
                center_x = int(detection[0] * frameWidth)
                center_y = int(detection[1] * frameHeight)
                width = int(detection[2] * frameWidth)
                height = int(detection[3] * frameHeight)
                left = int(center_x - width / 2)
                top = int(center_y - height / 2)
                classIds.append(classId)
                confidences.append(float(confidence))

                boxes.append([left, top, width, height])

    indices = cv2.dnn.NMSBoxes(boxes, confidences, confThreshold, nmsThreshold)
    for i in indices:

        i = i[0]
        box = boxes[i]
        left = box[0]
        top = box[1]
        width = box[2]
        height = box[3]
        drawPred(classIds[i], confidences[i], left-4, top-4, left + width+4, top + height+4)

def OutputNames(predictor):
    # Get the names of all the layers in the network
    layersNames = predictor.getLayerNames()

    # Get the names of the output layers, i.e. the layers with unconnected outputs
    return [layersNames[i[0] - 1] for i in predictor.getUnconnectedOutLayers()]

ret, img= cap.read()
while(True):
    if (ret):

        blob = cv2.dnn.blobFromImage(img, 1/255, (width, height), [0,0,0], 1, crop=False)
        predictor.setInput(blob)

        res=predictor.forward(OutputNames(predictor))
        process(img,res)
        cv2.imshow("IMG", img)
        if (args.image):
            cv2.imwrite(outputFile, img.astype(np.uint8))

        k = cv2.waitKey(0)
        if (k == 27):
            cv2.destroyAllWindows()
            break