#vaemodel
import copy
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.autograd as autograd
from torch.utils import data
from data_loader import DATA_LOADER as dataloader
from dataloader import Flickr30k as dataLoader
import final_classifier as  classifier
import models
from torchvision import models as torchModels
from sklearn.neighbors import NearestNeighbors
from sklearn.manifold import TSNE
from scipy.spatial import distance
import numpy as np

import matplotlib.pyplot as plt

class LINEAR_LOGSOFTMAX(nn.Module):
    def __init__(self, input_dim, nclass):
        super(LINEAR_LOGSOFTMAX, self).__init__()
        self.fc = nn.Linear(input_dim,nclass)
        self.logic = nn.LogSoftmax(dim=1)
        self.lossfunction =  nn.NLLLoss()

    def forward(self, x):
        o = self.logic(self.fc(x))
        return o

class Model(nn.Module):

    def __init__(self,hyperparameters):
        super(Model,self).__init__()

        self.device = hyperparameters['device']
        self.auxiliary_data_source = hyperparameters['auxiliary_data_source']
        self.attr = hyperparameters['attr']
        self.all_data_sources  = ['resnet_features', 'attributes']
        self.DATASET = hyperparameters['dataset']
        self.num_shots = hyperparameters['num_shots']
        self.latent_size = hyperparameters['latent_size']
        self.batch_size = hyperparameters['batch_size']
        self.hidden_size_rule = hyperparameters['hidden_size_rule']
        self.warmup = hyperparameters['model_specifics']['warmup']
        self.generalized = hyperparameters['generalized']
        self.classifier_batch_size = 32
        #self.img_seen_samples   = hyperparameters['samples_per_class'][self.DATASET][0]
        #self.att_seen_samples   = hyperparameters['samples_per_class'][self.DATASET][1]
        #self.att_unseen_samples = hyperparameters['samples_per_class'][self.DATASET][2]
       # self.img_unseen_samples = hyperparameters['samples_per_class'][self.DATASET][3]
        self.reco_loss_function = hyperparameters['loss']
        self.margin = hyperparameters['margin_loss']
        self.nepoch = hyperparameters['epochs']
        self.lr_cls = hyperparameters['lr_cls']
        self.cross_reconstruction = hyperparameters['model_specifics']['cross_reconstruction']
        self.cls_train_epochs = hyperparameters['cls_train_steps']
        #self.dataset = dataloader(self.DATASET, copy.deepcopy(self.auxiliary_data_source) , device= 'cuda')
        self.dataset = dataLoader(copy.deepcopy(self.auxiliary_data_source) , device= 'cuda', attr = self.attr)
        if self.DATASET=='CUB':
            self.num_classes=200
            self.num_novel_classes = 50
        elif self.DATASET=='SUN':
            self.num_classes=717
            self.num_novel_classes = 72
        elif self.DATASET=='AWA1' or self.DATASET=='AWA2':
            self.num_classes=50
            self.num_novel_classes = 10

        if self.attr == 'attributes':
            feature_dimensions = [2048, self.dataset.K]
        elif self.attr == 'bert':
            feature_dimensions = [2048, 768] #2048, 768

        # Here, the encoders and decoders for all modalities are created and put into dict

        self.fc_ft = nn.Linear(2048,2048)
        self.fc_ft.to(self.device)

        self.ft_bn = nn.BatchNorm1d(2048).to(self.device)

        self.fc_at = nn.Linear(self.dataset.K, self.dataset.K)
        self.fc_at.to(self.device)
        self.at_bn = nn.BatchNorm1d(self.dataset.K).to(self.device)

        self.encoder = {}

        for datatype, dim in zip(self.all_data_sources,feature_dimensions):

            self.encoder[datatype] = models.encoder_template(dim,self.latent_size,self.hidden_size_rule[datatype],self.device)

            print(str(datatype) + ' ' + str(dim))

        self.decoder = {}
        for datatype, dim in zip(self.all_data_sources,feature_dimensions):
            self.decoder[datatype] = models.decoder_template(self.latent_size,dim,self.hidden_size_rule[datatype],self.device)

        # An optimizer for all encoders and decoders is defined here
        parameters_to_optimize = list(self.parameters())
        for datatype in self.all_data_sources:
            parameters_to_optimize +=  list(self.encoder[datatype].parameters())
            parameters_to_optimize +=  list(self.decoder[datatype].parameters())
        parameters_to_optimize += list(self.fc_ft.parameters())
        parameters_to_optimize += list(self.fc_at.parameters())
        parameters_to_optimize += list(self.ft_bn.parameters())
        parameters_to_optimize += list(self.at_bn.parameters())

        self.optimizer  = optim.Adam( parameters_to_optimize ,lr=hyperparameters['lr_gen_model'], betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=True)

        if self.reco_loss_function=='l2':
            self.reconstruction_criterion = nn.MSELoss(size_average=False)

        elif self.reco_loss_function=='l1':
            self.reconstruction_criterion = nn.L1Loss(size_average=False)

    def reparameterize(self, mu, logvar):
        if self.reparameterize_with_noise:
            sigma = torch.exp(logvar)
            eps = torch.cuda.FloatTensor(logvar.size()[0],1).normal_(0,1)
            eps  = eps.expand(sigma.size())
            return mu + sigma*eps
        else:
            return mu

    def forward(self):
        pass

    def map_label(self,label, classes):
        mapped_label = torch.LongTensor(label.size()).to(self.device)
        for i in range(classes.size(0)):
            mapped_label[label==classes[i]] = i

        return mapped_label

    def trainstep(self, img, att):

        ##############################################
        # Encode image features and additional
        # features
        ##############################################

        img_in = F.normalize(img, p=2, dim=1)
        img_in = self.ft_bn(img_in)

        img_in = self.fc_ft(img_in)

        att_in = F.normalize(att, p=2, dim=1)
        att_in = self.at_bn(att_in)

        att_in = self.fc_at(att_in)

        #Add non-linearity?

        mu_img, logvar_img = self.encoder['resnet_features'](img_in)
        z_from_img = self.reparameterize(mu_img, logvar_img)

        mu_att, logvar_att = self.encoder[self.auxiliary_data_source](att_in)
        z_from_att = self.reparameterize(mu_att, logvar_att)

        ##############################################
        # Reconstruct inputs
        ##############################################

        img_from_img = self.decoder['resnet_features'](z_from_img)
        att_from_att = self.decoder[self.auxiliary_data_source](z_from_att)

        reconstruction_loss = self.reconstruction_criterion(img_from_img, img) \
                              + self.reconstruction_criterion(att_from_att, att)

        ##############################################
        # Cross Reconstruction Loss
        ##############################################
        img_from_att = self.decoder['resnet_features'](z_from_att)
        att_from_img = self.decoder[self.auxiliary_data_source](z_from_img)

        cross_reconstruction_loss = self.reconstruction_criterion(img_from_att, img) \
                                    + self.reconstruction_criterion(att_from_img, att)

        ##############################################
        # KL-Divergence
        ##############################################

        KLD = (0.5 * torch.sum(1 + logvar_att - mu_att.pow(2) - logvar_att.exp())) \
              + (0.5 * torch.sum(1 + logvar_img - mu_img.pow(2) - logvar_img.exp()))

        ##############################################
        # Distribution Alignment
        ##############################################
        distance = torch.sqrt(torch.sum((mu_img - mu_att) ** 2, dim=1) + \
                              torch.sum((torch.sqrt(logvar_img.exp()) - torch.sqrt(logvar_att.exp())) ** 2, dim=1))

        distance = distance.sum()

        #distanceI = torch.sum((mu_img - mu_att) ** 2, dim=1)
        #distanceT = torch.sum((mu_att - mu_img) ** 2, dim=1)

        #tripletsI = []
        #tripletsT = []

        #for i in range(0, mu_img.shape[0]):
        #    for j in range(0, mu_img.shape[0]):
        #        if i != j:
        #            distI = distanceI[i] - torch.sum((mu_img[i] - mu_att[j]) ** 2) + self.margin
        #            distI = torch.max(torch.FloatTensor([0]), distI)
        #            distI += torch.sum((torch.sqrt(logvar_img[i].exp()) - torch.sqrt(logvar_att[i].exp())) ** 2, dim=0)
        #            distI = torch.sqrt(distI)
        #            tripletsI.append(distI)

        #            distT = distanceT[i] - torch.sum((mu_att[i] - mu_img[j]) ** 2) + self.margin
        #            distT = torch.max(torch.FloatTensor([0]), distT)
        #            distT += torch.sum((torch.sqrt(logvar_att[i].exp()) - torch.sqrt(logvar_img[i].exp())) ** 2, dim=0)
        #            distT = torch.sqrt(distT)
        #            tripletsT.append(distT)

        #tripletI2t = sum(tripletsI)
        #tripletT2i = sum(tripletsT)

        #distance = tripletI2t + tripletT2i

        #distance = torch.FloatTensor([distance]).to(self.device)

        ##############################################
        # scale the loss terms according to the warmup
        # schedule
        ##############################################

        f1 = 1.0*(self.current_epoch - self.warmup['cross_reconstruction']['start_epoch'] )/(1.0*( self.warmup['cross_reconstruction']['end_epoch']- self.warmup['cross_reconstruction']['start_epoch']))
        f1 = f1*(1.0*self.warmup['cross_reconstruction']['factor'])
        cross_reconstruction_factor = torch.cuda.FloatTensor([min(max(f1,0),self.warmup['cross_reconstruction']['factor'])])

        f2 = 1.0 * (self.current_epoch - self.warmup['beta']['start_epoch']) / ( 1.0 * (self.warmup['beta']['end_epoch'] - self.warmup['beta']['start_epoch']))
        f2 = f2 * (1.0 * self.warmup['beta']['factor'])
        beta = torch.cuda.FloatTensor([min(max(f2, 0), self.warmup['beta']['factor'])])

        f3 = 1.0*(self.current_epoch - self.warmup['distance']['start_epoch'] )/(1.0*( self.warmup['distance']['end_epoch']- self.warmup['distance']['start_epoch']))
        f3 = f3*(1.0*self.warmup['distance']['factor'])
        distance_factor = torch.cuda.FloatTensor([min(max(f3,0),self.warmup['distance']['factor'])])

        ##############################################
        # Put the loss together and call the optimizer
        ##############################################

        self.optimizer.zero_grad()

        loss = reconstruction_loss - beta * KLD

        if cross_reconstruction_loss>0:
           loss += cross_reconstruction_factor*cross_reconstruction_loss
        if distance_factor >0:
           loss += distance_factor*distance

        loss.backward()

        self.optimizer.step()

        return loss.item(), reconstruction_loss, beta*KLD, cross_reconstruction_factor*cross_reconstruction_loss, distance_factor*distance

    def train_vae(self):

        elosses = []
        lossesR = []
        lossesK = []
        lossesC = []
        lossesD = []

        self.dataloader = data.DataLoader(self.dataset,batch_size=self.batch_size,shuffle= True,drop_last=True)#,num_workers = 4)

        #self.dataset.novelclasses =self.dataset.novelclasses.long().cuda()
        #self.dataset.seenclasses =self.dataset.seenclasses.long().cuda()

        #leave both statements
        self.train()
        self.fc_ft.train()
        self.fc_at.train()
        self.reparameterize_with_noise = True

        metricsI = []
        metricsT = []

        print('train for reconstruction')
        for epoch in range(0, self.nepoch ):
            self.train()

            self.current_epoch = epoch

            losses = []
            ilossesR = []
            ilossesK = []
            ilossesC = []
            ilossesD = []

            i=-1
            y = 0
            for iters in range(0, len(self.dataset), self.batch_size):

            #for iters in range(0, 1000, self.batch_size):
                i+=1

                features, attributes, idxs = self.dataset.next_batch(self.batch_size) #Si no és test treure la y

                data_from_modalities = [features, attributes.type(torch.FloatTensor)]

                for j in range(len(data_from_modalities)):
                    data_from_modalities[j] = data_from_modalities[j].to(self.device)
                    data_from_modalities[j].requires_grad = False

                loss, lossR, lossK, lossC, lossD = self.trainstep(data_from_modalities[0], data_from_modalities[1] )

                if i%10==0:

                    print('epoch ' + str(epoch) + ' | iter ' + str(i) + '\t'+
                    ' | loss ' +  str(loss))

                losses.append(loss)
                ilossesR.append(lossR)
                ilossesK.append(lossK)
                ilossesC.append(lossC)
                ilossesD.append(lossD)

                idxs = idxs.cpu()

                attributes = attributes.cpu()

                y += 1

            y = 0

            mean_loss = sum(losses)/len(losses)
            elosses.append(mean_loss)
            print('epoch ' + str(epoch) + 'Loss: ' + str(loss))

            lossesR.append(sum(ilossesR)/len(ilossesR))
            lossesK.append(sum(ilossesK)/len(ilossesK))
            lossesC.append(sum(ilossesC)/len(ilossesC))
            lossesD.append(sum(ilossesD)/len(ilossesD))

            for j in range(len(data_from_modalities)):
                    data_from_modalities[j] = data_from_modalities[j].cpu()

            print('Generating gallery set...')
            self.generate_gallery()

            print('Generating t-SNE plot...')
            z_imgs_embedded = TSNE(n_components=2).fit_transform(self.gallery_imgs_z.clone().cpu().detach())
            z_attrs_embedded = TSNE(n_components=2).fit_transform(self.gallery_attrs_z.clone().cpu().detach())

            plt.scatter(z_imgs_embedded[:,0], z_imgs_embedded[:,1], c = 'red')
            plt.scatter(z_attrs_embedded[:,0], z_attrs_embedded[:,1], c = 'blue')
            filename = 't-sne-plot-epoch'+str(epoch)+'.png'
            plt.savefig(filename)
            plt.clf()

            print('Evaluating retrieval...')
            metricsIepoch, metricsTepoch = self.retrieval()
            metricsI.append(metricsIepoch)
            metricsT.append(metricsTepoch)

            print('Evaluation Metrics for image retrieval')
            print("R@1: {}, R@5: {}, R@10: {}, R@50: {}, R@100: {}, MEDR: {}, MEANR: {}".format(metricsIepoch[0], metricsIepoch[1], metricsIepoch[2], metricsIepoch[3], metricsIepoch[4], metricsIepoch[5], metricsIepoch[6]))
            print('Evaluation Metrics for caption retrieval')
            print("R@1: {}, R@5: {}, R@10: {}, R@50: {}, R@100: {}, MEDR: {}, MEANR: {}".format(metricsTepoch[0], metricsTepoch[1], metricsTepoch[2], metricsTepoch[3], metricsTepoch[4], metricsTepoch[5], metricsTepoch[6]))

        # turn into evaluation mode:
        for key, value in self.encoder.items():
            self.encoder[key].eval()
        for key, value in self.decoder.items():
            self.decoder[key].eval()

        #Plot de les losses i desar-lo
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.plot(np.arange(self.nepoch), elosses)
        plt.plot(np.arange(self.nepoch), lossesR)
        plt.plot(np.arange(self.nepoch), lossesK)
        plt.plot(np.arange(self.nepoch), lossesC)
        plt.plot(np.arange(self.nepoch), lossesD)
        plt.legend()
        plt.show()
        plt.savefig('losses.png')
        plt.clf()

        #Plot de les metrics
        plt.xlabel('Epoch')
        plt.ylabel('Metric')
        plt.plot(np.arange(self.nepoch), metricsI)
        plt.plot(np.arange(self.nepoch), metricsT)
        plt.legend()
        plt.show()
        plt.savefig('metrics.png')
        plt.clf()

        return losses, metricsI, metricsT

    def retrieval(self):

        self.eval()

        def lda(self, x, y):
            distance = torch.sqrt(torch.sum((x[0] - y[0]) ** 2, dim=1) + \
                              torch.sum((torch.sqrt(x[1].exp()) - torch.sqrt(y[1].exp())) ** 2, dim=1))

            return distance

        #nbrsI = NearestNeighbors(n_neighbors=self.dataset.ntest, algorithm='auto').fit(self.gallery_imgs_z.cpu().detach().numpy())
        #nbrsI = NearestNeighbors(n_neighbors=1000, algorithm='auto').fit(self.gallery_imgs_z.cpu().detach().numpy())
        #nbrsT = NearestNeighbors(n_neighbors=self.dataset.ntest, algorithm='auto').fit(self.gallery_attrs_z.cpu().detach().numpy())
        #nbrsT = NearestNeighbors(n_neighbors=5000, algorithm='auto').fit(self.gallery_attrs_z.cpu().detach().numpy())

        distI_dict = {}
        distT_dict = {}

        ranksI = np.zeros((1,5*self.dataset.ntest))
        ranksT = np.zeros((1,self.dataset.ntest))

        for i in range(0, self.dataset.ntest):
        #for i in range(0, 500):

            mu_img = self.gallery_imgs_z[i,:].unsqueeze(0)
            mu_att = self.gallery_attrs_z[5*i:5*i + 5,:]

            #Add L2 norm and BatchNorm?
            #im_ft = F.normalize(im_ft, p=2, dim = 1)
            #im_ft = self.ft_bn(im_ft)

            #im_ft = self.fc_ft(im_ft)

            #attr = F.normalize(attr.type(torch.FloatTensor).to(self.device), p=2, dim=1)
            #attr = self.at_bn(attr)

            #attr = self.fc_at(attr)

            #mu_img, logvar_img = self.encoder['resnet_features'](im_ft)
            #z_from_img = self.reparameterize(mu_img, logvar_img)

            #mu_att, logvar_att = self.encoder['attributes'](attr.type(torch.FloatTensor).to(self.device))
            #z_from_att = self.reparameterize(mu_att, logvar_att)

            #img = [mu_img.cpu().detach().numpy(), logvar_img.cpu().detach().numpy()]
            #att = [mu_att.cpu().detach().numpy(), logvar_att.cpu().detach().numpy()]

            #distancesI, indicesI = nbrsI.kneighbors(mu_att.cpu().detach().numpy())
            #distancesT, indicesT = nbrsT.kneighbors(mu_img.cpu().detach().numpy())

            distancesI = distance.cdist(mu_att.cpu().detach().numpy(), self.gallery_imgs_z.cpu().detach().numpy(), 'correlation')
            distancesT = distance.cdist(mu_img.cpu().detach().numpy(), self.gallery_attrs_z.cpu().detach().numpy(), 'correlation')

            indicesI = np.argsort(distancesI)
            indicesT = np.argsort(distancesT[0,:])

            #Trobar l'índex de la distància mínima entre cada parella.

            for z in range(0,5):
                if len(indicesI[z] == i) != 0:
                    ranksI[:,(5*i) + z] = np.where(indicesI[z] == i)[0][0]
                else:
                    ranksI[:,(5*i) + z] = 1000

            if len(np.where((indicesT >= 5*i) & (indicesT <= ((5*i) + 4)))) != 0:
                ranksT[:,i] = np.where((indicesT >= 5*i) & (indicesT <= ((5*i) + 4)))[0][0]
            else:
                ranksT[:,i] = 1000

            '''
            for z in range(0,5):
                if len(np.where((indicesI[z] >= 5*i) & (indicesI[z] <= (5*i+4)))[0]) != 0:
                        ranksI[:,5*i + z] = np.where((indicesI[z] >= 5*i) & (indicesI[z] <= (5*i+4)))[0][0]
                else:
                    ranksI[:,5*i + z] = 5000

            if len(np.where((indicesT[0] > (5*i-1)) & (indicesT[0] < (5*i+5)))) != 0:
                ranksT[:,i] = np.where((indicesT[0] > (5*i-1)) & (indicesT[0] < (5*i+5)))[0][0]
            else:
                ranksT[:,i] = 1000

            '''
            '''
            for z in range(0,5):
                if len(np.where((indicesT[z] >= 5*i) & (indicesT[z] <= (5*i+4)))[0]) != 0:
                        ranksT[:,5*i + z] = np.where((indicesT[z] >= 5*i) & (indicesT[z] <= (5*i+4)))[0][0]
                else:
                    ranksT[:,5*i + z] = 5000

            if len(np.where(indicesI[0] == i)) != 0:
                ranksI[:,i] = np.where(indicesI[0] == i)[0][0]
            else:
                ranksI[:,i] = 1000

            '''

        r1im = 100.0 * len(np.where(ranksI < 1)[1]) / len(ranksI[0,:])
        r5im = 100.0 * len(np.where(ranksI < 5)[1]) / len(ranksI[0,:])
        r10im = 100.0 * len(np.where(ranksI < 10)[1]) / len(ranksI[0,:])
        r50im = 100.0 * len(np.where(ranksI < 50)[1]) / len(ranksI[0,:])
        r100im = 100.0 * len(np.where(ranksI < 100)[1]) / len(ranksI[0,:])

        r1t = 100.0 * len(np.where(ranksT < 1)[1]) / len(ranksT[0,:])
        r5t = 100.0 * len(np.where(ranksT < 5)[1]) / len(ranksT[0,:])
        r10t = 100.0 * len(np.where(ranksT < 10)[1]) / len(ranksT[0,:])
        r50t = 100.0 * len(np.where(ranksT < 50)[1]) / len(ranksT[0,:])
        r100t = 100.0 * len(np.where(ranksT < 100)[1]) / len(ranksT[0,:])

        medrI = np.floor(np.median(ranksI)) + 1
        meanrI = ranksI.mean() + 1

        medrT = np.floor(np.median(ranksT)) + 1
        meanrT = ranksT.mean() + 1

        metricsI = [r1im, r5im, r10im, r50im, r100im, medrI, meanrI]
        metricsT = [r1t, r5t, r10t, r50t, r100t, medrT, meanrT]

        return metricsI, metricsT

    def generate_gallery(self):

        self.eval()

        z_imgs = []
        z_vars_im = []
        z_attrs = []
        z_vars_att = []
        rec_imgs = []
        rec_attrs = []

        y = 0
        i=-1
        for iters in range(0, self.dataset.ntest, 50):
        #for iters in range(0, 500, 50):
            i+=1

            features, attributes, idxs = self.dataset.next_batch_test(50, y)
            idxs = idxs.cpu()

            data_from_modalities = [features, attributes.type(torch.FloatTensor)]

            for j in range(len(data_from_modalities)):
                data_from_modalities[j] = data_from_modalities[j].to(self.device)
                data_from_modalities[j].requires_grad = False
                if j== 0:
                    #Add L2 norm and BatchNorm?
                    data_from_modalities[j] = F.normalize(data_from_modalities[j], p=2, dim=1)
                    data_from_modalities[j] = self.ft_bn(data_from_modalities[j])
                    data_from_modalities[j] = self.fc_ft(data_from_modalities[j])
                elif j == 1:
                    #Add L2 norm and BatchNorm?
                    data_from_modalities[j] = F.normalize(data_from_modalities[j], p=2, dim=1)
                    data_from_modalities[j] = self.at_bn(data_from_modalities[j])
                    data_from_modalities[j] = self.fc_at(data_from_modalities[j])

            mu_img, logvar_img = self.encoder['resnet_features'](data_from_modalities[0])
            z_from_img = self.reparameterize(mu_img, logvar_img)

            mu_att, logvar_att = self.encoder['attributes'](data_from_modalities[1])
            z_from_att = self.reparameterize(mu_att, logvar_att)

            if y == 0:
                z_imgs = mu_img.cpu()
                z_vars_im = logvar_img.cpu()
                z_attrs = mu_att.cpu()
                z_vars_att = logvar_att.cpu()
            else:
                z_imgs = torch.cat((z_imgs.cpu(),mu_img.cpu()), dim = 0).cpu()
                z_vars_im = torch.cat((z_vars_im.cpu(),logvar_img.cpu()), dim = 0).cpu()
                z_attrs = torch.cat((z_attrs.cpu(),mu_att.cpu()), dim = 0).cpu()
                z_vars_att = torch.cat((z_vars_att.cpu(),logvar_att.cpu()), dim = 0).cpu()

            y = y + 1

            print('iter: '+str(iters))

        self.gallery_imgs_z = z_imgs.cpu()
        self.gallery_vars_im = z_vars_im.cpu()
        print(self.gallery_imgs_z.size())
        self.gallery_attrs_z = z_attrs.cpu()
        self.gallery_vars_att = z_vars_att.cpu()
        print(self.gallery_attrs_z.size())

    def train_classifier(self, show_plots=False):

        if self.num_shots > 0 :
            print('================  transfer features from test to train ==================')
            self.dataset.transfer_features(self.num_shots, num_queries='num_features')

        history = []  # stores accuracies

        cls_seenclasses = self.dataset.seenclasses
        cls_novelclasses = self.dataset.novelclasses

        train_seen_feat = self.dataset.data['train_seen']['resnet_features']
        train_seen_label = self.dataset.data['train_seen']['labels']

        novelclass_aux_data = self.dataset.novelclass_aux_data  # access as novelclass_aux_data['resnet_features'], novelclass_aux_data['attributes']
        seenclass_aux_data = self.dataset.seenclass_aux_data

        novel_corresponding_labels = self.dataset.novelclasses.long().to(self.device)
        seen_corresponding_labels = self.dataset.seenclasses.long().to(self.device)

        # The resnet_features for testing the classifier are loaded here
        novel_test_feat = self.dataset.data['test_unseen'][
            'resnet_features']  # self.dataset.test_novel_feature.to(self.device)
        seen_test_feat = self.dataset.data['test_seen'][
            'resnet_features']  # self.dataset.test_seen_feature.to(self.device)
        test_seen_label = self.dataset.data['test_seen']['labels']  # self.dataset.test_seen_label.to(self.device)
        test_novel_label = self.dataset.data['test_unseen']['labels']  # self.dataset.test_novel_label.to(self.device)

        train_unseen_feat = self.dataset.data['train_unseen']['resnet_features']
        train_unseen_label = self.dataset.data['train_unseen']['labels']

        # in ZSL mode:
        if self.generalized == False:
            # there are only 50 classes in ZSL (for CUB)
            # novel_corresponding_labels =list of all novel classes (as tensor)
            # test_novel_label = mapped to 0-49 in classifier function
            # those are used as targets, they have to be mapped to 0-49 right here:

            novel_corresponding_labels = self.map_label(novel_corresponding_labels, novel_corresponding_labels)

            if self.num_shots > 0:
                # not generalized and at least 1 shot means normal FSL setting (use only unseen classes)
                train_unseen_label = self.map_label(train_unseen_label, cls_novelclasses)

            # for FSL, we train_seen contains the unseen class examples
            # for ZSL, train seen label is not used
            # if self.num_shots>0:
            #    train_seen_label = self.map_label(train_seen_label,cls_novelclasses)

            test_novel_label = self.map_label(test_novel_label, cls_novelclasses)

            # map cls novelclasses last
            cls_novelclasses = self.map_label(cls_novelclasses, cls_novelclasses)

        if self.generalized:
            print('mode: gzsl')
            clf = LINEAR_LOGSOFTMAX(self.latent_size, self.num_classes)
        else:
            print('mode: zsl')
            clf = LINEAR_LOGSOFTMAX(self.latent_size, self.num_novel_classes)

        clf.apply(models.weights_init)

        with torch.no_grad():

            ####################################
            # preparing the test set
            # convert raw test data into z vectors
            ####################################

            self.reparameterize_with_noise = False

            mu1, var1 = self.encoder['resnet_features'](novel_test_feat)
            test_novel_X = self.reparameterize(mu1, var1).to(self.device).data
            test_novel_Y = test_novel_label.to(self.device)

            mu2, var2 = self.encoder['resnet_features'](seen_test_feat)
            test_seen_X = self.reparameterize(mu2, var2).to(self.device).data
            test_seen_Y = test_seen_label.to(self.device)

            ####################################
            # preparing the train set:
            # chose n random image features per
            # class. If n exceeds the number of
            # image features per class, duplicate
            # some. Next, convert them to
            # latent z features.
            ####################################

            self.reparameterize_with_noise = True

            def sample_train_data_on_sample_per_class_basis(features, label, sample_per_class):
                sample_per_class = int(sample_per_class)

                if sample_per_class != 0 and len(label) != 0:

                    classes = label.unique()

                    for i, s in enumerate(classes):

                        features_of_that_class = features[label == s, :]  # order of features and labels must coincide
                        # if number of selected features is smaller than the number of features we want per class:
                        multiplier = torch.ceil(torch.cuda.FloatTensor(
                            [max(1, sample_per_class / features_of_that_class.size(0))])).long().item()

                        features_of_that_class = features_of_that_class.repeat(multiplier, 1)

                        if i == 0:
                            features_to_return = features_of_that_class[:sample_per_class, :]
                            labels_to_return = s.repeat(sample_per_class)
                        else:
                            features_to_return = torch.cat(
                                (features_to_return, features_of_that_class[:sample_per_class, :]), dim=0)
                            labels_to_return = torch.cat((labels_to_return, s.repeat(sample_per_class)),
                                                         dim=0)

                    return features_to_return, labels_to_return
                else:
                    return torch.cuda.FloatTensor([]), torch.cuda.LongTensor([])

            # some of the following might be empty tensors if the specified number of
            # samples is zero :

            img_seen_feat,   img_seen_label   = sample_train_data_on_sample_per_class_basis(
                train_seen_feat,train_seen_label,self.img_seen_samples )

            img_unseen_feat, img_unseen_label = sample_train_data_on_sample_per_class_basis(
                train_unseen_feat, train_unseen_label, self.img_unseen_samples )

            att_unseen_feat, att_unseen_label = sample_train_data_on_sample_per_class_basis(
                    novelclass_aux_data,
                    novel_corresponding_labels,self.att_unseen_samples )

            att_seen_feat, att_seen_label = sample_train_data_on_sample_per_class_basis(
                seenclass_aux_data,
                seen_corresponding_labels, self.att_seen_samples)

            def convert_datapoints_to_z(features, encoder):
                if features.size(0) != 0:
                    mu_, logvar_ = encoder(features)
                    z = self.reparameterize(mu_, logvar_)
                    return z
                else:
                    return torch.cuda.FloatTensor([])

            z_seen_img   = convert_datapoints_to_z(img_seen_feat, self.encoder['resnet_features'])
            z_unseen_img = convert_datapoints_to_z(img_unseen_feat, self.encoder['resnet_features'])

            z_seen_att = convert_datapoints_to_z(att_seen_feat, self.encoder[self.auxiliary_data_source])
            z_unseen_att = convert_datapoints_to_z(att_unseen_feat, self.encoder[self.auxiliary_data_source])

            train_Z = [z_seen_img, z_unseen_img ,z_seen_att    ,z_unseen_att]
            train_L = [img_seen_label    , img_unseen_label,att_seen_label,att_unseen_label]

            # empty tensors are sorted out
            train_X = [train_Z[i] for i in range(len(train_Z)) if train_Z[i].size(0) != 0]
            train_Y = [train_L[i] for i in range(len(train_L)) if train_Z[i].size(0) != 0]

            train_X = torch.cat(train_X, dim=0)
            train_Y = torch.cat(train_Y, dim=0)

        ############################################################
        ##### initializing the classifier and train one epoch
        ############################################################

        cls = classifier.CLASSIFIER(clf, train_X, train_Y, test_seen_X, test_seen_Y, test_novel_X,
                                    test_novel_Y,
                                    cls_seenclasses, cls_novelclasses,
                                    self.num_classes, self.device, self.lr_cls, 0.5, 1,
                                    self.classifier_batch_size,
                                    self.generalized)

        for k in range(self.cls_train_epochs):
            if k > 0:
                if self.generalized:
                    cls.acc_seen, cls.acc_novel, cls.H = cls.fit()
                else:
                    cls.acc = cls.fit_zsl()

            if self.generalized:

                print('[%.1f]     novel=%.4f, seen=%.4f, h=%.4f , loss=%.4f' % (
                k, cls.acc_novel, cls.acc_seen, cls.H, cls.average_loss))

                history.append([torch.tensor(cls.acc_seen).item(), torch.tensor(cls.acc_novel).item(),
                                torch.tensor(cls.H).item()])

            else:
                print('[%.1f]  acc=%.4f ' % (k, cls.acc))
                history.append([0, torch.tensor(cls.acc).item(), 0])

        if self.generalized:
            return torch.tensor(cls.acc_seen).item(), torch.tensor(cls.acc_novel).item(), torch.tensor(
                cls.H).item(), history
        else:
            return 0, torch.tensor(cls.acc).item(), 0, history


