"""
NOTE: checkpoints can be loaded by python2 and pytorch 0.4.0
"""
import numpy as np
import os
import random
import shutil
import time
import warnings
from collections import defaultdict

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
from sacred import Experiment, Ingredient
from sacred.observers import FileStorageObserver

from model import E2E
from data_reader.dataset_v1 import SpoofDatsetSystemID, SpoofDatsetEval
from src.eval_metrics import compute_eer
from src.loss import FocalLoss

ex = Experiment('asvspoof19')
ex.observers.append(FileStorageObserver.create('snapshots'))

@ex.config
def my_config():
    model_params = {
        'MODEL_SELECT': 6,  # which model
        'NUM_SPOOF_CLASS': 2,  # 2,  # LA: 7 or PA: 10 x-class classification
        'FOCAL_GAMMA': None,  # gamma parameter for focal loss; if obj is not focal loss, set this to None
        'NUM_RESNET_BLOCK': 5,  # number of resnet blocks in ResNet
        'AFN_UPSAMPLE': 'Bilinear',  # upsampling method in AFNet: Conv or Bilinear
        'AFN_ACTIVATION': 'sigmoid',  # activation function in AFNet: sigmoid, softmaxF, softmaxT
        'NUM_HEADS': 3,  # number of heads for multi-head att in SAFNet
        'SAFN_HIDDEN': 10,  # hidden dim for SAFNet
        'SAFN_DIM': 'T',  # SAFNet attention dim: T or F
        'RNN_HIDDEN': 128,  # hidden dim for RNN
        'RNN_LAYERS': 4,  # number of hidden layers for RNN
        'RNN_BI': True,  # bidirecitonal/unidirectional for RNN
        'DROPOUT_R': 0.0,  # dropout rate
    }
    data_files = {  # training
        # ASV data settings
        # 'train_scp': '/home/vano/wrkdir/projects/antispoofing_speech/kaldi_feats/data/ASVspoof2019_PA_train_fbank/feats.scp', # VAD
        # 'train_utt2index': 'data_reader/utt2systemID/pa_train_utt2index_8',
        # 'dev_scp': '/home/vano/wrkdir/projects/antispoofing_speech/kaldi_feats/data/ASVspoof2019_PA_dev_fbank/feats.scp',
        # 'dev_utt2index': 'data_reader/utt2systemID/pa_dev_utt2index_8',
        # 'dev_utt2systemID': 'data_reader/utt2systemID/pa_dev_utt2systemID',

        # I2R users dataset
        'train_scp': '/home/vano/wrkdir/projects/antispoofing_speech/kaldi_feats/data/tasks_spec_cm/feats.scp',  # VAD
        'train_utt2index': 'data_reader/utt2systemID/tasks_utt2index_8',

        'dev_scp': '/home/vano/wrkdir/projects/antispoofing_speech/kaldi_feats/data/tasks_spec_cm/feats.scp',
        'dev_utt2index': 'data_reader/utt2systemID/tasks_utt2index_8',
        'dev_utt2systemID': 'data_reader/utt2systemID/tasks_utt2systemID',

        'eval_utt2systemID': 'data_reader/utt2systemID/pa_eval_utt2systemID',
        'scoring_dir': 'scoring/tasks_scores/',
    }

    leave_one_out = False  # leave one out during train and val
    eer_criteria = True  # train by dev acc or eer
    batch_size = 64  # 64
    test_batch_size = 1  # 64
    epochs = 10  # 20 for PA, 10 for LA
    start_epoch = 1
    n_warmup_steps = 1000
    log_interval = 100
    pretrained = '../pretrained/pa/senet34_py3'  # None
    pretrained_model_id = 30 #18  # PA: 18, PA_tune: 30 # LA: 20  # for forward pass
    class_labels = [  # for post analysis
        'bonafide', 'AB', 'AC',
        'BA', 'BB', 'BC', 'CA', 'CB', 'CC', 'AA',
    ]

best_acc1 = 0
best_eer = 10000

@ex.capture
def work(_run, pretrained, batch_size, test_batch_size, epochs, start_epoch, log_interval, n_warmup_steps, data_files,
         model_params, eer_criteria, leave_one_out):
    global best_acc1
    global best_eer

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    tmp = torch.tensor([2]).to(device)

    # model is trained for binary classification (for datalaoder)
    if model_params['NUM_SPOOF_CLASS'] == 2:
        binary_class = True
    else:
        binary_class = False

    # model is RNN
    if model_params['MODEL_SELECT'] == 4:
        use_rnn = True
    else:
        use_rnn = False

    if model_params['FOCAL_GAMMA']:
        print('training with focal loss')
        focal_obj = FocalLoss(gamma=model_params['FOCAL_GAMMA'])
    else:
        focal_obj = None

    kwargs = {'num_workers': 2, 'pin_memory': True} if use_cuda else {}

    # create model
    # cnx
    model = E2E(**model_params).to(device)
    model_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('===> Model total parameter: {}'.format(model_params))

    # Wrap model for multi-GPUs, if necessary
    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
        print('multi-gpu')
        model = nn.DataParallel(model).cuda()

    # define loss function (criterion) and optimizer
    # cnx
    optimizer = ScheduledOptim(
        torch.optim.Adam(
            filter(lambda p: p.requires_grad, model.parameters()),
            betas=(0.9, 0.98), eps=1e-09, weight_decay=1e-4, amsgrad=True),
        n_warmup_steps)

    # optionally resume from a checkpoint
    if pretrained:
        if os.path.isfile(pretrained):
            print("===> loading checkpoint '{}'".format(pretrained))
            checkpoint = torch.load(pretrained)
            start_epoch = checkpoint['epoch']
            if eer_criteria and 'best_eer' in checkpoint:
                best_eer = checkpoint['best_eer']
            else:
                best_acc1 = checkpoint['best_acc1']
                print("===> Best accuracy: %f" % best_acc1)
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            print("===> loaded checkpoint '{}' (epoch {})"
                  .format(pretrained, checkpoint['epoch']))
        else:
            print("===> no checkpoint found at '{}'".format(pretrained))

    cudnn.benchmark = True  # It enables benchmark mode in cudnn.

    # Data loading code
    train_data = SpoofDatsetSystemID(data_files['train_scp'], data_files['train_utt2index'],
                                     rnd_nslides=True,
                                     binary_class=binary_class,
                                     leave_one_out=leave_one_out)
    val_data = SpoofDatsetSystemID(data_files['dev_scp'], data_files['dev_utt2index'],
                                   binary_class=binary_class,
                                   leave_one_out=leave_one_out)
    train_loader = torch.utils.data.DataLoader(
        train_data, batch_size=batch_size, shuffle=True, **kwargs)
    val_loader = torch.utils.data.DataLoader(
        val_data, batch_size=test_batch_size, shuffle=False, **kwargs)
    if leave_one_out:
        eval_data = SpoofDatsetSystemID(data_files['dev_scp'], data_files['dev_utt2index'], binary_class, False)
        eval_loader = torch.utils.data.DataLoader(
            eval_data, batch_size=test_batch_size, shuffle=False, **kwargs)
    else:
        eval_loader = val_loader

    best_epoch = 0
    early_stopping, max_patience = 0, 5  # for early stopping
    for epoch in range(start_epoch, start_epoch + epochs):

        # train for one epoch
        train(train_loader, model, optimizer, epoch, device, log_interval, use_rnn, focal_obj)

        # evaluate on validation set
        acc1 = validate(val_loader, data_files['dev_utt2systemID'], model, device, log_interval,
                        use_rnn, eer_criteria, focal_obj)

        # remember best acc@1/eer and save checkpoint
        if eer_criteria:
            is_best = acc1 < best_eer
            best_eer = min(acc1, best_eer)
        else:
            is_best = acc1 > best_acc1
            best_acc1 = max(acc1, best_acc1)

        # adjust learning rate + early stopping
        if is_best:
            early_stopping = 0
            best_epoch = epoch + 1
        else:
            early_stopping += 1
            if epoch - best_epoch > 2:
                optimizer.increase_delta()
                best_epoch = epoch + 1
        if early_stopping == max_patience:
            break

        # save model
        if not is_best:
            continue
        if eer_criteria:
            save_checkpoint({
                'epoch': epoch,
                'state_dict': model.state_dict(),
                'best_eer': best_eer,
                'optimizer': optimizer.state_dict(),
            }, is_best, "snapshots/" + str(_run._id), str(epoch) + ('_%.3f' % acc1) + ".pth.tar")
        else:
            save_checkpoint({
                'epoch': epoch,
                'state_dict': model.state_dict(),
                'best_acc1': best_acc1,
                'optimizer': optimizer.state_dict(),
            }, is_best, "snapshots/" + str(_run._id), str(epoch) + ('_%.3f' % acc1) + ".pth.tar")

    # load best model
    best_model_pth = os.path.join('snapshots', str(_run._id), 'model_best.pth.tar')
    score_file_pth = os.path.join(data_files['scoring_dir'], str(_run._id) + '-scores.txt')
    print("===> loading best model for scoring: '{}'".format(best_model_pth))
    checkpoint = torch.load(best_model_pth)
    model.load_state_dict(checkpoint['state_dict'])

    # compute EER
    print("===> scoring file saved at: '{}'".format(score_file_pth))
    prediction(eval_loader, model, device, score_file_pth, data_files['dev_utt2systemID'], use_rnn, focal_obj)

@ex.capture
def post(_run, pretrained, test_batch_size, data_files, model_params, eer_criteria, class_labels):
    """ what are the classes that performed well, and the classes that did not perform well """

    use_cuda = torch.cuda.is_available()  # use cpu
    device = torch.device("cuda" if use_cuda else "cpu")
    tmp = torch.tensor([2]).to(device)

    # model is RNN
    if model_params['MODEL_SELECT'] == 4:
        use_rnn = True
    else:
        use_rnn = False

    kwargs = {'num_workers': 2, 'pin_memory': True} if use_cuda else {}

    # create model
    # cnx
    model = E2E(**model_params).to(device)
    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('===> Model total parameter: {}'.format(num_params))

    if pretrained:
        if os.path.isfile(pretrained):
            print("===> loading checkpoint '{}'".format(pretrained))
            checkpoint = torch.load(pretrained, map_location=lambda storage, loc: storage)  # load for cpu
            if eer_criteria:
                best_eer = checkpoint['best_eer']
            else:
                best_acc1 = checkpoint['best_acc1']
            model.load_state_dict(checkpoint['state_dict'])
            print("===> loaded checkpoint '{}' (epoch {})"
                  .format(pretrained, checkpoint['epoch']))
        else:
            print("===> no checkpoint found at '{}'".format(pretrained))

    # Data loading code (class analysis for multi-class classification only)
    train_data = SpoofDatsetSystemID(data_files['train_scp'], data_files['train_utt2index'], binary_class=False)
    val_data = SpoofDatsetSystemID(data_files['dev_scp'], data_files['dev_utt2index'], binary_class=False)
    train_loader = torch.utils.data.DataLoader(
        train_data, batch_size=test_batch_size, shuffle=False, **kwargs)
    val_loader = torch.utils.data.DataLoader(
        val_data, batch_size=test_batch_size, shuffle=False, **kwargs)

    # class analysis for train and dev
    print("===> class analysis for train set")
    class_analysis(train_loader, model, device, use_rnn, class_labels, use_focal=model_params['FOCAL_GAMMA'])
    print("===> class analysis for dev set")
    class_analysis(val_loader, model, device, use_rnn, class_labels, use_focal=model_params['FOCAL_GAMMA'])