import collections
import configparser
import hashlib
import os
import re
import zlib

git_folder = ".wyag" # ".git"

class GitRepository(object):
    """A git repository"""

    worktree = None
    gitdir = None
    conf = None

    def __init__(self, path, force=False):
        self.worktree = path
        self.gitdir = os.path.join(path, git_folder)

        if not (force or os.path.isdir(self.gitdir)):
            raise Exception("Not a Git repository {:s}".format(self.worktree))

        # read configuration file .git/config
        self.conf = configparser.ConfigParser()
        cf = repo_file(self, "config")

        if cf and os.path.exists(cf):
            self.conf.read([cf])
        elif not force:
            raise Exception("Configuration file is missing")

        if not force:
            vers = int(self.conf.get("core", "repositoryformatversion"))
            if vers != 0:
                raise Exception("Unsupported repositoryformatversion {}".format(vers))

def repo_path(repo, *path):
    """Comput path under repo's gitdir"""
    return os.path.join(repo.gitdir, *path)

def repo_file(repo, *path, mkdir=False):
    """Same as repo_path, but create dirname(*path) if absent. For example,
    repo_file(r, refs, remotes, origin, HEAD) will create .git/refs/remotes/origin."""

    if repo_dir(repo, *path[:-1], mkdir=mkdir):
        return repo_path(repo, *path)

def repo_dir(repo, *path, mkdir=False):
    """Same as repo_path, but mkdir *path if absent if mkdir."""

    path = repo_path(repo, *path)

    if os.path.exists(path):
        if os.path.isdir(path):
            return path
        else:
            raise Exception("Not a directory {}".format(path))

    if mkdir:
        os.makedirs(path)
        return path
    else:
        return None

def repo_create(path):
    """Create a new repository at path."""

    repo = GitRepository(path, True)

    # make sure path doesn't exist or is empty

    if os.path.exists(repo.worktree):
        if not os.path.isdir(repo.worktree):
            raise Exception("{} is not a directroy".format(path))
        if os.listdir(repo.worktree):
            raise Exception("{} is not empty".format(path))
    else:
        os.makedirs(repo.worktree)

    assert(repo_dir(repo, "branches", mkdir=True))
    assert(repo_dir(repo, "objects", mkdir=True))
    assert(repo_dir(repo, "refs", "tags", mkdir=True))
    assert(repo_dir(repo, "refs", "head", mkdir=True))

    # .git/description
    with open(repo_file(repo, "description"), "w") as f:
        f.write("Unnamed repository; edit this file 'description' to name the repository.\n")

    # .git/HEAD
    with open(repo_file(repo, "HEAD"), "w") as f:
        f.write("refs: refs/heads/master\n")

    with open(repo_file(repo, "config"), "w") as f:
        config = repo_default_config()
        config.write(f)

    return repo