@ex.capture
def forward_pass(_run, pretrained_model_id, test_batch_size, data_files, model_params, eer_criteria, class_labels):
    """ forward pass dev and eval data to trained model  """

    use_cuda = torch.cuda.is_available()  # use cpu
    device = torch.device("cuda" if use_cuda else "cpu")
    tmp = torch.tensor([2]).to(device)

    # model is RNN
    if model_params['MODEL_SELECT'] == 4:
        use_rnn = True
    else:
        use_rnn = False

    # model is trained with focal loss objective
    if model_params['FOCAL_GAMMA']:
        print('training with focal loss')
        focal_obj = FocalLoss(gamma=model_params['FOCAL_GAMMA'])
    else:
        focal_obj = None

    kwargs = {'num_workers': 2, 'pin_memory': True} if use_cuda else {}

    # create model
    # cnx
    model = E2E(**model_params).to(device)
    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('===> Model total parameter: {}'.format(num_params))

    if pretrained_model_id:
        pretrain_pth = 'snapshots/' + str(pretrained_model_id) + '/model_best.pth.tar'
        #pretrain_pth = '../pretrained/pa/senet34_py3'
        if os.path.isfile(pretrain_pth):
            print("===> loading checkpoint '{}'".format(pretrain_pth))
            # python 2
            checkpoint = torch.load(pretrain_pth, map_location=lambda storage, loc: storage)  # load for cpu
            model.load_state_dict(checkpoint['state_dict'], strict=False)
            # python 3
            # checkpoint = torch.load(pretrain_pth)
            # model.load_state_dict(checkpoint, strict=False)
            print("===> loaded checkpoint '{}'"
                  .format(pretrain_pth))
        else:
            print("===> no checkpoint found at '{}'".format(pretrain_pth))
            exit()

    # Data loading code (class analysis for multi-class classification only)
    val_data = SpoofDatsetSystemID(data_files['dev_scp'], data_files['dev_utt2index'], binary_class=False)
    # val_data = SpoofDatsetEval(data_files['dev_scp'])
    # eval_data = SpoofDatsetEval(data_files['eval_scp'])
    val_loader = torch.utils.data.DataLoader(val_data,
                                             batch_size=test_batch_size,
                                             shuffle=False, **kwargs)
    # eval_loader = torch.utils.data.DataLoader(
    #     eval_data, batch_size=test_batch_size, shuffle=False, **kwargs)

    # forward pass for dev
    print("===> forward pass for dev set")
    score_file_pth = os.path.join(data_files['scoring_dir'], str(pretrained_model_id) + '-dev_scores.txt')
    print("===> dev scoring file saved at: '{}'".format(score_file_pth))
    prediction(val_loader, model, device, score_file_pth, data_files['dev_utt2systemID'], use_rnn, focal_obj)
    # forward pass for eval
    # print("===> forward pass for eval set")
    # score_file_pth = os.path.join(data_files['scoring_dir'], str(pretrained_model_id) + '-eval_scores.txt')
    # print("===> eval scoring file saved at: '{}'".format(score_file_pth))
    # eval_prediction(eval_loader, model, device, score_file_pth, data_files['eval_utt2systemID'], use_rnn, focal_obj)

def train(train_loader, model, optimizer, epoch, device, log_interval, rnn, focal_obj):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()

    # switch to train mode
    model.train()

    end = time.time()
    for i, (_, input, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        # Create vaiables
        # TODO fix sampling accross utterances!!!
        # input = input[0].to(device)
        # target = torch.tensor(target).to(device).view((-1,))
        # TODO note, features comes trasposed
        input = torch.reshape(input, (input.size(0) * input.size(1),) + input.shape[2:])
        input = input.to(device)
        target = torch.stack(target)
        target = target.reshape(-1)
        target = target.to(device).view((-1,))

        # compute output
        if rnn:
            hidden = model.init_hidden(input.size(0))
            output = model(input, hidden)
        else:
            output = model(input)

        # loss
        if focal_obj:
            loss = focal_obj(output, target)
        else:
            loss = F.nll_loss(output, target)

        # measure accuracy and record loss
        acc1, = accuracy(output, target, topk=(1,))
        losses.update(loss.item(), input.size(0))
        top1.update(acc1[0], input.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        lr = optimizer.update_learning_rate()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % log_interval == 0:
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'LR {lr:.6f}\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Acc@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                epoch, i, len(train_loader), batch_time=batch_time,
                data_time=data_time, lr=lr, loss=losses, top1=top1))

def validate(val_loader, utt2systemID_file, model, device, log_interval, rnn, eer, focal_obj):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()

    # switch to evaluate mode
    model.eval()
    utt2scores = defaultdict(list)
    with torch.no_grad():
        end = time.time()
        for i, (utt_list, input, target) in enumerate(val_loader):
            # TODO fix later the collacation function
            input = input[0].to(device)
            target = torch.tensor(target).to(device).view((-1,))
            # input = input.to(device, non_blocking=True)
            # target = target.to(device, non_blocking=True).view((-1,))

            # compute output
            if rnn:
                hidden = model.init_hidden(input.size(0))
                output = model(input, hidden)
            else:
                output = model(input)

            # loss
            if focal_obj:
                loss = focal_obj(output, target)
            else:
                loss = F.nll_loss(output, target)

            # measure accuracy and record loss
            acc1, = accuracy(output, target, topk=(1,))
            losses.update(loss.item(), input.size(0))
            top1.update(acc1[0], input.size(0))

            # for model selection criteria == dev_eer only
            if eer:
                output = output.cpu().numpy()
                score =  1. - 1. / (1. + np.exp(-output[:, 0])) #output[:, 0]  # use log-probability of the bonafide class for scoring
                for index, utt_id in enumerate(utt_list):
                    utt2scores[utt_id[0]].append(score[index].item())

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % log_interval == 0:
                print('Test: [{0}/{1}]\t'
                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                      'Acc@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                    i, len(val_loader), batch_time=batch_time, loss=losses,
                    top1=top1))

    print('===> Acc@1 {top1.avg:.3f}\n'.format(top1=top1))

    if eer:
        # first do averaging
        with open(utt2systemID_file, 'r') as f:
            temp = f.readlines()
        content = [x.strip() for x in temp]
        utt_list = [x.split()[0] for x in content]
        id_list = [x.split()[1] for x in content]

        spoof_cm, bona_cm = [], []
        for index, utt_id in enumerate(utt_list):
            if utt_id not in utt2scores.keys():  # condition for leave on out
                continue
            score_list = utt2scores[utt_id]
            avg_score = np.mean(score_list)
            spoof_id = id_list[index]
            if spoof_id == 'bonafide':
                bona_cm.append(avg_score)
            else:
                spoof_cm.append(avg_score)

        spoof_cm, bona_cm = np.array(spoof_cm), np.array(bona_cm)
        eer_cm = 100. * compute_eer(bona_cm, spoof_cm)[0]

        print('===> EER_CM: {}\n'.format(eer_cm))

    if eer:
        return eer_cm
    else:
        return top1.avg

def prediction(val_loader, model, device, output_file, utt2systemID_file, rnn, focal_obj):
    # switch to evaluate mode
    utt2scores = defaultdict(list)
    model.eval()

    with torch.no_grad():
        for i, (utt_list, input, target) in enumerate(val_loader):
            input = input[0].to(device)
            # target = target.to(device).view((-1,))
            # print(i, utt_list)

            # compute output
            if rnn:
                hidden = model.init_hidden(input.size(0))
                output = model(input, hidden)
            else:
                output = model(input)

            # get score
            if focal_obj:
                output = F.log_softmax(output, dim=-1)  # apply softmax if model trained with focal loss
            output = output.cpu().numpy()
            score = 1. - 1. / (1. + np.exp(-output[:, 0])) # output[:, 0] #1. - 1. / (1. + np.exp(-output[:, 0]))  # output[:, 0]  # use log-probability of the bonafide class for scoring

            for index, utt_id in enumerate(utt_list):
                # curr_utt = ''.join(utt_id.split('-')[0] + '-' + utt_id.split('-')[1])
                utt2scores[utt_id[0]].append(score[index].item())
                # first do averaging
        with open(utt2systemID_file, 'r') as f:
            temp = f.readlines()
        content = [x.strip() for x in temp]
        utt_list = [x.split()[0] for x in content]
        id_list = [x.split()[1] for x in content]

        # os.mkdir(os.path.split(output_file)[0], True)
        with open(output_file, 'w') as f:
            for index, utt_id in enumerate(utt_list):
                score_list = utt2scores[utt_id]
                avg_score = np.mean(score_list)  # reduce(lambda x, y: x + y, score_list) / len(score_list)
                spoof_id = id_list[index]
                if spoof_id == 'bonafide':
                    f.write('%s %s %s %f\n' % (utt_id, '-', 'bonafide', avg_score))
                else:
                    f.write('%s %s %s %f\n' % (utt_id, spoof_id, 'spoof', avg_score))