def repo_default_config():
    ret = configparser.ConfigParser()

    ret.add_section("core")
    ret.set("core", "repositoryformatversion", "0")
    ret.set("core", "filemode", "false")
    ret.set("core", "bare", "false")

    return ret

def repo_find(path=".", required=True):
    path = os.path.realpath(path)

    if os.path.isdir(os.path.join(path, git_folder)):
        return GitRepository(path)

    # search recursive in parent directories
    parent = os.path.realpath(os.path.join(path, ".."))

    if parent == path:
        # re reached root

        if required:
            raise Exception("No git repository")
        else:
            return None

    return repo_find(parent, required)

class GitObject (object):

    repo = None

    def __init__(self, repo, data=None):
        self.repo = repo

        if data != None:
            self.deserialize(data)

    def serialize(self):
        """This function MUST be implemented by subclassses.

        It must read the obkect's contens from self.data, a byte string, and
        do whaterver it takes to conber it into a meaningful representation.
        What exactly that means depends on each subclass."""

        raise Exception("Unimplemented!")

    def deserialize(self, data):
        raise Exception("Unimplemented!")

def object_find(repo, name, fmt=None, follow=True):
    return name

def object_read(repo, sha):
    """Read object object_id from Git repository repo. Return a GitObject
    whose exact type depends on the object."""

    path = repo_file(repo, "objects", sha[:2], sha[2:])

    with open(path, "rb") as f:
        raw = zlib.decompress(f.read())

        # read object type
        x = raw.find(b' ')
        fmt = raw[0:x]

        # read and validate  object size
        y = raw.find(b'\x00', x)
        size = int(raw[x:y].decode("ascii"))
        if size != len(raw)-y-1:
            raise Exception("Malformed object {}: bad length".format(sha))

        # pick correct constructor
        if fmt==b'commit' : c=GitCommit
        elif fmt==b'tree' : c=GitTree
        elif fmt==b'tag' : c=GitTag
        elif fmt==b'blob' : c=GitBlob
        else:
            raise Exception("Unknown type {} for object {}".format(fmt.decode("ascii"), sha))

        return c(repo, raw[y+1:])

def object_write(obj, actually_write=True):
    # serialize object data
    data = obj.serialize()

    # add header
    result = obj.fmt + b' ' + str(len(data)).encode() + b'\x00' + data

    # get hash
    sha = hashlib.sha1(result).hexdigest()

    if actually_write:
        # compute path
        path = repo_file(obj.repo, "objects", sha[:2], sha[2:], mkdir=actually_write)

        with open(path, 'wb') as f:
            # compress and write
            f.write(zlib.compress(result))

    return sha

def object_hash(f, fmt, repo=None):
    data = f.read()

    # Choose constructor depnedo on object type foound in header
    if fmt == b'commit' : obj=GitCommit(repo, data)
    elif fmt == b'tree' : obj=GitTree(repo, data)
    elif fmt == b'tag' : obj=GitTag(repo, data)
    elif fmt == b'blob' : obj=GitBlob(repo, data)
    else:
        raise Exception("Unknown type {}".format(fmt))

    return object_write(obj, repo)