def eval_prediction(eval_loader, model, device, output_file, utt2systemID_file, rnn, focal_obj):
    """ eval data's utt2systemID does not have ID """

    # switch to evaluate mode
    utt2scores = defaultdict(list)
    model.eval()

    with torch.no_grad():
        for i, (utt_list, input) in enumerate(eval_loader):
            # input = input.to(device, non_blocking=True)
            input = input.to(device)

            # compute output
            if rnn:
                hidden = model.init_hidden(input.size(0))
                output = model(input, hidden)
            else:
                output = model(input)

            # get score
            if focal_obj:
                output = F.log_softmax(output, dim=-1)  # apply softmax if model trained with focal loss
            score = output[:, 0]  # use log-probability of the bonafide class for scoring

            for index, utt_id in enumerate(utt_list):
                curr_utt = ''.join(utt_id.split('-')[0] + '-' + utt_id.split('-')[1])
                utt2scores[curr_utt].append(score[index].item())

                # first do averaging
        with open(utt2systemID_file, 'r') as f:
            temp = f.readlines()
        content = [x.strip() for x in temp]
        utt_list = content

        with open(output_file, 'w') as f:
            for index, utt_id in enumerate(utt_list):
                score_list = utt2scores[utt_id]
                avg_score = reduce(lambda x, y: x + y, score_list) / len(score_list)
                f.write('%s %f\n' % (utt_id, avg_score))

def class_analysis(val_loader, model, device, rnn, class_labels, use_focal):
    model.eval()

    # for recording acc. for each class
    class_correct = list(0. for i in range(len(class_labels)))
    class_total = list(0. for i in range(len(class_labels)))

    with torch.no_grad():
        for i, (utt_list, input, target) in enumerate(val_loader):
            input = input.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True).view((-1,))

            # compute output
            if rnn:
                hidden = model.init_hidden(input.size(0))
                output = model(input, hidden)
            else:
                output = model(input)

            # get score
            if use_focal:
                output = F.log_softmax(output, dim=-1)  # apply softmax if model trained with focal loss

            # measure accuracy
            _, predicted = torch.max(output, 1)
            c = (predicted == target).squeeze()

            # record results
            for i in range(input.size(0)):
                label = target[i]  # get which class
                class_correct[label] += c[i].item()
                class_total[label] += 1

    for i in range(len(class_labels)):
        print("===> Results for each class")
        print('Accuracy of %8s : %6f %%' % (class_labels[i],
                                            100. * class_correct[i] / class_total[i]))

def save_checkpoint(state, is_best, path, filename='checkpoint.pth.tar'):
    torch.save(state, path + '/' + filename)
    if is_best:
        print("===> save to checkpoint at {}\n".format(path + '/' + 'model_best.pth.tar'))
        shutil.copyfile(path + '/' + filename, path + '/' + 'model_best.pth.tar')

class AverageMeter(object):
    """ Computes and stores the average and current value """

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def accuracy(output, target, topk=(1,)):
    """ Computes the accuracy over the k top predictions for the specified values of k """
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res

class ScheduledOptim(object):
    """ A simple wrapper class for learning rate scheduling """

    def __init__(self, optimizer, n_warmup_steps):
        self.optimizer = optimizer
        self.d_model = 64
        self.n_warmup_steps = n_warmup_steps
        self.n_current_steps = 0
        self.delta = 1

    def step(self):
        "Step by the inner optimizer"
        self.optimizer.step()

    def zero_grad(self):
        "Zero out the gradients by the inner optimizer"
        self.optimizer.zero_grad()

    def increase_delta(self):
        self.delta *= 2

    def update_learning_rate(self):
        "Learning rate scheduling per step"

        self.n_current_steps += self.delta
        new_lr = np.power(self.d_model, -0.5) * np.min([
            np.power(self.n_current_steps, -0.5),
            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])

        for param_group in self.optimizer.param_groups:
            param_group['lr'] = new_lr
        return new_lr

    def state_dict(self):
        ret = {
            'd_model': self.d_model,
            'n_warmup_steps': self.n_warmup_steps,
            'n_current_steps': self.n_current_steps,
            'delta': self.delta,
        }
        ret['optimizer'] = self.optimizer.state_dict()
        return ret

    def load_state_dict(self, state_dict):
        self.d_model = state_dict['d_model']
        self.n_warmup_steps = state_dict['n_warmup_steps']
        self.n_current_steps = state_dict['n_current_steps']
        self.delta = state_dict['delta']
        self.optimizer.load_state_dict(state_dict['optimizer'])

@ex.automain
def main():
    #work()
    # post()
    forward_pass()


from __future__ import print_function
import numpy
import pyRMSD.calculators
from pyRMSD.availableCalculators import availableCalculators
from pyRMSD.symmTools import symm_permutations, swap_atoms, min_rmsd_of_rmsds_list,\
    symm_groups_validation