class GitBlob(GitObject):
    fmt=b'blob'

    def serialize(self):
        return self.blobdata

    def deserialize(self, data):
        self.blobdata = data

# Key-Value List with Message
def kvlm_parse(raw, start=0, dct=None):
    if not dct:
        dct = collections.OrderedDict()
        # you cannot declare the argument as dct=OrderedDict() or all
        # call to the function will endlessly grom the same dict

    # we search for the next space and the next newline
    spc = raw.find(b' ', start)
    nl = raw.find(b'\n', start)

    # id space appears before the newline, we have a keyword

    # base case
    # ============
    # If newline appears first(or there is no space at all, in which
    # case find returns -1), we assume a blank line. A blank line
    # means the remainder of the data is the message.
    if (spc < 0) or (nl < spc):
        assert(nl == start)
        dct[b''] = raw[start+1:]
        return dct

    # Recusive case
    # ================
    # we read a key-value pair and recurse for the next.
    key = raw[start:spc]

    # Find the end of the value. Continuatuon lines begin with a
    # space, so loop until we find a \n not followed by a space.
    end = start
    while True:
        end = raw.find(b'\n', end + 1)
        if raw[end + 1] != ord(' '):
            break

    # Grab the value
    # Also, drop the leading space on continuation lines
    value = raw[spc + 1:end].replace(b'\n ', b'\n')

    # do not overwrite existing data contents
    if key in dct:
        if type(dct[key]) == list:
            dct[key].append(value)
        else:
            dct[key] = [dct[key], value]
    else:
        dct[key] = value

    return kvlm_parse(raw, start=end+1, dct=dct)

def kvlm_serialize(kvlm):
    ret = b''

    # output fields
    for k in kvlm.keys():
        # skip message itself
        if k==b'':
            continue
        val = kvlm[k]
        # Normalize to a list
        if type(val) != list:
            val = [val]

        for v in val:
            ret += k + b' ' + v.replace(b'\n', b'\n ') + b'\n'

    # append message
    ret += b'\n' + kvlm[b'']

    return ret

class GitCommit(GitObject):
    fmt = b'commit'

    def deserialize(self, data):
        self.kvlm = kvlm_parse(data)

    def serialize(self):
        return kvlm_serialize(self.kvlm)

class GitTreeLeaf(object):
    def __init__(self, mode, path, sha):
        self.mode = mode
        self.path = path
        self.sha = sha

def tree_parse_one(raw, start=0):
    # find the space terminator of the mode
    x = raw.find(b' ', start)
    assert(x-start == 5 or x-start == 6)

    # read the mode
    mode = raw[start:x]

    # fin the NULL terminator of the path
    y = raw.find(b'\x00', x)

    # read the path
    path = raw[x+1:y]

    # read SHA as binary and convert to hex string
    sha = hex(
        int.from_bytes(
            raw[y+1:y+21], "big"))[2:] # hex add 0x in front, we need to remove that

    return y+21, GitTreeLeaf(mode, path, sha)

def tree_parse(raw):
    pos = 0
    max = len(raw)
    ret = list()
    while pos < max:
        pos, data = tree_parse_one(raw, pos)
        ret.append(data)

    return ret

def tree_serialize(obj):
    ret = b''
    for i in obj.items:
        ret += i.mode
        ret += b' '
        ret += i.path
        ret += b'\x00'
        sha = int(i.sha, 16)
        ret += sha.to_bytes(20, byteorder="big")
    return ret

class GitTree(GitObject):
    fmt = b'tree'

    def deserialize(self, data):
        self.items = tree_parse(data)

    def serialize(self):
        return tree_serialize(self)


from django.template import loader, Context,RequestContext
from django.http import HttpResponse
from django.shortcuts import render_to_response, get_object_or_404,redirect
from Deanery.models import Groups, Students
from django.shortcuts import render
from Deanery.forms import StudForm, GroupForm
from django.http import HttpResponseRedirect
from django.contrib.auth.decorators import login_required

@login_required
def archive(request):
    ls_groups = Groups.objects.all()
    data = dict(ls_groups=ls_groups)
    return render_to_response('tm_group.html',
        data,
        context_instance=RequestContext(request))

#    ls_groups = Groups.objects.all()
#    t=loader.get_template("tm_group.html")
#    c=Context({'ls_groups':ls_groups})
#    return HttpResponse(t.render(c))

@login_required
def students_in_this_group(request,group_id):
    current_group = get_object_or_404(Groups, id=group_id)
    variables = RequestContext(request, {
        'current_group': current_group,
        })
    return render_to_response(
        'tm_students.html',
        variables
    )

@login_required
def group_add(request):
    group = Groups()
    if request.method == 'POST':
        form = GroupForm(data=request.POST, instance=group)
        if form.is_valid():
            form.save()
            return HttpResponseRedirect('/Deanery/')
    else:
        form = GroupForm(instance=group)
        variables = RequestContext(request, {
        'type': 'Deanery',
        'form': form,
        })
    return render_to_response(
        'tm_group_add.html',
        variables
    )

@login_required
def group_delete(request, group_id):
    group = get_object_or_404(Groups, id=group_id)
    for Students in group.rel_students.all():
        Students.group = None
        Students.save()
    group.delete()
    return HttpResponseRedirect('/Deanery/')

@login_required
def group_edit(request, group_id):
    group = get_object_or_404(Groups, id=group_id)
    if request.method == 'POST':
        form = GroupForm(data=request.POST, instance=group)
        if form.is_valid():
            form.save()
            return HttpResponseRedirect('/Deanery/')
    else:
        form = GroupForm(instance=group)
    variables = RequestContext(request, {
        'obj': group,
        'type': 'Deanery',
        'form': form,
        })
    return render_to_response(
        'tm_group_edit.html',
        variables
    )

@login_required
def student_add(request):
    std = Students()
    if request.method == 'POST':
        form = StudForm(data=request.POST, instance=std)
        if form.is_valid():
            form.save()
            return HttpResponseRedirect(std.get_absolute_url())
    else:
        form = StudForm(instance=std)
    variables = RequestContext(request, {
        'type': 'Deanery',
        'form': form,
        })
    return render_to_response(
        'tm_group_add.html',
        variables
    )

@login_required
def student_view(request, student_id):
    std = get_object_or_404(Students, id=student_id)
    variables = RequestContext(request, {
        'student': std,
        })
    return render_to_response(
        'tm_std_view.html',
        variables
    )

@login_required
def student_edit(request, student_id):
    std = get_object_or_404(Students, id=student_id)
    if request.method == 'POST':
        form = StudForm(data=request.POST, instance=std)
        if form.is_valid():
            form.save()
            return HttpResponseRedirect(std.get_absolute_url())
    else:
        form = StudForm(instance=std)
        variables = RequestContext(request, {
        'obj': std,
        'type': 'Deanery',
        'form': form,
        })
    return render_to_response(
        'tm_group_edit.html',
        variables
    )

@login_required
def student_delete(request, student_id):
    std = get_object_or_404(Students, id=student_id)
    group_id = std.groups.id
    std.delete()
    return HttpResponseRedirect('/Deanery/students-in-this-group/%s/' % group_id)


#!/usr/bin/python

import xml.etree.ElementTree as xml
import random

team = [("president", 1)]
vices = random.randint(4,7)
v=0
while v < vices:
    v+=1
    team.append(('vice president %s' % v, 2 , v) )
    members = random.randint(2,10)
    me=0
    while me < members:
        me+=1
        team.append(('team member %s' % me , 3, v))
print(team)

model = xml.Element('mxGraphModel')
root = xml.Element('root')
model.append(root)
m=0
zero = xml.Element('mxCell',id='0')
one = xml.Element('mxCell',id='1', parent='0')
root.append(zero)
root.append(one)

for member in team:
    if member[1] == 1:
        newCell=xml.Element('mxCell', id=member[0]+'_vertex', value=member[0], vertex='1', parent='1')
        newMxGeometry=xml.Element('mxGeometry', x='0', y='0', width='80', height='70')
        newMxGeometry.set('as','geometry')
        newCell.append(newMxGeometry)

    elif member[1] == 2:
        newCell=xml.Element('mxCell', id=member[0]+'_vertex', value=member[0], vertex='1', parent='1')
        newMxGeometry=xml.Element('mxGeometry', x='0', y='0', width='70', height='60')
        newMxGeometry.set('as','geometry')
        newCell.append(newMxGeometry)
        newEdge=xml.Element('mxCell', id=member[0]+"_edge",
                            value='', edge="1", parent="1", source='president_vertex', target=member[0]+'_vertex')
        newEdgeGeo=xml.Element('mxGeometry', relative='1')
        newEdgeGeo.set('as','geometry')
        newEdge.append(newEdgeGeo)
        root.append(newEdge)
    elif member[1] == 3:
        m+=1
        newCell=xml.Element('mxCell', id=member[0]+'-'+str(m)+'_vertex', value=member[0],
                            vertex='1', parent='1')
        newMxGeometry=xml.Element('mxGeometry', x='0', y='0', width='60', height='50')
        newMxGeometry.set('as','geometry')
        newCell.append(newMxGeometry)
        newEdge=xml.Element('mxCell', id=member[0]+'-'+str(m)
                            +'_edge', value='', edge="1", parent="1",
                            target=member[0]+'-'+str(m)+'_vertex', source='vice president '+str(member[2])+'_vertex')
        newEdgeGeo=xml.Element('mxGeometry', relative='1')
        newEdgeGeo.set('as','geometry')
        newEdge.append(newEdgeGeo)
        root.append(newEdge)
    root.append(newCell)

file = open("example.xml", 'w')
xml.ElementTree(model).write(file)
file.close()


# -*- coding: utf-8 -*-
"""
Created on Tue Apr 21 14:31:58 2020

@author: ninjaac
"""
from keras.datasets import mnist
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

(x_train,y_train), (x_test,y_test)=mnist.load_data()
some_digit=x_train[0]
#plt.imshow(some_digit,cmap = plt.cm.binary)
#plt.show()
y_train_5=(y_train==5)
print(np.shape(x_train))
x_train=x_train.reshape((60000,784))
print(np.shape(x_train))
# select a classifier and fit the model
from sklearn.linear_model import SGDClassifier
sgd=SGDClassifier(random_state=42)
sgd.fit(x_train,y_train_5)

#model performance mesure
from sklearn.model_selection import cross_val_score
score=cross_val_score(sgd,x_train,y_train_5,cv=5,scoring='accuracy')
print(score)

# more better prediction will be found by presicion and recall score
from sklearn.model_selection import cross_val_predict
y_train_pred = cross_val_predict(sgd, x_train, y_train_5, cv=3)
from sklearn.metrics  import precision_score,recall_score,f1_score
presicion=precision_score(y_train_5,y_train_pred)
print(f"precision {presicion}")
recall_sc=recall_score(y_train_5,y_train_pred)
print(f"Recall value {recall_sc}")
f1=f1_score(y_train_5,y_train_pred)
print(f" fi score value {f1}")
#presicion recall curve
y_scores = cross_val_predict(sgd, X_train, y_train_5, cv=3,
                             method="decision_function")
from sklearn.metrics import precision_recall_curve
precisions,recalls,thresholds=precision_recall_curve(y_train_5,y_scores)
def plot_precision_recall_vs_threshold(precisions, recalls, thresholds):
    plt.plot(thresholds, precisions[:-1], "b--", label="Precision")
    plt.plot(thresholds, recalls[:-1], "g-", label="Recall")
    plt.legend()
    plt.show()

plot_precision_recall_vs_threshold(precisions, recalls, thresholds)


'''Train CIFAR10 with PyTorch.'''
from __future__ import print_function

from torch.autograd import Variable
import sys
import torch
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn

import torchvision
import torchvision.transforms as transforms

import os
import argparse
import time
import numpy as np
import models
import datasets
import math

from utils import *

parser = argparse.ArgumentParser(description='PyTorch Seen Testing Category Training')
parser.add_argument('--lr', default=3e-4, type=float, help='learning rate')
parser.add_argument('--log_dir', default='log/', type=str,
                    help='log save path')
parser.add_argument('--model_dir', default='checkpoint/', type=str,
                    help='model save path')
parser.add_argument('--test_epoch', default=1, type=int,
                    metavar='E', help='test every N epochs')
parser.add_argument('--test-only', action='store_true', help='test only')
parser.add_argument('--low-dim', default=128, type=int,
                    metavar='D', help='feature dimension')
parser.add_argument('--batch-t', default=0.1, type=float,
                    metavar='T', help='temperature parameter for softmax')
parser.add_argument('--batch-m', default=1, type=float,
                    metavar='N', help='m for negative sum')
parser.add_argument('--batch-size', default=128, type=int,
                    metavar='B', help='training batch size')
parser.add_argument('--gpu', default='0,1,2,3', type=str,
                      help='gpu device ids for CUDA_VISIBLE_DEVICES')
parser.add_argument('--logistic_epochs', default=1000, type=int, metavar='B', help='training batch size')
parser.add_argument('--logistic_batch_size', default=128, type=int, metavar='B', help='training batch size')

parser.add_argument('--dataset', default='cifar10',  help='cifar10, cifar100')
parser.add_argument('--trial', type=int, help='trial')
parser.add_argument('--resnet', default='resnet18',  help='resnet18, resnet34, resnet50, resnet101')
parser.add_argument('--adv', default=False, action='store_true', help='adversarial exmaple')
parser.add_argument('--eps', default=0.03, type=float, help='eps for adversarial')
parser.add_argument('--bn_adv_momentum', default=0.01, type=float, help='eps for adversarial')
parser.add_argument('--alpha', default=1.0, type=float, help='stregnth for regularization')
parser.add_argument('--debug', default=False, action='store_true', help='test_both_adv')
args = parser.parse_args()

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

dataset = args.dataset

class LogisticRegression(nn.Module):

    def __init__(self, n_features, n_classes):
        super(LogisticRegression, self).__init__()

        self.model = nn.Linear(n_features, n_classes)

    def forward(self, x):
        return self.model(x)

if args.dataset =='cifar10' or args.dataset =='cifar100':
    img_size = 32
    pool_len = 4

log_dir = args.log_dir  + args.dataset + '_eval_log/'
test_epoch = args.test_epoch
if not os.path.isdir(log_dir):
    os.makedirs(log_dir)

if not os.path.isdir(args.model_dir):
    os.makedirs(args.model_dir)

if not os.path.isdir(args.model_dir + '/' + dataset + '_eval'):
    os.makedirs(args.model_dir + '/' + dataset + '_eval')

suffix = args.dataset + '_{}_batch_{}_embed_dim_{}'.format(args.resnet, args.batch_size, args.low_dim)

if args.adv:
    suffix = suffix + '_adv_eps_{}_alpha_{}'.format(args.eps, args.alpha)
    suffix = suffix + '_bn_adv_momentum_{}_trial_{}'.format(args.bn_adv_momentum, args.trial)
else:
    suffix = suffix + '_trial_{}'.format(args.trial)

print(suffix)
# log the output
test_log_file = open(log_dir + suffix + '.txt', "w")

device = 'cuda' if torch.cuda.is_available() else 'cpu'
best_acc = 0  # best test accuracy
start_epoch = 0  # start from epoch 0 or last checkpoint epoch

# Data Preparation
print('==> Preparing data..')

transform = transforms.Compose([
    transforms.Resize(size=img_size),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

root='../datasets'
if args.dataset == "cifar10":
    train_dataset = torchvision.datasets.CIFAR10(
        root, train=True, download=True, transform=transform)
    test_dataset = torchvision.datasets.CIFAR10(
        root, train=False, download=True, transform=transform
    )
elif args.dataset == "cifar100":
    train_dataset = torchvision.datasets.CIFAR100(
        root, train=True, download=True, transform=transform
    )
    test_dataset = torchvision.datasets.CIFAR100(
        root, train=False, download=True, transform=transform
    )
else:
    raise NotImplementedError

train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=args.logistic_batch_size,
        shuffle=True,
        drop_last=True,
        num_workers=4,
    )

test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=args.logistic_batch_size,
    shuffle=False,
    drop_last=True,
    num_workers=4,
)

ndata = train_dataset.__len__()

print('==> Building model..')
if args.adv:
    net = models.__dict__[args.resnet + '_cifar'](pool_len = pool_len, low_dim=args.low_dim, bn_adv_flag=True, bn_adv_momentum=args.bn_adv_momentum)
else:
    net = models.__dict__[args.resnet + '_cifar'](pool_len = pool_len, low_dim=args.low_dim, bn_adv_flag=False)

# define leminiscate: inner product within each mini-batch (Ours)

if device == 'cuda':
    net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
    cudnn.benchmark = True

# Load checkpoint.
model_path = args.model_dir + '/' + dataset + '/' + suffix + '_best.t'
print('==> Load pretrained model {}'.format(model_path))

assert os.path.isdir(args.model_dir), 'Error: no checkpoint directory found!'
checkpoint = torch.load(model_path)
net.load_state_dict(checkpoint['net'])
net.to(device)
net.eval()

if args.dataset == "cifar100":
    n_classes = 100 # stl-10
else:
    n_classes = 10
model = LogisticRegression(args.low_dim, n_classes)
model = model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
criterion = torch.nn.CrossEntropyLoss()