class RMSDCalculator(object):
    """
    Using the calculator will sometimes modify your input coordinates (shape and value). If the user wants to preserve the
    input coordinates, the best way would be to save them to disk and recover them after the operation (i.e. numpy.save or
    similar functions).
    """
    def __init__(self, calculatorType,
                 fittingCoordsets,
                 calculationCoordsets = None,
                 fitSymmetryGroups = [],
                 calcSymmetryGroups = []):
        """
        Class constructor.

        @param calculatorType: One of the calculators returned by 'availableCalculators()'. i.e. KABSCH_OMP_CALCULATOR

        @param fittingCoordsets: An array containing the used coordinates of each conformation. It has the following form:
            coordsets: [Conformation 1, Conformation 2, ..., Conformation N]
            Conformation: [Atom 1, Atom 2,..., Atom M]
            Atom: [x,y,z]

            This coordinates will be used for both structural superposition and RMSD calculation if the 'calculation
            coordinates' parameter is not defined.

            The array type is constrained to be a numpy.array object with dtype = numpy.float64 (dtype will not be always double
            by default).

            Input coordinates are modified after each operation (usually centered and superposed into reference conformation).

        @param calculationCoordsets: An array containing the coordinates used to calculate the RMSD. Must have the same structure
            than 'fittingCoordinates'.

        @param fitSymmetryGroups: List of symmetry groups.
            Symm. groups are a low-level structure-agnostic of the type of symmetries that can be found in some ligands,
            i.e. in rotating benzene groups. It can also be used in symmetries of bigger selections though.
            See :py:func:`pyRMSD.symmTools.symm_permutations`
        @param calcSymmetryGroups: As in 'fitSymmetryGroups', a list of symmetry groups.

        @author: vgil
        @date: 26/11/2012
        """
        if not calculatorType in availableCalculators():
            print("Calculator ", calculatorType, " is not an available calculator.")
            raise ValueError
        else:
            self.fitting_coordinates = fittingCoordsets
            self.calculator_type = calculatorType
            self.number_of_conformations = self.fitting_coordinates.shape[0]
            self.number_of_fitting_atoms = self.fitting_coordinates.shape[1]

            self.calculation_coordinates = calculationCoordsets
            if self.calculation_coordinates is not None:
                self.calculation_coordinates = calculationCoordsets
                if self.number_of_conformations != self.calculation_coordinates.shape[0]:
                    print("Calculation coordinates must hold the same number of conformations than fitting coordinates.")
                    raise ValueError
                self.number_of_calculation_atoms = self.calculation_coordinates.shape[1]
            else:
                self.number_of_calculation_atoms = 0

            # Default values for openMP and CUDA flags
            self.__threads_per_block = 32
            self.__blocks_per_grid = 8
            self.__number_of_threads = 8

            # Symmetry group handling
            symm_groups_validation(fitSymmetryGroups)
            symm_groups_validation(calcSymmetryGroups)
            self.fit_symmetry_groups = fitSymmetryGroups
            self.calc_symmetry_groups = calcSymmetryGroups

    def pairwise(self, first_conformation_number, second_conformation_number, get_superposed_coordinates = False):
        """
        Calculates the rmsd of two conformations. As it must create a copy of the conformations, it does not modify the
        input coordinates. However input coordinates constantness is never guaranteed.

        @param first_conformation_number: Id of the reference conformation.

        @param second_conformation_number: Id of the conformation that will be superposed into the reference conformation.

        @param get_superposed_coordinates: If true, the function will also return the superposed coordinates.

        @return: The RMSD value or, if get_superposed_coordinates is True, a tuple consisting of the RMSD of both conformations
        along with the superposed fitting coordinates (copy), and superposed calculation coordinates if defined (also an unbound
        copy). Superposed fitting and calculation resulting coordinates is a 2 conformation array.

        @author: vgil
        @date: 26/11/2012
        """
        # Working with fitting coordinates
        first_coords = self.fitting_coordinates[first_conformation_number]
        second_coords = self.fitting_coordinates[second_conformation_number]
        tmp_coordsets = numpy.copy(numpy.array([first_coords,second_coords]))

        # Then with calculation coordinates (if available)
        tmp_calculation_coordsets = None
        if  self.calculation_coordinates is not None:
            first_calculation_coords = self.calculation_coordinates[first_conformation_number]
            second_calculation_coords = self.calculation_coordinates[second_conformation_number]
            tmp_calculation_coordsets = numpy.array([first_calculation_coords,second_calculation_coords])

        rmsd = RMSDCalculator(self.calculator_type,
                              tmp_coordsets,
                              tmp_calculation_coordsets,
                              self.fit_symmetry_groups,
                              self.calc_symmetry_groups).oneVsFollowing(0)[0]

        if get_superposed_coordinates:
            if tmp_calculation_coordsets is None:
                return rmsd, tmp_coordsets
            else:
                return rmsd, tmp_coordsets, tmp_calculation_coordsets
        else:
            return rmsd

    def oneVsTheOthers(self, conformation_number, get_superposed_coordinates = False):
        """
        Calculates the RMSD between a reference conformation and all other conformations in the set.

        @param conformation_number: The id of the reference structure.

        @return: The array of RMSD values or a tuple consisting of the RMSD array, a superposed copy of fitting coordinates, and a
        superposed copy of calculation coordinates (this last only if calculation coordinates were defined). Unless "conformation_number"
        is 0, in the returned coordinate sets the first conformation is the one with "conformation_number" in the original input coordinates.
        All the other conformations with index < "conformation_number" will be indexed with its original index + 1.

        @author: vgil
        @date: 26/11/2012
        """
        # Rearrange fitting coordinates
        previous_coords = self.fitting_coordinates[:conformation_number]
        following_coords = self.fitting_coordinates[conformation_number+1:]
        rearranged_coords_list = [numpy.copy(self.fitting_coordinates[conformation_number])]

        for coords in previous_coords:
            rearranged_coords_list.append(numpy.copy(coords))

        for coords in following_coords:
            rearranged_coords_list.append(numpy.copy(coords))

        rearranged_coords = numpy.array(rearranged_coords_list)

        # Rearrange calculation coordinates
        rearranged_calculation_coords = None
        if self.calculation_coordinates is not None:
            previous_coords = self.calculation_coordinates[:conformation_number]
            following_coords = self.calculation_coordinates[conformation_number+1:]
            rearranged_calculation_coords_list = [self.calculation_coordinates[conformation_number]]

            for coords in previous_coords:
                rearranged_calculation_coords_list.append(coords)

            for coords in following_coords:
                rearranged_calculation_coords_list.append(coords)

            rearranged_calculation_coords = numpy.array(rearranged_calculation_coords_list)

        rmsd_array= RMSDCalculator(self.calculator_type,
                                   rearranged_coords,
                                   rearranged_calculation_coords,
                                   self.fit_symmetry_groups,
                                   self.calc_symmetry_groups).oneVsFollowing(0)

        if get_superposed_coordinates:
            if rearranged_calculation_coords is None:
                return rmsd_array, rearranged_coords
            else:
                return rmsd_array, rearranged_coords, rearranged_calculation_coords
        else:
            return rmsd_array

    def oneVsFollowing(self, conformation_number):
        """
        Calculates the RMSD between a reference conformation and all other conformations with an id greater than it.
        If fitting symmetry groups are used, input coordinates won't be modified (as it works with coordinate copies).

        @param conformation_number: The id of the reference structure.

        @return: A numpy array of RMSD values.

        @author: vgil
        @date: 26/11/2012
        """
        np_coords_fit, np_coords_calc = self.__coords_reshaping()

        if self.fit_symmetry_groups == []:
            rmsds =  pyRMSD.calculators.oneVsFollowing(availableCalculators()[self.calculator_type],
                                                         np_coords_fit,
                                                         self.number_of_fitting_atoms,
                                                         np_coords_calc,
                                                         self.number_of_calculation_atoms,
                                                         conformation_number,
                                                         self.number_of_conformations,
                                                         self.calc_symmetry_groups,
                                                         self.__number_of_threads,
                                                         self.__threads_per_block,
                                                         self.__blocks_per_grid)
        else:
            # If we have fitting symmetry groups, we have to try with all possible combinations.
            # Calculation symmetry groups are applied at C level, changing the way RMSD is calculated.
            symm_rmsds = []
            for permutation in symm_permutations(self.fit_symmetry_groups):

                # Copy the coordinates and convert to matrix form for ease of indexing
                coords_copy = numpy.array(np_coords_fit, copy= True, dtype = numpy.float64)
                coords_copy.shape = (self.number_of_conformations, self.number_of_fitting_atoms,3)

                # Apply the changes to reference
                for symm_group in permutation:
                    # Do it only if the symm. group is not permuted. Otherwise we would always permute!
                    if not symm_group in self.fit_symmetry_groups:
                        for symm_pair in symm_group:
                            swap_atoms(coords_copy[conformation_number], symm_pair[0], symm_pair[1])

                # Flatten again to feed the C calculator
                coords_copy.shape = (self.number_of_conformations*self.number_of_fitting_atoms*3)

                # And calculate the RMSD of this permutation
                symm_rmsds.append(pyRMSD.calculators.oneVsFollowing(availableCalculators()[self.calculator_type],
                                                         coords_copy,
                                                         self.number_of_fitting_atoms,
                                                         np_coords_calc,
                                                         self.number_of_calculation_atoms,
                                                         conformation_number,
                                                         self.number_of_conformations,
                                                         self.calc_symmetry_groups,
                                                         self.__number_of_threads,
                                                         self.__threads_per_block,
                                                         self.__blocks_per_grid))
            # Pick the minimum rmsd of all possibilities.
            rmsds = min_rmsd_of_rmsds_list(numpy.array(symm_rmsds))

        return rmsds

    def pairwiseRMSDMatrix(self):
        """
        Calculates the pairwise RMSD matrix for all conformations in the coordinates set.

        @return: A numpy array with the upper triangle of the matrix, in row major format.

        @author: vgil
        @date: 26/11/2012
        """
        np_coords_fit, np_coords_calc = self.__coords_reshaping()

        rmsd_values = pyRMSD.calculators.calculateRMSDCondensedMatrix(availableCalculators()[self.calculator_type],
                                                               np_coords_fit,
                                                               self.number_of_fitting_atoms,
                                                               np_coords_calc,
                                                               self.number_of_calculation_atoms,
                                                               self.number_of_conformations,
                                                               self.calc_symmetry_groups,
                                                               self.__number_of_threads,
                                                               self.__threads_per_block,
                                                               self.__blocks_per_grid)

        # It has been necessary to add this hack to the function. When using symmetries, in one Test case,
        # the only way to make it work was to add this list() hack and changing return PyArray_Return(rmsds_list_obj);
        # by Py_INCREF(rmsds_list_obj); + return (PyObject*) rmsds_list_obj; in calculateRMSDCondensedMatrix
        # This problem only happens in some cases and I've been unable to guess why.
        # The error says: SystemError: ../Objects/listobject.c:169: bad argument to internal function
        return list(rmsd_values)

    def iterativeSuperposition(self):
        """
        Calculates an iterative superposition of a set of conformations. When using this function,
        the input coordinates are changed (it wouldn't have too much sense otherwise).
        In this case calculation coordinates are not used as such, but as a secondary coordinates set
        that is rotate along with the primary fitting set (which allows to use different selections
        to do the iterative fit and any other calculation). Fit symmetry groups have no effect in iterative
        superposition.

        @author: vgil
        @date: 27/03/2013
        """
        np_coords_fit, np_coords_calc = self.__coords_reshaping()

        pyRMSD.calculators.iterativeSuperposition(availableCalculators()[self.calculator_type],
                                                   np_coords_fit,
                                                   self.number_of_fitting_atoms,
                                                   np_coords_calc,
                                                   self.number_of_calculation_atoms,
                                                   self.number_of_conformations,
                                                   self.calc_symmetry_groups,
                                                   self.__number_of_threads,
                                                   self.__threads_per_block,
                                                   self.__blocks_per_grid)

        return numpy.reshape(np_coords_fit, (self.number_of_conformations, self.number_of_fitting_atoms, 3))

    def __coords_reshaping(self):
        """
        Flattens the coordinates.
        """
        # This kind of reshaping should not do a copy
        np_coords_fit = numpy.reshape(self.fitting_coordinates, self.number_of_conformations*self.number_of_fitting_atoms*3)
        np_coords_calc = numpy.array([])
        if (self.calculation_coordinates is not None):
            np_coords_calc = numpy.reshape(self.calculation_coordinates, self.number_of_conformations*self.number_of_calculation_atoms*3)
        return np_coords_fit, np_coords_calc

    def setNumberOfOpenMPThreads(self, number_of_threads):
        """
        Sets the number of threads to be used by an OpenMP calculator.

        @param number_of_threads: The number of threads to be used by OpenMP runtime.

        @author: vgil
        @date: 26/11/2012
        """
        self.__number_of_threads = number_of_threads

    def setCUDAKernelThreadsPerBlock(self, number_of_threads, number_of_blocks):
        """
        Sets the number of threads per block and blocks per grid in CUDA calculators.

        @param number_of_threads: Number of threads per block to be used when launching CUDA Kernels.

        @param number_of_blocks: Number of blocks per grid to be used when launching CUDA Kernels.

        @author: vgil
        @date: 26/11/2012
        """
        self.__threads_per_block = number_of_threads
        self.__blocks_per_grid = number_of_blocks


from abc import ABC, abstractmethod
from typing import Optional, List, Tuple, Dict

import numpy as np
import torch as th
import torch.nn as nn
from gym import spaces
from sb3_contrib.common.maskable.distributions import MaskableCategorical
from stable_baselines3.common.distributions import Distribution

class VATDistribution(Distribution, ABC):
    @abstractmethod
    def apply_vat(self, vat: Optional[np.ndarray]) -> None:
        pass

class VATMultiCategoricalDistribution(VATDistribution):
    """
    MultiCategorical distribution for multi discrete actions. Supports invalid action masking.

    :param action_dims: List of sizes of discrete action spaces
    """

    def __init__(self, action_dims: List[int], device: th.device):
        super().__init__()
        self.distributions: List[MaskableCategorical] = []
        self.valid_action_trees: Optional[np.ndarray] = None
        self.action_dims = action_dims
        self.device = device

    def proba_distribution_net(self, latent_dim: int) -> nn.Module:
        """
        Create the layer that represents the distribution:
        it will be the logits (flattened) of the MultiCategorical distribution.
        You can then get probabilities using a softmax on each sub-space.

        :param latent_dim: Dimension of the last layer
            of the policy network (before the action layer)
        :return:
        """

        action_logits = nn.Linear(latent_dim, sum(self.action_dims))
        return action_logits

    def proba_distribution(self, action_logits: th.Tensor) -> "VATMultiCategoricalDistribution":
        # Restructure shape to align with logits
        reshaped_logits = action_logits.view(-1, sum(self.action_dims))

        self.distributions = [
            MaskableCategorical(logits=split) for split in th.split(reshaped_logits, tuple(self.action_dims), dim=1)
        ]
        return self

    def log_prob(self, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
        assert len(self.distributions) > 0, "Must set distribution parameters"

        logp = th.zeros([self.valid_action_trees.shape[0]]).to(self.device)
        entropy = th.zeros([self.valid_action_trees.shape[0]]).to(self.device)
        split_dists = [VATMultiCategoricalDistribution._split_categorical(d) for d in self.distributions]
        # transpose it
        split_dists = list(map(list, zip(*split_dists)))
        for idx, (action, vat, dists) in enumerate(zip(actions, self.valid_action_trees, split_dists)):
            action_logp, action_entropy = self._logp_and_entropy_for_action(action, vat, dists)
            logp[idx] = action_logp.sum()
            entropy[idx] = action_logp.sum()

        return logp, entropy

    def _logp_and_entropy_for_action(self, action: th.Tensor, vat: np.ndarray, dists: List[MaskableCategorical]):
        assert isinstance(vat[0], Dict)
        logp = th.zeros([len(action)]).to(self.device)
        entropy = th.zeros([len(action)]).to(self.device)

        subtree = vat[0]
        action_type, source_node, target_node, local_vuln, remote_vuln, port, cred = range(7)
        sel_action = action[action_type]
        options = list(subtree.keys())
        # mask = np.full_like(dists[action_type].logits, False, dtype=bool)
        mask = np.full(tuple(dists[action_type].logits.shape), False, dtype=bool)
        mask[options] = True
        dists[action_type].apply_masking(mask)

        sel_action_logp = dists[action_type].log_prob(sel_action)
        logp[action_type] = sel_action_logp
        entropy[action_type] = dists[action_type].entropy()

        subtree = subtree[int(sel_action)]
        other_parts = {
            0: [source_node, local_vuln],
            1: [source_node, target_node, remote_vuln],
            2: [source_node, target_node, port, cred],
        }
        for part in other_parts[int(sel_action)]:
            options = list(subtree.keys())
            selection = action[part]
            # mask = np.full_like(dists[part].logits, False, dtype=bool)
            mask = np.full(tuple(dists[part].logits.shape), False, dtype=bool)
            mask[options] = True
            dists[part].apply_masking(mask)

            lp = dists[part].log_prob(selection)
            logp[part] = lp
            entropy[part] = dists[part].entropy()

            subtree = subtree[int(selection)]

        return logp, entropy

    # def _logp_and_entropy_for_action_part(self, subtree: Dict, sampled: th.Tensor, dist: MaskableCategorical) -> Tuple[th.Tensor, th.Tensor]:
    #     options = list(subtree.keys())
    #     mask = np.full_like(dist.logits[0], False, dtype=bool)
    #     mask[options] = True
    #     self.distributions[dist_pos].apply_masking(mask)
    #
    #     logp = self.distributions[dist_pos].log_prob(sampled)
    #     entropy = self.distributions[dist_pos].entropy()
    #     return logp, entropy

    def entropy(self) -> th.Tensor:
        # need an action for that one
        raise NotImplementedError

    def sample(self) -> th.Tensor:
        assert len(self.distributions) > 0, "Must set distribution parameters"
        return self._mask_and_sample(False)

    def mode(self) -> th.Tensor:
        assert len(self.distributions) > 0, "Must set distribution parameters"
        return self._mask_and_sample(True)

    def actions_from_params(self, action_logits: th.Tensor, deterministic: bool = False) -> th.Tensor:
        # Update the proba distribution
        self.proba_distribution(action_logits)
        return self.get_actions(deterministic=deterministic)

    def log_prob_from_params(self, action_logits: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
        actions = self.actions_from_params(action_logits)
        log_prob = self.log_prob(actions)
        return actions, log_prob[0]

    def apply_vat(self, vat: Optional[np.ndarray]) -> None:
        self.valid_action_trees = vat

    def _mask_and_sample(self, deterministic=False) -> th.Tensor:
        assert self.valid_action_trees is not None

        actions = th.zeros([self.valid_action_trees.shape[0], len(self.action_dims)], dtype=th.int).to(self.device)
        logp = th.zeros([len(self.action_dims)]).to(self.device)
        split_dists = [VATMultiCategoricalDistribution._split_categorical(d) for d in self.distributions]
        # transpose it
        split_dists = list(map(list, zip(*split_dists)))

        for idx, (vat, dists) in enumerate(zip(self.valid_action_trees, split_dists)):
            assert isinstance(vat[0], Dict)
            action, action_logp = self._mask_and_sample_full_action(dists, vat, deterministic)
            actions[idx] = action
            logp[idx] = action_logp

        return actions

    def _mask_and_sample_full_action(self, dists: List[MaskableCategorical], vat: np.ndarray,
                                     deterministic: bool) -> Tuple[th.Tensor, th.Tensor]:
        assert isinstance(vat[0], Dict)

        action = th.zeros(len(self.action_dims), dtype=th.int).to(self.device)
        logp = th.zeros(len(self.action_dims)).to(self.device)
        subtree: Dict = vat[0]
        action_type, source_node, target_node, local_vuln, remote_vuln, port, cred = range(7)
        sampled_type, sampled_type_logp, _ = self._mask_and_sample_action_part(dists[action_type], subtree,
                                                                               deterministic)
        action[action_type] = sampled_type
        logp[action_type] = sampled_type_logp

        subtree = subtree[int(sampled_type)]
        other_action_parts = {
            0: [source_node, local_vuln],
            1: [source_node, target_node, remote_vuln],
            2: [source_node, target_node, port, cred],
        }
        for part in other_action_parts[int(sampled_type)]:
            sampled, sampled_logp, _ = self._mask_and_sample_action_part(dists[part], subtree, deterministic)
            action[part] = sampled
            logp[part] = sampled_logp
            subtree = subtree[int(sampled)]

        return action, logp.sum()

    @staticmethod
    def _mask_and_sample_action_part(dist: MaskableCategorical, vat_subtree: Dict,
                                     deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, np.ndarray]:
        options = list(vat_subtree.keys())
        # mask = np.full_like(dist.logits, False, dtype=bool)
        mask = np.full(tuple(dist.logits.shape), False, dtype=bool)
        mask[options] = True
        dist.apply_masking(mask)

        if deterministic:
            sampled = th.argmax(dist.probs)
        else:
            sampled = dist.sample()
        logp = dist.log_prob(sampled)
        return sampled, logp, mask

    @staticmethod
    def _split_categorical(dist: MaskableCategorical) -> List[MaskableCategorical]:
        return [MaskableCategorical(logits=logits) for logits in dist.logits]

def make_vat_proba_distribution(action_space: spaces.Space, device: th.device) -> VATDistribution:
    if isinstance(action_space, spaces.MultiDiscrete):
        return VATMultiCategoricalDistribution(action_space.nvec, device)
    else:
        raise NotImplementedError


from numpy.random import seed
seed(42)
from tensorflow import set_random_seed
set_random_seed(42)

import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)

import string, sys
import numpy as np
import pickle

import matplotlib.pyplot as plt
import json, pathlib

import keras as k
from keras.callbacks import ModelCheckpoint
from keras_contrib.metrics import crf_viterbi_accuracy
from keras.models import Model, Input
from keras.layers import LSTM, Embedding, Dense, TimeDistributed,concatenate,  Dropout, Bidirectional, Lambda, Layer, Conv1D, MaxPooling1D, Flatten
from keras_contrib.layers import CRF
from keras_contrib.losses import crf_loss

sys.path.append("../")
import evaluator

class Learner():
    def __init__(self):
        print("[WELCOME NEURAL NETWORKS DDI]... Init learning progress")
        # import nltk CoreNLP module (just once)

    def load_data(self, path):
        with open(path) as outfile:
            data = json.load(outfile)
        return data

    def create_indexs(self, dataset, max_length):
        '''
        Create index dictionaries both for input ( words ) and output ( labels ) from given dataset .
        '''
        words = ['<PAD>', '<UNK>']
        lemmas = ['<PAD>', '<UNK>']
        tags = ['<PAD>', '<UNK>']
        rel_pos1 = ['<PAD>', '<UNK>']
        rel_pos2 = ['<PAD>', '<UNK>']

        labels = []
        # create reverse index for words, lemmas and tags
        for data in dataset:
            dditype = data["dditype"]
            if dditype not in labels:
                labels.append(dditype)
            for word in data["feats"]:
                feat_word, feat_lemma, feat_tags, feat_pos1, feat_pos2 = word
                if feat_word not in words:
                    words.append(feat_word)
                if feat_lemma not in lemmas:
                    lemmas.append(feat_lemma)
                if feat_tags not in tags:
                    tags.append(feat_tags)
                if feat_pos1 not in rel_pos1:
                    rel_pos1.append(feat_pos1)
                if feat_pos2 not in rel_pos2:
                    rel_pos2.append(feat_pos2)
        words = {k: v for v, k in enumerate(words)}
        lemmas = {k: v for v, k in enumerate(lemmas)}
        tags = {k: v for v, k in enumerate(tags)}
        rel_pos1 = {k: v for v, k in enumerate(rel_pos1)}
        rel_pos2 = {k: v for v, k in enumerate(rel_pos2)}
        labels = {k: v for v, k in enumerate(labels)}

        result = {}
        result['words'] = words
        result['lemmas'] = lemmas
        result['tags'] = tags
        result['rel_pos1'] = rel_pos1
        result['rel_pos2'] = rel_pos2
        result['labels'] = labels
        result['maxlen'] = max_length
        return result

    def encode_words(self, dataset, idx):
        '''
        Encode the words in a sentence dataset formed by lists of tokens
        into lists of indexes suitable for NN input .

        The dataset encoded as a list of sentence , each of them is a list of
        word indices . If the word is not in the index , <UNK > code is used . If
        the sentence is shorter than max_len it is padded with <PAD > code .
        '''

        def getIndex(feat, index_specific):
            if feat in index_specific:
                index = index_specific[feat]
            else:
                index = index_specific['<UNK>']
            return index

        def pad(array, max_len, tag):
            while len(array) < max_len:
                array.append(tag)
            return array

        results_words = []
        results_lemmas = []
        results_tags = []
        results_pos1 = []
        results_pos2 = []

        for data in dataset:
            encoded_sentence_words = []
            encoded_sentence_lemmas = []
            encoded_sentence_tags = []
            encoded_sentence_pos_1 = []
            encoded_sentence_pos_2 = []
            for word in data['feats']:
                if len(word) != 5:
                    print(word)
                feat_word, feat_lemma, feat_tags, feat_pos_1, feat_pos_2 = word
                encoded_sentence_words.append(getIndex(feat_word, idx["words"]))
                encoded_sentence_lemmas.append(getIndex(feat_lemma, idx["lemmas"]))
                encoded_sentence_tags.append(getIndex(feat_tags, idx["tags"]))
                encoded_sentence_pos_1.append(getIndex(feat_pos_1, idx["rel_pos1"]))
                encoded_sentence_pos_2.append(getIndex(feat_pos_2, idx["rel_pos2"]))

            encoded_sentence_words = pad(encoded_sentence_words, idx["maxlen"], idx["words"]['<PAD>'])
            encoded_sentence_lemmas = pad(encoded_sentence_lemmas, idx["maxlen"], idx["lemmas"]['<PAD>'])
            encoded_sentence_tags = pad(encoded_sentence_tags, idx["maxlen"], idx["tags"]['<PAD>'])
            encoded_sentence_pos_1 = pad(encoded_sentence_pos_1, idx["maxlen"], idx["rel_pos1"]['<PAD>'])
            encoded_sentence_pos_2 = pad(encoded_sentence_pos_2, idx["maxlen"], idx["rel_pos2"]['<PAD>'])

            results_words.append(np.array(encoded_sentence_words))
            results_lemmas.append(np.array(encoded_sentence_lemmas))
            results_tags.append(np.array(encoded_sentence_tags))
            results_pos1.append(np.array(encoded_sentence_pos_1))
            results_pos2.append(np.array(encoded_sentence_pos_2))

        return [np.array(results_words), np.array(results_lemmas), np.array(results_tags), np.array(results_pos1), np.array(results_pos2)]

    def encode_labels(self, dataset, idx):
        '''
        Encode the ground truth labels in a sentence dataset formed by lists of
        tokens into lists of indexes suitable for NN output .
        '''
        results = []
        for data in dataset:
            index = idx["labels"][data['dditype']]
            results.append(np.array(index))
        n_tags = len(idx["labels"])

        def to_categorical(y, num_classes):
            return np.eye(num_classes)[y]

        results = [to_categorical(i, num_classes=n_tags) for i in results]
        results = np.array(results)
        return results

    def save_model_and_indexs(self, model, idx, filename):
        '''
        Save given model and indexs to disk
        '''
        model.save_weights(filename + '.hdf5')
        with open(filename + '.idx', 'wb') as fp:
            pickle.dump(idx, fp, protocol=pickle.HIGHEST_PROTOCOL)

    def load_model_and_indexs(self, filename):
        '''
        Save given model and indexs to disk
        '''
        with open(filename + '.idx', 'rb') as fp:
            data = pickle.load(fp)
        n_words = len(data['words'])
        n_labels = len(data['labels'])
        max_len = data['maxlen']
        model = self.defineModel(n_words, n_labels, max_len)
        model.load_weights(filename + '.hdf5')
        return model, data

    def output_interactions(self, dataset, preds, outfile):
        '''
        Output detected entities in the format expected by the evaluator
        '''
        # if it's not waiting will print the BI elements without the marks
        # in order to not print the O's or print together the BI
        f = open(outfile, "w+")

        for i, pair in enumerate(dataset):
            sid= pair['sid']
            e1 = pair['id_e1']
            e2 = pair['id_e2']
            type_pred = preds[i]
            if type_pred != "null":
                f.write(sid + '|' + e1 + '|' +e2 + '|' + type_pred + '\n')
        f.close()

    def predict(self, modelname, datadir, outfile):
        '''
        Loads a NN model from file ’modelname ’ and uses it to extract drugs
        in datadir . Saves results to ’outfile ’ in the appropriate format .
        '''
        print("[INFO]... Model in inference process")
        # load model and associated encoding data
        model, idx = self.load_model_and_indexs(modelname)
        # load data to annotate
        filename = "features_"+ pathlib.Path(datadir).stem+".txt"
        testdata = self.load_data(filename)
        # encode dataset
        X = self.encode_words(testdata, idx)

        # tag sentences in dataset
        Y = model.predict(X)
        #[s.numpy().argmax() for s in y_val_pred]
        reverse_labels= {y: x for x, y in idx['labels'].items()}
        Y = [reverse_labels[np.argmax(s)] for s in Y]
        # extract entities and dump them to output file
        self.output_interactions(testdata, Y, outfile)

        # evaluate using official evaluator
        self.evaluation(datadir, outfile)

    def checkOutputs(self, modelname, datadir, outfile):
        print("[INFO]... Model in checking process")
        # load model and associated encoding data
        model, idx = self.load_model_and_indexs(modelname)
        # load data to annotate
        filename = "features_"+ pathlib.Path(datadir).stem+".txt"
        testdata = self.load_data(filename)
        # encode dataset
        Y = self.encode_labels(testdata, idx)
        reverse_labels = {y: x for x, y in idx['labels'].items()}
        Y = [reverse_labels[np.argmax(s)] for s in Y]

        # extract entities and dump them to output file
        self.output_interactions(testdata, Y, outfile)

        # evaluate using official evaluator
        self.evaluation(datadir, outfile)

    def evaluation(self, datadir, outfile):
        evaluator.evaluate("DDI", datadir, outfile)

    def learn(self, traindir, validationdir, modelname, finetune=False):
        '''
        Learns a NN model using traindir as training data , and validationdir
        as validation data . Saves learnt model in a file named modelname
        '''
        print("[INFO]... Model architecture in training process")

        # load train and validation data in a suitable form
        traindata = self.load_data(traindir)
        valdata = self.load_data(validationdir)

        # create indexes from training data
        max_len = 200
        idx = self.create_indexs(traindata, max_len)

        n_words=len(idx['words'])
        # load the whole embedding into memory
        embeddings_index = dict()
        f = open('../data/glove.6B/glove.6B.100d.txt', encoding="utf8")
        for line in f:
            values = line.split()
            word = values[0]
            coefs = np.asarray(values[1:], dtype='float32')
            embeddings_index[word] = coefs
        f.close()

        embedding_matrix = np.zeros((n_words, 100))
        h=0
        for word in idx['words']:
            embedding_vector = embeddings_index.get(word)
            if embedding_vector is not None:
                embedding_matrix[h] = embedding_vector
            h+=1

        f = open("./embedding_matrix.txt", 'w')
        for row in embedding_matrix:
            np.savetxt(f,row)
        f.close()

        # build network
        model = self.build_network(idx)

        # encode datasets
        Xtrain = self.encode_words(traindata, idx)
        Ytrain = self.encode_labels(traindata, idx)
        Xval = self.encode_words(valdata, idx)
        Yval = self.encode_labels(valdata, idx)

        # train model
        # Saving the best model only
        filepath = modelname+"-{val_acc:.3f}.hdf5"
        checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True, mode='max')
        callbacks_list = [checkpoint]

        # Fit the best model
        history = model.fit(Xtrain, Ytrain, validation_data=(Xval, Yval), batch_size=256, epochs=20, verbose=1, callbacks=callbacks_list)

        # save model and indexs , for later use in prediction
        self.save_model_and_indexs(model, idx, modelname)

        self.plot(history)

    def plot(self, history):
        # Plot the graph
        print(history.history)
        plt.style.use('ggplot')
        accuracy = history.history['acc']
        val_accuracy = history.history['val_acc']
        loss = history.history['loss']
        val_loss = history.history['val_loss']
        x = range(1, len(accuracy) + 1)

        plt.figure(figsize=(12, 5))
        plt.subplot(1, 2, 1)
        plt.plot(x, accuracy, 'b', label='Training acc')
        plt.plot(x, val_accuracy, 'r', label='Validation acc')
        plt.title('Training and validation accuracy')
        plt.legend()
        plt.subplot(1, 2, 2)
        plt.plot(x, loss, 'b', label='Training loss')
        plt.plot(x, val_loss, 'r', label='Validation loss')
        plt.title('Training and validation loss')
        plt.legend()
        plt.savefig("History_model.jpg")

    def defineModel(self, input_dim, n_labels, max_len):

        embedding_matrix=np.loadtxt("./embedding_matrix.txt").reshape(input_dim, 100)

        word_in = Input(shape=(max_len,))
        word_emb = Embedding(input_dim=input_dim, output_dim=100, input_length=max_len, trainable=False, weights=[embedding_matrix])(
            word_in)  # 20-dim embedding

        lemma_in = Input(shape=(max_len,))
        lemma_emb = Embedding(input_dim=input_dim, output_dim=100, input_length=max_len)(word_in)

        tags_in = Input(shape=(max_len,))
        tags_emb = Embedding(input_dim=input_dim, output_dim=100, input_length=max_len)(lemma_in)

        pos1_in = Input(shape=(max_len,))
        pos1_emb = Embedding(input_dim=input_dim, output_dim=100, input_length=max_len)(tags_in)

        pos2_in = Input(shape=(max_len,))
        pos2_emb = Embedding(input_dim=input_dim, output_dim=100, input_length=max_len)(pos1_in)

        concat = concatenate([word_emb, lemma_emb, tags_emb, pos1_emb, pos2_emb])
        model = Dropout(0.55)(concat)
        model = Conv1D(128, 5, padding='same', activation='relu')(model)
        model = MaxPooling1D(pool_size=2)(model)
        model = Bidirectional(LSTM(units=100,return_sequences=False,recurrent_dropout=0.3,))(model)  # variational biLSTM
        #model = Flatten()(model))
        #model = LSTM(100)(model)
        out = Dense(n_labels, activation='softmax')(model)

        #out = TimeDistributed(Dense(n_labels, activation="softmax"))(model)  # a dense layer as suggested by neuralNer

        # create and compile model
        model = Model([word_in, lemma_in, tags_in, pos1_in, pos2_in], out)
        return model

    def build_network(self,idx):
        '''
        Create network for the learner
        '''
        # sizes
        n_words = len(idx['words'])
        n_lemmas = len(idx['lemmas'])
        n_tags = len(idx['tags'])
        n_pos1 = len(idx['rel_pos1'])
        n_pos2 = len(idx['rel_pos2'])
        input_dim = max(n_words, n_lemmas, n_tags, n_pos1, n_pos2)

        n_labels = len(idx['labels'])
        max_len = idx['maxlen']

        # create network layers
        model = self.defineModel(input_dim, n_labels, max_len)

        # set appropriate parameters (optimizer, loss, etc)
        model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

        model.summary()
        return model