def train(loader, net, model, criterion, optimizer):
    print('\nEpoch: %d' % epoch)
    # switch to train mode
    net.eval()
    model.train()
    loss_epoch = 0
    accuracy_epoch = 0
    for batch_idx, (x, y) in enumerate(loader):
        optimizer.zero_grad()
        x = x.to(device)
        y = y.to(device)
        with torch.no_grad():
            feat = net(x)
        output = model(feat)
        loss = criterion(output, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        predicted = output.argmax(1)
        acc = (predicted == y).sum().item() / y.size(0)
        accuracy_epoch += acc

        if args.debug:
            break

    return loss_epoch, accuracy_epoch

def test(loader, net, model, criterion, optimizer):
    loss_epoch = 0
    accuracy_epoch = 0
    model.eval()
    net.eval
    for step, (x, y) in enumerate(loader):
        optimizer.zero_grad()
        x = x.to(device)
        y = y.to(device)
        with torch.no_grad():
            feat = net(x)
            output = model(feat)
            loss = criterion(output, y)
            predicted = output.argmax(1)
            acc = (predicted == y).sum().item() / y.size(0)
            accuracy_epoch += acc
            loss_epoch += loss.item()
        if args.debug:
            break
    return loss_epoch, accuracy_epoch

best_acc = 0
for epoch in range(args.logistic_epochs):
    loss_epoch, accuracy_epoch = train(train_loader, net, model, criterion, optimizer)
    print("Train Epoch [{}]\t Loss: {}\t Accuracy: {}".format(epoch, loss_epoch / len(train_loader), accuracy_epoch / len(train_loader)), file = test_log_file)
    print("Train Epoch [{}]\t Loss: {}\t Accuracy: {}".format(epoch, loss_epoch / len(train_loader), accuracy_epoch / len(train_loader)))
    test_log_file.flush()
    # final testing
    test_loss_epoch, test_accuracy_epoch = test(test_loader, net, model, criterion, optimizer)
    test_current_acc = test_accuracy_epoch / len(test_loader)
    if test_current_acc > best_acc:
        best_acc = test_current_acc
        if not os.path.isdir(args.model_dir):
            os.mkdir(args.model_dir)
        torch.save(net, args.model_dir + '/' + dataset + '_eval/' + suffix + '_eval_best.t')
    print("Test Epoch [{}]\t Loss: {}\t Accuracy: {}\t Best Accuracy: {}".format(epoch, test_loss_epoch / len(test_loader), test_current_acc, best_acc), file = test_log_file)
    print("Test Epoch [{}]\t Loss: {}\t Accuracy: {}\t Best Accuracy: {}".format(epoch, test_loss_epoch / len(test_loader), test_current_acc, best_acc))
    test_log_file.flush()

    if args.debug:
        break

print("Final \t Best Accuracy: {}".format(epoch, best_acc), file = test_log_file)
test_log_file.flush()


import datetime, json, logging, os, codecs, re
import unicodecsv as csv
import textmining, mediacloud, stopwords
from ConfigParser import SafeConfigParser

ROWS_PER_QUERY = 500000

# Load config data
parser = SafeConfigParser()
parser.read('config.txt')
MY_API_KEY = parser.get('API','MY_API_KEY')
mc = mediacloud.api.AdminMediaCloud(MY_API_KEY) #AdminMediaCloud, rather than MediaCloud

logging.basicConfig(level=logging.DEBUG)
logging.info("-----------------------------------------------------------------")
logging.info("Starting QB data gathering")

# build stopwords
my_stopwords = [word.lower() for word in stopwords.getStopWords()]
qb_table = csv.reader(codecs.open('qb-table.csv', 'r', 'utf-8'))
qb_table.next()
team_stopwords = []
qb_stopwords = []
for row in qb_table:
    [ team_stopwords.append(word.lower()) for word in row[0].split() ]
    [ qb_stopwords.append(word.lower()) for word in row[1].split() ]
logging.debug(" Added qb names to stopwords: %s" % qb_stopwords)
logging.debug(" Added team names to stopwords: %s" % team_stopwords)
my_stopwords = my_stopwords + qb_stopwords + team_stopwords

# load media sources
m = codecs.open('sources.csv','r','utf-8')
media_reader = csv.reader(m)
media = [x[1] for x in media_reader][1:]
media_id_str = " ".join(media)
logging.info("Searching in %d media" % len(media))
logging.debug("media ids = %s" % media_id_str)

def fetch_corpus_from_mc(team,qb): #MC query, returns list of words
    '''
    Query MC for coverage of the QB specified, return a large string corpus
    '''
    more = True
    start = 0
    while more:
        logging.debug('    starting at %d' % start)
        sentences = mc.sentenceList(solr_query=str('"'+qb+'"'),
            solr_filter=[mc.publish_date_query(datetime.date(2015,9,9), datetime.date(2016,1,4)),
                         '+media_id:('+media_id_str+')'], rows = ROWS_PER_QUERY, start=0)
        more = len(sentences['response']['docs'])==ROWS_PER_QUERY
        start = start + ROWS_PER_QUERY
    logging.info('    found %d sentences',len(sentences['response']['docs']))
    response = sentences['response']
    docs = response['docs']
    logging.info('  done')
    return " ".join([d['sentence'] for d in docs])

def load_qb_corpus(team, qb):
    '''
    Grab the corpus for a QB from local data, or from MC if we don't have it yet
    '''
    corpora_dir = os.path.join('data','corpora')
    qb_corpora_path = os.path.join(corpora_dir,qb+'.txt')
    if not os.path.exists(corpora_dir):
        os.makedirs(corpora_dir)
    if os.path.exists(qb_corpora_path):
        f = codecs.open(qb_corpora_path, 'r', encoding='utf-8')
        corpus = f.read()
    else:
        corpus = fetch_corpus_from_mc(team,qb)
        f = codecs.open(qb_corpora_path, 'w', encoding='utf-8')
        f.write(corpus)
    return corpus

def tokenize_and_remove_stopwords(document):
    '''
    Callback to help the TDM creation via the textmining package
    '''
    document = document.lower() # do everything in lowercase
    document = re.sub('[^a-z]', ' ', document)  # remove non-alpha-numeric chars
    words = document.strip().split()
    before_word_count = len(words)
    # Remove stopwords
    words = [word for word in words if word not in my_stopwords]
    after_word_count = len(words)
    logging.debug("    %d (removed %d)" % (after_word_count,before_word_count-after_word_count))
    return words

def write_csv(cols, doc_iterator, filename):
    '''
    Write the results to a CSV in a human-usable format
    '''
    words = doc_iterator.next()
    word_freqs = [ row for row in doc_iterator ]
    corpus_word_count = [sum(row) for row in word_freqs]
    output_csv = csv.writer( codecs.open( os.path.join('data',filename), 'wb', 'utf-8') )
    output_csv.writerow(['word']+cols+[c+" pct" for c in cols])
    for idx in range(0,len(words)):
        word_counts = [ r[idx] for r in word_freqs ]
        normalized_word_counts = [ float(r[idx])/float(corpus_word_count[i]) for i,r in enumerate(word_freqs) ]
        output_csv.writerow([words[idx]]+word_counts+normalized_word_counts)

def get_and_write_data():
    '''
    Call this to collect and output the data
    '''
    corpus = {}
    count_corpus = {}
    white_doc = ""
    other_doc = ""
    logging.info('Loading sentences...')
    tdm_names = []
    tdm = textmining.TermDocumentMatrix(tokenize_and_remove_stopwords)
    race_tdm_names = []
    race_tdm = textmining.TermDocumentMatrix(tokenize_and_remove_stopwords)
    qb_table = csv.reader(codecs.open('qb-table.csv', 'r', 'utf-8'))
    qb_table.next()
    for row in qb_table:
        team = row[0]
        qb = row[1]
        race = row[2]
        logging.info('  %s (%s) - %s' % (qb,team,race))
        file_label = str(qb+' ('+team+')')
        qb_corpus = load_qb_corpus(team,qb)
        tdm.add_doc(qb_corpus)
        tdm_names.append(qb)
        if race == 'white':
            white_doc += qb_corpus + " "
        else:
            other_doc += qb_corpus + " "
    logging.info('done')
    # write the results
    logging.info('Writing results...')
    write_csv(tdm_names, tdm.rows(cutoff=1), 'word_freq_by_quarterback.csv')
    logging.info("Building white corpus")
    race_tdm.add_doc(white_doc)
    logging.info("Building non-white corpus")
    race_tdm.add_doc(other_doc)
    write_csv(['white','other'], race_tdm.rows(cutoff=1), 'word_freq_by_race.csv')
    logging.info('done')

if __name__ == "__main__":
    get_and_write_data()


from paging.paginators import EndlessPaginator, BetterPaginator
from jingo import register

@register.function
def paginate(request, queryset_or_list, per_page=25, endless=True, range_gap=2):
    if endless:
        paginator_class = EndlessPaginator
    else:
        paginator_class = BetterPaginator

    paginator = paginator_class(queryset_or_list, per_page)

    query_dict = request.GET.copy()
    if 'page' in query_dict:
        del query_dict['page']

    try:
        page = int(request.GET.get('page', 1))
    except (ValueError, TypeError):
        page = 1
    if page < 1:
        page = 1

    context = {
        'query_string': query_dict.urlencode(),
        'paginator': paginator.get_context(page, range_gap=range_gap),
    }
    return context


from mrcnn import model as modellib, visualize, utils, config
import numpy as np
import argparse
import imgaug
import os
import statistics
from xml.etree import ElementTree

class myMaskRCNNConfig(config.Config):
    # Give the configuration a recognizable name
    NAME = "MaskRCNN_config"

    # Set the number of GPUs to use along with the number of images per GPU
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    # Number of classes (we would normally add +1 for the background)
    NUM_CLASSES = 6 + 1

    # Number of training steps per epoch
    STEPS_PER_EPOCH = 300

    # Learning rate
    LEARNING_RATE = 0.001

    # Skip detections with < 90% confidence
    DETECTION_MIN_CONFIDENCE = 0.3

    # Setting max ground truth instances
    MAX_GT_INSTANCES = 10

    # BACKBONE                       resnet101
    # BACKBONE_STRIDES               [4, 8, 16, 32, 64]
    # BATCH_SIZE                     1
    # BBOX_STD_DEV                   [0.1 0.1 0.2 0.2]
    # COMPUTE_BACKBONE_SHAPE         None
    DETECTION_MAX_INSTANCES = 10
    # DETECTION_NMS_THRESHOLD        0.3
    # FPN_CLASSIF_FC_LAYERS_SIZE     1024
    # GPU_COUNT                      1
    # GRADIENT_CLIP_NORM             5.0
    # IMAGES_PER_GPU                 1
    # IMAGE_CHANNEL_COUNT            3

    IMAGE_MAX_DIM = 128
    # IMAGE_META_SIZE                19
    IMAGE_MIN_DIM = 128
    # IMAGE_MIN_SCALE                0
    IMAGE_RESIZE_MODE = "crop"
    IMAGE_SHAPE = [132, 132, 3]

    # LEARNING_MOMENTUM              0.9
    # LEARNING_RATE = 0.003
    LOSS_WEIGHTS = {'rpn_class_loss': 1.0, 'rpn_bbox_loss': 1.0, 'mrcnn_class_loss': 1.0, 'mrcnn_bbox_loss': 1.0, 'mrcnn_mask_loss': 0.01}
    # MASK_POOL_SIZE                 14
    # MASK_SHAPE                     [28, 28]
    # MAX_GT_INSTANCES               10

    # MEAN_PIXEL = [212.0, 47.0, 17.0]  # for
    MEAN_PIXEL = [24.9, 128.3, 235.6]  # for D1
    # MINI_MASK_SHAPE                (56, 56)
    # NAME                           MaskRCNN_config
    # NUM_CLASSES                    7
    # POOL_SIZE                      7

    POST_NMS_ROIS_INFERENCE = 10 #100
    POST_NMS_ROIS_TRAINING = 200

    PRE_NMS_LIMIT = 600
    # ROI_POSITIVE_RATIO             0.33
    # RPN_ANCHOR_RATIOS = [1]

    RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)
    # RPN_ANCHOR_STRIDE              1
    # RPN_BBOX_STD_DEV               [0.1 0.1 0.2 0.2]
    # RPN_NMS_THRESHOLD              0.7

    RPN_TRAIN_ANCHORS_PER_IMAGE = 64
    # STEPS_PER_EPOCH                131
    # TOP_DOWN_PYRAMID_SIZE          256
    # TRAIN_BN                       False
    TRAIN_ROIS_PER_IMAGE = 66
    USE_MINI_MASK = False
    # USE_RPN_ROIS                   True
    VALIDATION_STEPS = 100
    # WEIGHT_DECAY                   0.0001