if __name__ == '__main__':
    learner = Learner()
    learner.learn("features_train.txt", "features_devel.txt", "adecuate")
    #learner.checkOutputs("original", "../data/train", "results.txt")
    print("TRAIN")
    learner.predict("adecuate", "../data/train", "results.txt")
    print("\nDEVEL")
    learner.predict("adecuate", "../data/devel", "results.txt")
    print("\nTEST")
    learner.predict("adecuate", "../data/test", "results.txt")


import os
import re
import sys
import cv2
import math
import time
import scipy
import argparse
import matplotlib
import numpy as np
import pylab as plt
#from joblib import Parallel, delayed
import util
import torch
import torch as T
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from collections import OrderedDict
from config_reader import config_reader
from scipy.ndimage.filters import gaussian_filter
#parser = argparse.ArgumentParser()
#parser.add_argument('--t7_file', required=True)
#parser.add_argument('--pth_file', required=True)
#args = parser.parse_args()

torch.set_num_threads(torch.get_num_threads())
weight_name = './model/pose_model.pth'

blocks = {}

# find connection in the specified sequence, center 29 is in the position 15
limbSeq = [[2,3], [2,6], [3,4], [4,5], [6,7], [7,8], [2,9], [9,10], \
           [10,11], [2,12], [12,13], [13,14], [2,1], [1,15], [15,17], \
           [1,16], [16,18], [3,17], [6,18]]

# the middle joints heatmap correpondence
mapIdx = [[31,32], [39,40], [33,34], [35,36], [41,42], [43,44], [19,20], [21,22], \
          [23,24], [25,26], [27,28], [29,30], [47,48], [49,50], [53,54], [51,52], \
          [55,56], [37,38], [45,46]]

# visualize
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
          [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
          [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]

block0  = [{'conv1_1':[3,64,3,1,1]},{'conv1_2':[64,64,3,1,1]},{'pool1_stage1':[2,2,0]},{'conv2_1':[64,128,3,1,1]},{'conv2_2':[128,128,3,1,1]},{'pool2_stage1':[2,2,0]},{'conv3_1':[128,256,3,1,1]},{'conv3_2':[256,256,3,1,1]},{'conv3_3':[256,256,3,1,1]},{'conv3_4':[256,256,3,1,1]},{'pool3_stage1':[2,2,0]},{'conv4_1':[256,512,3,1,1]},{'conv4_2':[512,512,3,1,1]},{'conv4_3_CPM':[512,256,3,1,1]},{'conv4_4_CPM':[256,128,3,1,1]}]

blocks['block1_1']  = [{'conv5_1_CPM_L1':[128,128,3,1,1]},{'conv5_2_CPM_L1':[128,128,3,1,1]},{'conv5_3_CPM_L1':[128,128,3,1,1]},{'conv5_4_CPM_L1':[128,512,1,1,0]},{'conv5_5_CPM_L1':[512,38,1,1,0]}]

blocks['block1_2']  = [{'conv5_1_CPM_L2':[128,128,3,1,1]},{'conv5_2_CPM_L2':[128,128,3,1,1]},{'conv5_3_CPM_L2':[128,128,3,1,1]},{'conv5_4_CPM_L2':[128,512,1,1,0]},{'conv5_5_CPM_L2':[512,19,1,1,0]}]

for i in range(2,7):
    blocks['block%d_1'%i]  = [{'Mconv1_stage%d_L1'%i:[185,128,7,1,3]},{'Mconv2_stage%d_L1'%i:[128,128,7,1,3]},{'Mconv3_stage%d_L1'%i:[128,128,7,1,3]},{'Mconv4_stage%d_L1'%i:[128,128,7,1,3]},
{'Mconv5_stage%d_L1'%i:[128,128,7,1,3]},{'Mconv6_stage%d_L1'%i:[128,128,1,1,0]},{'Mconv7_stage%d_L1'%i:[128,38,1,1,0]}]
    blocks['block%d_2'%i]  = [{'Mconv1_stage%d_L2'%i:[185,128,7,1,3]},{'Mconv2_stage%d_L2'%i:[128,128,7,1,3]},{'Mconv3_stage%d_L2'%i:[128,128,7,1,3]},{'Mconv4_stage%d_L2'%i:[128,128,7,1,3]},
{'Mconv5_stage%d_L2'%i:[128,128,7,1,3]},{'Mconv6_stage%d_L2'%i:[128,128,1,1,0]},{'Mconv7_stage%d_L2'%i:[128,19,1,1,0]}]

def make_layers(cfg_dict):
    layers = []
    for i in range(len(cfg_dict)-1):
        one_ = cfg_dict[i]
        for k,v in one_.iteritems():
            if 'pool' in k:
                layers += [nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2] )]
            else:
                conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], kernel_size=v[2], stride = v[3], padding=v[4])
                layers += [conv2d, nn.ReLU(inplace=True)]
    one_ = cfg_dict[-1].keys()
    k = one_[0]
    v = cfg_dict[-1][k]
    conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], kernel_size=v[2], stride = v[3], padding=v[4])
    layers += [conv2d]
    return nn.Sequential(*layers)

layers = []
for i in range(len(block0)):
    one_ = block0[i]
    for k,v in one_.iteritems():
        if 'pool' in k:
            layers += [nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2] )]
        else:
            conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], kernel_size=v[2], stride = v[3], padding=v[4])
            layers += [conv2d, nn.ReLU(inplace=True)]

models = {}
models['block0']=nn.Sequential(*layers)

for k,v in blocks.items():
    models[k] = make_layers(v)

class pose_model(nn.Module):
    def __init__(self,model_dict,transform_input=False):
        super(pose_model, self).__init__()
        self.model0   = model_dict['block0']
        self.model1_1 = model_dict['block1_1']
        self.model2_1 = model_dict['block2_1']
        self.model3_1 = model_dict['block3_1']
        self.model4_1 = model_dict['block4_1']
        self.model5_1 = model_dict['block5_1']
        self.model6_1 = model_dict['block6_1']

        self.model1_2 = model_dict['block1_2']
        self.model2_2 = model_dict['block2_2']
        self.model3_2 = model_dict['block3_2']
        self.model4_2 = model_dict['block4_2']
        self.model5_2 = model_dict['block5_2']
        self.model6_2 = model_dict['block6_2']

    def forward(self, x):
        out1 = self.model0(x)

        out1_1 = self.model1_1(out1)
        out1_2 = self.model1_2(out1)
        out2  = torch.cat([out1_1,out1_2,out1],1)

        out2_1 = self.model2_1(out2)
        out2_2 = self.model2_2(out2)
        out3   = torch.cat([out2_1,out2_2,out1],1)

        out3_1 = self.model3_1(out3)
        out3_2 = self.model3_2(out3)
        out4   = torch.cat([out3_1,out3_2,out1],1)

        out4_1 = self.model4_1(out4)
        out4_2 = self.model4_2(out4)
        out5   = torch.cat([out4_1,out4_2,out1],1)

        out5_1 = self.model5_1(out5)
        out5_2 = self.model5_2(out5)
        out6   = torch.cat([out5_1,out5_2,out1],1)

        out6_1 = self.model6_1(out6)
        out6_2 = self.model6_2(out6)

        return out6_1,out6_2

model = pose_model(models)
model.load_state_dict(torch.load(weight_name))
model.cuda()
model.float()
model.eval()

param_, model_ = config_reader()

def handle_one(oriImg):

    # for visualize
    canvas = np.copy(oriImg)
    imageToTest = Variable(T.transpose(T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(),0),2,3),1,2),volatile=True).cuda()
    print(oriImg.shape)
    scale = model_['boxsize'] / float(oriImg.shape[0])
    print(scale)
    h = int(oriImg.shape[0]*scale)
    w = int(oriImg.shape[1]*scale)
    pad_h = 0 if (h%model_['stride']==0) else model_['stride'] - (h % model_['stride'])
    pad_w = 0 if (w%model_['stride']==0) else model_['stride'] - (w % model_['stride'])
    new_h = h+pad_h
    new_w = w+pad_w

    imageToTest = cv2.resize(oriImg, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
    imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_['stride'], model_['padValue'])
    imageToTest_padded = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,2,0,1))/256 - 0.5

    feed = Variable(T.from_numpy(imageToTest_padded)).cuda()

    output1,output2 = model.forward(feed)

    heatmap = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output2)

    paf = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output1)

    print(heatmap.size())
    print(paf.size())
    print(type(heatmap))
    heatmap_avg = T.transpose(T.transpose(heatmap[0],0,1),1,2).data.cpu().numpy()
    paf_avg = T.transpose(T.transpose(paf[0],0,1),1,2).data.cpu().numpy()

    all_peaks = []
    peak_counter = 0

    #maps =
    for part in range(18):
        map_ori = heatmap_avg[:,:,part]
        map = gaussian_filter(map_ori, sigma=3)

        map_left = np.zeros(map.shape)
        map_left[1:,:] = map[:-1,:]
        map_right = np.zeros(map.shape)
        map_right[:-1,:] = map[1:,:]
        map_up = np.zeros(map.shape)
        map_up[:,1:] = map[:,:-1]
        map_down = np.zeros(map.shape)
        map_down[:,:-1] = map[:,1:]

        peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param_['thre1']))
    #    peaks_binary = T.eq(
    #    peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])

        peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0])) # note reverse

        peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]
        id = range(peak_counter, peak_counter + len(peaks))
        peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))]

        all_peaks.append(peaks_with_score_and_id)
        peak_counter += len(peaks)

    connection_all = []
    special_k = []
    mid_num = 10

    for k in range(len(mapIdx)):
        score_mid = paf_avg[:,:,[x-19 for x in mapIdx[k]]]
        candA = all_peaks[limbSeq[k][0]-1]
        candB = all_peaks[limbSeq[k][1]-1]
        nA = len(candA)
        nB = len(candB)
        indexA, indexB = limbSeq[k]
        if(nA != 0 and nB != 0):
            connection_candidate = []
            for i in range(nA):
                for j in range(nB):
                    vec = np.subtract(candB[j][:2], candA[i][:2])
                    norm = math.sqrt(vec[0]*vec[0] + vec[1]*vec[1])
                    vec = np.divide(vec, norm)

                    startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
                                   np.linspace(candA[i][1], candB[j][1], num=mid_num))

                    vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
                                      for I in range(len(startend))])
                    vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
                                      for I in range(len(startend))])

                    score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
                    score_with_dist_prior = sum(score_midpts)/len(score_midpts) + min(0.5*oriImg.shape[0]/norm-1, 0)
                    criterion1 = len(np.nonzero(score_midpts > param_['thre2'])[0]) > 0.8 * len(score_midpts)
                    criterion2 = score_with_dist_prior > 0
                    if criterion1 and criterion2:
                        connection_candidate.append([i, j, score_with_dist_prior, score_with_dist_prior+candA[i][2]+candB[j][2]])

            connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
            connection = np.zeros((0,5))
            for c in range(len(connection_candidate)):
                i,j,s = connection_candidate[c][0:3]
                if(i not in connection[:,3] and j not in connection[:,4]):
                    connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
                    if(len(connection) >= min(nA, nB)):
                        break

            connection_all.append(connection)
        else:
            special_k.append(k)
            connection_all.append([])

    # last number in each row is the total parts number of that person
    # the second last number in each row is the score of the overall configuration
    subset = -1 * np.ones((0, 20))
    candidate = np.array([item for sublist in all_peaks for item in sublist])

    for k in range(len(mapIdx)):
        if k not in special_k:
            partAs = connection_all[k][:,0]
            partBs = connection_all[k][:,1]
            indexA, indexB = np.array(limbSeq[k]) - 1

            for i in range(len(connection_all[k])): #= 1:size(temp,1)
                found = 0
                subset_idx = [-1, -1]
                for j in range(len(subset)): #1:size(subset,1):
                    if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
                        subset_idx[found] = j
                        found += 1

                if found == 1:
                    j = subset_idx[0]
                    if(subset[j][indexB] != partBs[i]):
                        subset[j][indexB] = partBs[i]
                        subset[j][-1] += 1
                        subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
                elif found == 2: # if found 2 and disjoint, merge them
                    j1, j2 = subset_idx
                    print("found = 2")
                    membership = ((subset[j1]>=0).astype(int) + (subset[j2]>=0).astype(int))[:-2]
                    if len(np.nonzero(membership == 2)[0]) == 0: #merge
                        subset[j1][:-2] += (subset[j2][:-2] + 1)
                        subset[j1][-2:] += subset[j2][-2:]
                        subset[j1][-2] += connection_all[k][i][2]
                        subset = np.delete(subset, j2, 0)
                    else: # as like found == 1
                        subset[j1][indexB] = partBs[i]
                        subset[j1][-1] += 1
                        subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]

                # if find no partA in the subset, create a new subset
                elif not found and k < 17:
                    row = -1 * np.ones(20)
                    row[indexA] = partAs[i]
                    row[indexB] = partBs[i]
                    row[-1] = 2
                    row[-2] = sum(candidate[connection_all[k][i,:2].astype(int), 2]) + connection_all[k][i][2]
                    subset = np.vstack([subset, row])

    # delete some rows of subset which has few parts occur
    deleteIdx = [];
    for i in range(len(subset)):
        if subset[i][-1] < 4 or subset[i][-2]/subset[i][-1] < 0.4:
            deleteIdx.append(i)
    subset = np.delete(subset, deleteIdx, axis=0)