# Define the prediction configuration
class PredictionConfig(config.Config):
    # Define the name of the configuration
    NAME = "galaxy inference"

    # Number of classes (background + 6 galaxy classes)
    NUM_CLASSES = 6 + 1

    # Simplify GPU config
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    # Minimum probability value to accept a detected instance
    # ROIs below this threshold are skipped
    DETECTION_MIN_CONFIDENCE = 0.0

    # Non-maximum suppression threshold for detection
    DETECTION_NMS_THRESHOLD = 0.0

class GalaxyDataset(utils.Dataset):
    # Load the dataset definitions
    def load_dataset(self, dataset_dir, is_train=True):

        # Add classes. We have only one class to add.
        self.add_class("dataset", 1, "1_1")
        self.add_class("dataset", 2, "1_2")
        self.add_class("dataset", 3, "1_3")
        self.add_class("dataset", 4, "2_2")
        self.add_class("dataset", 5, "2_3")
        self.add_class("dataset", 6, "3_3")

        # Define data locations for images and annotations
        images_dir = dataset_dir + 'images/'
        annotations_dir = dataset_dir + 'annots/'

        # Iterate through all files in the folder to add class, images and annotaions
        for filename in os.listdir(images_dir):
            if filename[0] != 'F':
                continue

            # Extract image id
            image_id = filename[:-4]

            # Setting image file
            img_path = images_dir + filename

            # Setting annotations file
            ann_path = annotations_dir + image_id + '.xml'

            # Adding images and annotations to dataset
            self.add_image('dataset', image_id=image_id, path=img_path, annotation=ann_path)

    # Extract bounding boxes from an annotation file
    def extract_boxes(self, filename):
        # Load and parse the file
        tree = ElementTree.parse(filename)

        # Get the root of the document
        root = tree.getroot()

        # Extract each bounding box
        boxes = list()

        for box in root.findall('.//object'):
            xmin = int(box.find('bndbox//xmin').text)
            ymin = int(box.find('bndbox//ymin').text)
            xmax = int(box.find('bndbox//xmax').text)
            ymax = int(box.find('bndbox//ymax').text)
            name = box.find('name').text
            coors = [xmin, ymin, xmax, ymax, name]
            boxes.append(coors)

        # Extract image dimensions
        width = int(root.find('.//size/width').text)
        height = int(root.find('.//size/height').text)

        return boxes, width, height

    # Load the masks for an image
    def load_mask(self, image_id):
        """
            Generate instance masks for an image.
            Returns:
                masks: A bool array of shape [height, width, instance count] with one mask per instance.
                class_ids: a 1D array of class IDs of the instance masks.
        """
        # Get details of image
        info = self.image_info[image_id]

        # Define annotation  file location
        path = info['annotation']

        # Load XML
        boxes, w, h = self.extract_boxes(path)

        # Create one array for all masks, each on a different channel
        masks = np.zeros([h, w, len(boxes)], dtype='uint8')

        # Create masks
        class_ids = list()
        for i in range(len(boxes)):
            box = boxes[i]
            row_s, row_e = box[1], box[3]
            col_s, col_e = box[0], box[2]
            masks[row_s:row_e, col_s:col_e, i] = 1
            class_ids.append(self.class_names.index(box[4]))
        return masks, np.asarray(class_ids, dtype='int32')

    # Load an image reference
    def image_reference(self, image_id):
        """Return the path of the image."""
        info = self.image_info[image_id]
        print(info)
        return info['path']

# Calculate the mAP for a model on a given dataset
def evaluate_model(dataset, model, cfg):
    APs = list()

    for image_id in dataset.image_ids:
        # Load image, bounding boxes and masks for the image id
        image, image_meta, gt_class_id, gt_bbox, gt_mask = modellib.load_image_gt(dataset, cfg, image_id,
                                                                                use_mini_mask=False)
        # Convert pixel values (e.g. center)
        scaled_image = modellib.mold_image(image, cfg)

        # Convert image into one sample
        sample = np.expand_dims(scaled_image, 0)

        # Make prediction
        yhat = model.detect(sample, verbose=0)

        # Extract results for first sample
        r = yhat[0]

        # Calculate statistics, including AP
        AP, _, _, _ = utils.compute_ap(gt_bbox, gt_class_id, gt_mask, r["rois"], r["class_ids"], r["scores"],
                                       r['masks'])

        # Store
        APs.append(AP)

    # Calculate the mean AP across all images
    mAP = statistics.mean(APs)
    return mAP

if __name__ == '__main__':
    # Get command to decide between train and test
    parser = argparse.ArgumentParser(
                        description='Train Mask R-CNN on MS COCO.')
    parser.add_argument("command",
                        metavar="<command>",
                        help="'train' or 'test'")
    parser.add_argument('--weights', required=False,
                        default='mask_rcnn_.1577911080.033216.h5',
                        metavar="<weights>",
                        help='.h5 weights file to test ')

    args = parser.parse_args()
    print("Command: ", args.command)
    print("Weights: ", args.weights)

    # Configure network
    #if args.command == 'train':
    config = myMaskRCNNConfig()
    #else:
    #    config = PredictionConfig()
    config.display()

    # Prepare train set
    train_set = GalaxyDataset()
    train_set.load_dataset('D1_train/', is_train=True)
    train_set.prepare()
    print('Train: %d' % len(train_set.image_ids))

    # Prepare test/val set
    test_set = GalaxyDataset()
    test_set.load_dataset('D1_test/', is_train=False)
    test_set.prepare()
    print('Test: %d' % len(test_set.image_ids))

    if args.command == "train":
        print("Loading Mask R-CNN model...")
        model = modellib.MaskRCNN(mode="training", config=config, model_dir='./')

        # Load the weights
        model.load_weights('../../mask_rcnn_coco.h5',
                           by_name=True,
                           exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"])

        # Horizontal flip
        augmentation = imgaug.augmenters.Fliplr(0.5)

        # Training - Stage 1
        print("Training network heads...")
        model.train(train_set, test_set,
                    learning_rate=config.LEARNING_RATE,
                    epochs=80,
                    layers='heads',
                    augmentation=augmentation)

        # Training - Stage 2
        # Fine tune layers from ResNet stage 4 and up
        print("Fine tune ResNet stage 4 and up")
        model.train(train_set, test_set,
                    learning_rate=config.LEARNING_RATE,
                    epochs=120,
                    layers='4+',
                    augmentation=augmentation)

        # Training - Stage 3
        # Fine tune all layers
        print("Fine tune all layers")
        model.train(train_set, test_set,
                    learning_rate=config.LEARNING_RATE / 10,
                    epochs=160,
                    layers='all',
                    augmentation=augmentation)

        history = model.keras_model.history.history

        # Save final weights
        # import time
        # model_path = 'mask_rcnn_' + '.' + str(time.time()) + '.h5'
        # model.keras_model.save_weights(model_path)
    else:
        model = modellib.MaskRCNN(mode="inference", config=config, model_dir='./')
        model.load_weights(args.weights[:], by_name=True)

        if 0:
            # Evaluate model on training dataset
            train_mAP = evaluate_model(train_set, model, config)
            print("Train mAP: %.3f" % train_mAP)
            # Evaluate model on test dataset
            test_mAP = evaluate_model(test_set, model, config)
            print("Test mAP: %.3f" % test_mAP)

        if 0:
            from keras.preprocessing.image import load_img
            from keras.preprocessing.image import img_to_array

            # Loading the model in the inference mode
            model = modellib.MaskRCNN(mode="inference", config=config, model_dir='./')
            # Loading the trained weights o the custom dataset
            # model.load_weights(model_path, by_name=True)
            img = load_img("D4_test/images/FIRSTJ111721.0+342714_infraredctmask.png")
            img = img_to_array(img)
            # Detecting objects in the image
            result = model.detect([img])
            print(result)

        if 1:
            for image_id in range(0, 100, 3):
                image, image_meta, gt_class_id, gt_bbox, gt_mask = modellib.load_image_gt(test_set, config, image_id,
                                                                                use_mini_mask=False)
                info = test_set.image_info[image_id]
                print("image ID: {}.{} ({}) {}".format(info["source"], info["id"], image_id,
                                                       test_set.image_reference(image_id)))

                visualize.display_instances(image, gt_bbox, gt_mask, gt_class_id, test_set.class_names)

                # Run object detection
                results = model.detect([image], verbose=1)

                # Display results
                r = results[0]
                visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
                                            test_set.class_names, r['scores'],
                                            title="Predictions")