#    canvas = cv2.imread(test_image) # B,G,R order
    for i in range(18):
        for j in range(len(all_peaks[i])):
            cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)

    stickwidth = 4

    for i in range(17):
        for n in range(len(subset)):
            index = subset[n][np.array(limbSeq[i])-1]
            if -1 in index:
                continue
            cur_canvas = canvas.copy()
            Y = candidate[index.astype(int), 0]
            X = candidate[index.astype(int), 1]
            mX = np.mean(X)
            mY = np.mean(Y)
            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
            polygon = cv2.ellipse2Poly((int(mY),int(mX)), (int(length/2), stickwidth), int(angle), 0, 360, 1)
            cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
            canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)

    return canvas

if __name__ == "__main__":
    print('warming up')
    _ = handle_one(np.ones((320,320,3)))

    #video_capture = cv2.VideoCapture(0)
    video_capture = cv2.VideoCapture("./%%%.avi")
    while True:
        # Capture frame-by-frame
        ret, frame = video_capture.read()

        canvas = handle_one(frame)

        # Display the resulting frame
        cv2.imshow('Video', canvas)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # When everything is done, release the capture
    video_capture.release()
    cv2.destroyAllWindows()


from __future__ import print_function, absolute_import
import numpy as np
import torch
import copy
import os.path as osp
from collections import defaultdict
from opts import market1501_test_map, duke_test_map
import sys

def make_results(qf, gf, lqf, lgf, q_pids, g_pids, q_camids, g_camids, targetmodel, ak_typ, attr_matrix=None, dataset_name=None, attr=None):
  qf, gf = featureNormalization(qf, gf, targetmodel)
  m, n = qf.size(0), gf.size(0)
  distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
  distmat.addmm_(1, -2, qf, gf.t())
  distmat = distmat.numpy()

  if targetmodel == 'aligned':
    from .distance import low_memory_local_dist
    lqf, lgf = lqf.permute(0,2,1), lgf.permute(0,2,1)
    local_distmat = low_memory_local_dist(lqf.numpy(),lgf.numpy(), aligned=True)
    distmat = local_distmat+distmat

  if ak_typ > 0:
    distmat, all_hit, ignore_list = evaluate_attr(distmat, q_pids, g_pids, attr_matrix, dataset_name, attr)
    return distmat, all_hit, ignore_list
  else:
    cmc, mAP = evaluate(distmat, q_pids, g_pids, q_camids, g_camids, use_metric_cuhk03=False)
    return distmat, cmc, mAP

def featureNormalization(qf, gf, targetmodel):
  if targetmodel in ['aligned', 'densenet121', 'hacnn', 'mudeep', 'ide', 'cam', 'lsro', 'hhl', 'spgan']:
    qf = 1. * qf / (torch.norm(qf, p=2, dim=-1, keepdim=True).expand_as(qf) + 1e-12)
    gf = 1. * gf / (torch.norm(gf, p=2, dim=-1, keepdim=True).expand_as(gf) + 1e-12)

  elif targetmodel in ['pcb']:
    qf = (qf / (np.sqrt(6) * torch.norm(qf, p=2, dim=1, keepdim=True).expand_as(qf))).view(qf.size(0), -1)
    gf = (gf / (np.sqrt(6) * torch.norm(gf, p=2, dim=1, keepdim=True).expand_as(gf))).view(gf.size(0), -1)

  return qf, gf

def eval_cuhk03(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, N=100):
    """Evaluation with cuhk03 metric
    Key: one image for each gallery identity is randomly sampled for each query identity.
    Random sampling is performed N times (default: N=100).
    """
    num_q, num_g = distmat.shape
    if num_g < max_rank:
        max_rank = num_g
        print("Note: number of gallery samples is quite small, got {}".format(num_g))
    indices = np.argsort(distmat, axis=1)
    matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)

    # compute cmc curve for each query
    all_cmc = []
    all_AP = []
    num_valid_q = 0. # number of valid query
    for q_idx in range(num_q):
        # get query pid and camid
        q_pid = q_pids[q_idx]
        q_camid = q_camids[q_idx]

        # remove gallery samples that have the same pid and camid with query
        order = indices[q_idx]
        remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)
        keep = np.invert(remove)

        # compute cmc curve
        orig_cmc = matches[q_idx][keep] # binary vector, positions with value 1 are correct matches
        if not np.any(orig_cmc):
            # this condition is true when query identity does not appear in gallery
            continue

        kept_g_pids = g_pids[order][keep]
        g_pids_dict = defaultdict(list)
        for idx, pid in enumerate(kept_g_pids):
            g_pids_dict[pid].append(idx)

        cmc, AP = 0., 0.
        for repeat_idx in range(N):
            mask = np.zeros(len(orig_cmc), dtype=np.bool)
            for _, idxs in g_pids_dict.items():
                # randomly sample one image for each gallery person
                rnd_idx = np.random.choice(idxs)
                mask[rnd_idx] = True
            masked_orig_cmc = orig_cmc[mask]
            _cmc = masked_orig_cmc.cumsum()
            _cmc[_cmc > 1] = 1
            cmc += _cmc[:max_rank].astype(np.float32)
            # compute AP
            num_rel = masked_orig_cmc.sum()
            tmp_cmc = masked_orig_cmc.cumsum()
            tmp_cmc = [x / (i+1.) for i, x in enumerate(tmp_cmc)]
            tmp_cmc = np.asarray(tmp_cmc) * masked_orig_cmc
            AP += tmp_cmc.sum() / num_rel
        cmc /= N
        AP /= N
        all_cmc.append(cmc)
        all_AP.append(AP)
        num_valid_q += 1.

    assert num_valid_q > 0, "Error: all query identities do not appear in gallery"

    all_cmc = np.asarray(all_cmc).astype(np.float32)
    all_cmc = all_cmc.sum(0) / num_valid_q
    mAP = np.mean(all_AP)

    return all_cmc, mAP

def eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank):
    """Evaluation with market1501 metric
    Key: for each query identity, its gallery images from the same camera view are discarded.
    """
    num_q, num_g = distmat.shape
    if num_g < max_rank:
        max_rank = num_g
        print("Note: number of gallery samples is quite small, got {}".format(num_g))
    indices = np.argsort(distmat, axis=1)
    matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)

    # compute cmc curve for each query
    all_cmc = []
    all_AP = []
    num_valid_q = 0. # number of valid query
    for q_idx in range(num_q):
        # get query pid and camid
        q_pid = q_pids[q_idx]
        q_camid = q_camids[q_idx]

        # remove gallery samples that have the same pid and camid with query
        order = indices[q_idx]
        remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)
        keep = np.invert(remove)

        # compute cmc curve
        orig_cmc = matches[q_idx][keep] # binary vector, positions with value 1 are correct matches
        if not np.any(orig_cmc):
            # this condition is true when query identity does not appear in gallery
            continue

        cmc = orig_cmc.cumsum()
        cmc[cmc > 1] = 1

        all_cmc.append(cmc[:max_rank])
        num_valid_q += 1.

        # compute average precision
        # reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision
        num_rel = orig_cmc.sum()
        tmp_cmc = orig_cmc.cumsum()
        tmp_cmc = [x / (i+1.) for i, x in enumerate(tmp_cmc)]
        tmp_cmc = np.asarray(tmp_cmc) * orig_cmc
        AP = tmp_cmc.sum() / num_rel
        all_AP.append(AP)

    assert num_valid_q > 0, "Error: all query identities do not appear in gallery"

    all_cmc = np.asarray(all_cmc).astype(np.float32)
    all_cmc = all_cmc.sum(0) / num_valid_q
    mAP = np.mean(all_AP)

    return all_cmc, mAP

def evaluate(distmat, q_pids, g_pids, q_camids, g_camids, max_rank=20, use_metric_cuhk03=False):
  if use_metric_cuhk03: return eval_cuhk03(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)
  else: return eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)

def evaluate_attr(distmat, q_pids, g_pids, attr_matrix, dataset_name, attr_list, max_rank=20):
  attr_key, attr_value = attr_list
  attr_name = 'duke_attribute' if dataset_name == 'dukemtmcreid' else 'market_attribute'
  offset = 0 if dataset_name == 'dukemtmcreid' else 1
  mapping = duke_test_map if dataset_name == 'dukemtmcreid' else market1501_test_map
  column = attr_matrix[attr_name][0]['test'][0][0][attr_key][0][0]

  num_q, num_g = distmat.shape
  indices = np.argsort(distmat, axis=1)

  all_hit = []
  ignore_list = []
  num_valid_q = 0. # number of valid query
  for q_idx in range(num_q):
    q_pid = q_pids[q_idx]
    if column[mapping[q_pid]-offset] == attr_value:
      ignore_list.append(q_idx)
      continue

    order = indices[q_idx]
    matches = np.zeros_like(order)

    for i in range(len(order)):
      if column[mapping[g_pids[order[i]]]-offset] == attr_value:
        matches[i] = 1

    hit = matches.cumsum()
    hit[hit > 1] = 1
    all_hit.append(hit[:max_rank])
    num_valid_q += 1. # number of valid query

  assert num_valid_q > 0
  all_hit = np.asarray(all_hit).astype(np.float32)
  all_hit = all_hit.sum(0) / num_valid_q

  # distmat = np.delete(distmat, ignore_list, axis=0)

  return distmat, all_hit, ignore_list