import time
import numpy as np
import torch
import pickle
import warnings
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from torchvision import datasets, transforms
from scipy.ndimage.interpolation import rotate as scipyrotate
from networks import MLP, ConvNet, LeNet, AlexNet, VGG11BN, VGG11, ResNet18, ResNet18BN_AP
from PIL import Image
from torchvision import datasets, transforms
from typing import Any, Callable, List, Optional, Union, Tuple, cast, Dict

import os.path
from urllib.error import URLError
from torchvision.datasets.utils import download_and_extract_archive, extract_archive, verify_str_arg, check_integrity, download_url
from torchvision.datasets.vision import VisionDataset
from torchvision.transforms import functional as FF

import argparse

def get_dataset(dataset, data_path, args):
    if dataset == 'MNIST':
        channel = 1
        im_size = (32, 32)
        num_classes = 10
        mean = [0.1307]
        std = [0.3081]
        transform = transforms.Compose([transforms.Resize([32, 32]), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)])
        dst_train = datasets.MNIST(data_path, train=True, download=True, transform=transform) # no augmentation
        dst_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)
        class_names = [str(c) for c in range(num_classes)]

    elif dataset == 'FashionMNIST':
        channel = 1
        im_size = (32, 32)
        num_classes = 10
        mean = [0.2861]
        std = [0.3530]
        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)])
        dst_train = FashionMNIST_BADNETS(data_path, train=True, download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        dst_test = FashionMNIST_BADNETS(data_path, train=False, download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        # dst_train = datasets.FashionMNIST(data_path, train=True, download=True, transform=transform) # no augmentation
        # dst_test = datasets.FashionMNIST(data_path, train=False, download=True, transform=transform)
        class_names = dst_train.classes

    elif dataset == 'SVHN':
        channel = 3
        im_size = (32, 32)
        num_classes = 10
        mean = [0.4377, 0.4438, 0.4728]
        std = [0.1980, 0.2010, 0.1970]
        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)])
        dst_train = SVHN_BADNETS(data_path, split="train", download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        dst_test = SVHN_BADNETS(data_path, split="test", download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        class_names = [str(c) for c in range(num_classes)]

    elif dataset == 'STL10':
        channel = 3
        im_size = (32, 32)
        num_classes = 10
        mean = [0.4914, 0.4822, 0.4465]
        std = [0.2023, 0.1994, 0.2010]
        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)])
        dst_train = STL10_BADNETS(data_path, split="train", download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        dst_test = STL10_BADNETS(data_path, split="test", download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        class_names = dst_train.classes

    elif dataset == 'CIFAR10':
        channel = 3
        im_size = (32, 32)
        num_classes = 10
        mean = [0.4914, 0.4822, 0.4465]
        std = [0.2023, 0.1994, 0.2010]
        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)])
        dst_train = CIFAR10_BADNETS(data_path, train=True, download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        dst_test = CIFAR10_BADNETS(data_path, train=False, download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        # dst_train = datasets.CIFAR10(data_path, train=True, download=True, transform=transform) # no augmentation
        # dst_test = datasets.CIFAR10(data_path, train=False, download=True, transform=transform)
        class_names = dst_train.classes

    elif dataset == 'CIFAR100':
        channel = 3
        im_size = (32, 32)
        num_classes = 100
        mean = [0.5071, 0.4866, 0.4409]
        std = [0.2673, 0.2564, 0.2762]
        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)])
        dst_train = CIFAR100_BADNETS(data_path, train=True, download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        dst_test = CIFAR100_BADNETS(data_path, train=False, download=True,
                    transform=transform, trigger_label=0, portion=args.portion, backdoor_size=args.backdoor_size, backdoor=args.naive, clean_test=args.clean)
        class_names = dst_train.classes

    else:
        exit('unknown dataset: %s'%dataset)

    testloader = torch.utils.data.DataLoader(dst_test, batch_size=256, shuffle=True, num_workers=3)
    return channel, im_size, num_classes, class_names, mean, std, dst_train, dst_test, testloader

class TensorDataset(Dataset):
    def __init__(self, images, labels): # images: n x c x h x w tensor
        self.images = images.detach().float()
        self.labels = labels.detach()

    def __getitem__(self, index):
        return self.images[index], self.labels[index]

    def __len__(self):
        return self.images.shape[0]

def get_default_convnet_setting():
    net_width, net_depth, net_act, net_norm, net_pooling = 128, 3, 'relu', 'instancenorm', 'avgpooling'
    return net_width, net_depth, net_act, net_norm, net_pooling

def get_network(model, channel, num_classes, im_size=(32, 32)):
    torch.random.manual_seed(int(time.time() * 1000) % 100000)
    net_width, net_depth, net_act, net_norm, net_pooling = get_default_convnet_setting()

    if model == 'MLP':
        net = MLP(channel=channel, num_classes=num_classes)
    elif model == 'ConvNet':
        net = ConvNet(channel=channel, num_classes=num_classes)
    elif model == 'LeNet':
        net = LeNet(channel=channel, num_classes=num_classes)
    elif model == 'AlexNet':
        net = AlexNet(channel=channel, num_classes=num_classes)
    elif model == 'VGG11':
        net = VGG11( channel=channel, num_classes=num_classes)
    elif model == 'VGG11BN':
        net = VGG11BN(channel=channel, num_classes=num_classes)
    elif model == 'ResNet18':
        net = ResNet18(channel=channel, num_classes=num_classes)
    elif model == 'ResNet18BN_AP':
        net = ResNet18BN_AP(channel=channel, num_classes=num_classes)

    elif model == 'ConvNetD1':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=1, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetD2':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=2, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetD3':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=3, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetD4':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=4, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)

    elif model == 'ConvNetW32':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=32, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetW64':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=64, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetW128':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=128, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetW256':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=256, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)

    elif model == 'ConvNetAS':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act='sigmoid', net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetAR':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act='relu', net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetAL':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act='leakyrelu', net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)

    elif model == 'ConvNetNN':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='none', net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetBN':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='batchnorm', net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetLN':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='layernorm', net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetIN':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='instancenorm', net_pooling=net_pooling, im_size=im_size)
    elif model == 'ConvNetGN':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='groupnorm', net_pooling=net_pooling, im_size=im_size)

    elif model == 'ConvNetNP':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling='none', im_size=im_size)
    elif model == 'ConvNetMP':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling='maxpooling', im_size=im_size)
    elif model == 'ConvNetAP':
        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling='avgpooling', im_size=im_size)

    else:
        net = None
        exit('unknown model: %s'%model)

    gpu_num = torch.cuda.device_count()
    if gpu_num>0:
        device = 'cuda'
        if gpu_num>1:
            net = nn.DataParallel(net)
    else:
        device = 'cpu'
    net = net.to(device)

    return net

def get_time():
    return str(time.strftime("[%Y-%m-%d %H:%M:%S]", time.localtime()))

def distance_wb(gwr, gws):
    shape = gwr.shape
    if len(shape) == 4: # conv, out*in*h*w
        gwr = gwr.reshape(shape[0], shape[1] * shape[2] * shape[3])
        gws = gws.reshape(shape[0], shape[1] * shape[2] * shape[3])
    elif len(shape) == 3:  # layernorm, C*h*w
        gwr = gwr.reshape(shape[0], shape[1] * shape[2])
        gws = gws.reshape(shape[0], shape[1] * shape[2])
    elif len(shape) == 2: # linear, out*in
        tmp = 'do nothing'
    elif len(shape) == 1: # batchnorm/instancenorm, C; groupnorm x, bias
        gwr = gwr.reshape(1, shape[0])
        gws = gws.reshape(1, shape[0])
        return torch.tensor(0, dtype=torch.float, device=gwr.device)

    dis_weight = torch.sum(1 - torch.sum(gwr * gws, dim=-1) / (torch.norm(gwr, dim=-1) * torch.norm(gws, dim=-1) + 0.000001))
    dis = dis_weight
    return dis

def match_loss(gw_syn, gw_real, args):
    dis = torch.tensor(0.0).to(args.device)

    if args.dis_metric == 'ours':
        for ig in range(len(gw_real)):
            gwr = gw_real[ig]
            gws = gw_syn[ig]
            dis += distance_wb(gwr, gws)

    elif args.dis_metric == 'mse':
        gw_real_vec = []
        gw_syn_vec = []
        for ig in range(len(gw_real)):
            gw_real_vec.append(gw_real[ig].reshape((-1)))
            gw_syn_vec.append(gw_syn[ig].reshape((-1)))
        gw_real_vec = torch.cat(gw_real_vec, dim=0)
        gw_syn_vec = torch.cat(gw_syn_vec, dim=0)
        dis = torch.sum((gw_syn_vec - gw_real_vec)**2)

    elif args.dis_metric == 'cos':
        gw_real_vec = []
        gw_syn_vec = []
        for ig in range(len(gw_real)):
            gw_real_vec.append(gw_real[ig].reshape((-1)))
            gw_syn_vec.append(gw_syn[ig].reshape((-1)))
        gw_real_vec = torch.cat(gw_real_vec, dim=0)
        gw_syn_vec = torch.cat(gw_syn_vec, dim=0)
        dis = 1 - torch.sum(gw_real_vec * gw_syn_vec, dim=-1) / (torch.norm(gw_real_vec, dim=-1) * torch.norm(gw_syn_vec, dim=-1) + 0.000001)

    else:
        exit('unknown distance function: %s'%args.dis_metric)

    return dis

def get_loops(ipc):
    # Get the two hyper-parameters of outer-loop and inner-loop.
    # The following values are empirically good.
    if ipc == 1:
        outer_loop, inner_loop = 1, 1
    elif ipc == 10:
        outer_loop, inner_loop = 10, 50
    elif ipc == 20:
        outer_loop, inner_loop = 20, 25
    elif ipc == 30:
        outer_loop, inner_loop = 30, 20
    elif ipc == 40:
        outer_loop, inner_loop = 40, 15
    elif ipc == 50:
        outer_loop, inner_loop = 50, 10
    else:
        outer_loop, inner_loop = 0, 0
        exit('loop hyper-parameters are not defined for %d ipc'%ipc)
    return outer_loop, inner_loop

def epoch(mode, dataloader, net, optimizer, criterion, args, aug):
    loss_avg, acc_avg, num_exp = 0, 0, 0
    net = net.to(args.device)
    criterion = criterion.to(args.device)

    if mode == 'train':
        net.train()
    else:
        net.eval()

    for i_batch, datum in enumerate(dataloader):
        img = datum[0].float().to(args.device)
        if aug:
            if args.dsa:
                img = DiffAugment(img, args.dsa_strategy, param=args.dsa_param)
            else:
                img = augment(img, args.dc_aug_param, device=args.device)
        lab = datum[1].long().to(args.device)

        if args.doorping_trigger:
            img[:] = img[:] * (1 - args.mask) + args.init_trigger[0] * args.mask
            lab[:] = args.trigger_label

        if args.invisible_trigger:
            img[:] = img[:] + args.init_trigger[0]
            lab[:] = args.trigger_label

        n_b = lab.shape[0]

        output = net(img)
        loss = criterion(output, lab)
        acc = np.sum(np.equal(np.argmax(output.cpu().data.numpy(), axis=-1), lab.cpu().data.numpy()))

        loss_avg += loss.item()*n_b
        acc_avg += acc
        num_exp += n_b

        if mode == 'train':
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    loss_avg /= num_exp
    acc_avg /= num_exp

    return loss_avg, acc_avg

def evaluate_synset(it_eval, net, images_train, labels_train, testloader, testloader_trigger, args):
    net = net.to(args.device)
    images_train = images_train.to(args.device)
    labels_train = labels_train.to(args.device)
    lr = float(args.lr_net)
    Epoch = int(args.epoch_eval_train)
    lr_schedule = [Epoch//2+1]
    optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005)
    criterion = nn.CrossEntropyLoss().to(args.device)

    dst_train = TensorDataset(images_train, labels_train)
    trainloader = torch.utils.data.DataLoader(dst_train, batch_size=args.batch_train, shuffle=True, num_workers=0)

    start = time.time()
    for ep in range(Epoch+1):
        loss_train, acc_train = epoch('train', trainloader, net, optimizer, criterion, args, aug = True)
        if ep in lr_schedule:
            lr *= 0.1
            optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005)

    time_train = time.time() - start
    args.doorping_trigger = False
    loss_test, acc_test = epoch('test', testloader, net, optimizer, criterion, args, aug = False)

    if args.doorping:
        args.doorping_trigger = True
        loss_test_trigger, acc_test_trigger = epoch('test', testloader, net, optimizer, criterion, args, aug = False)
        args.doorping_trigger = False

    elif args.invisible:
        args.invisible_trigger = True
        loss_test_trigger, acc_test_trigger = epoch('test', testloader, net, optimizer, criterion, args, aug = False)
        args.invisible_trigger = False

    else:
        loss_test_trigger, acc_test_trigger = epoch('test', testloader_trigger, net, optimizer, criterion, args, aug = False)

    print('%s Evaluate_%02d: epoch = %04d train time = %d s train loss = %.6f train acc = %.4f, clean test acc = %.4f, trigger test acc = %.4f' % (get_time(), it_eval, Epoch, int(time_train), loss_train, acc_train, acc_test, acc_test_trigger))
    return net, acc_train, acc_test, acc_test_trigger

def augment(images, dc_aug_param, device):
    # This can be sped up in the future.

    if dc_aug_param != None and dc_aug_param['strategy'] != 'none':
        scale = dc_aug_param['scale']
        crop = dc_aug_param['crop']
        rotate = dc_aug_param['rotate']
        noise = dc_aug_param['noise']
        strategy = dc_aug_param['strategy']

        shape = images.shape
        mean = []
        for c in range(shape[1]):
            mean.append(float(torch.mean(images[:,c])))

        def cropfun(i):
            im_ = torch.zeros(shape[1],shape[2]+crop*2,shape[3]+crop*2, dtype=torch.float, device=device)
            for c in range(shape[1]):
                im_[c] = mean[c]
            im_[:, crop:crop+shape[2], crop:crop+shape[3]] = images[i]
            r, c = np.random.permutation(crop*2)[0], np.random.permutation(crop*2)[0]
            images[i] = im_[:, r:r+shape[2], c:c+shape[3]]

        def scalefun(i):
            h = int((np.random.uniform(1 - scale, 1 + scale)) * shape[2])
            w = int((np.random.uniform(1 - scale, 1 + scale)) * shape[2])
            tmp = F.interpolate(images[i:i + 1], [h, w], )[0]
            mhw = max(h, w, shape[2], shape[3])
            im_ = torch.zeros(shape[1], mhw, mhw, dtype=torch.float, device=device)
            r = int((mhw - h) / 2)
            c = int((mhw - w) / 2)
            im_[:, r:r + h, c:c + w] = tmp
            r = int((mhw - shape[2]) / 2)
            c = int((mhw - shape[3]) / 2)
            images[i] = im_[:, r:r + shape[2], c:c + shape[3]]

        def rotatefun(i):
            im_ = scipyrotate(images[i].cpu().data.numpy(), angle=np.random.randint(-rotate, rotate), axes=(-2, -1), cval=np.mean(mean))
            r = int((im_.shape[-2] - shape[-2]) / 2)
            c = int((im_.shape[-1] - shape[-1]) / 2)
            images[i] = torch.tensor(im_[:, r:r + shape[-2], c:c + shape[-1]], dtype=torch.float, device=device)

        def noisefun(i):
            images[i] = images[i] + noise * torch.randn(shape[1:], dtype=torch.float, device=device)

        augs = strategy.split('_')

        for i in range(shape[0]):
            choice = np.random.permutation(augs)[0] # randomly implement one augmentation
            if choice == 'crop':
                cropfun(i)
            elif choice == 'scale':
                scalefun(i)
            elif choice == 'rotate':
                rotatefun(i)
            elif choice == 'noise':
                noisefun(i)

    return images

def get_daparam(dataset, model, model_eval, ipc):
    # We find that augmentation doesn't always benefit the performance.
    # So we do augmentation for some of the settings.

    dc_aug_param = dict()
    dc_aug_param['crop'] = 4
    dc_aug_param['scale'] = 0.2
    dc_aug_param['rotate'] = 45
    dc_aug_param['noise'] = 0.001
    dc_aug_param['strategy'] = 'none'

    if dataset == 'MNIST':
        dc_aug_param['strategy'] = 'none'

    if model_eval in ['ConvNetBN']: # Data augmentation makes model training with Batch Norm layer easier.
        dc_aug_param['strategy'] = 'crop_noise'

    return dc_aug_param

def get_eval_pool(eval_mode, model, model_eval):
    if eval_mode == 'M': # multiple architectures
        # model_eval_pool = ['MLP', 'ConvNet', 'LeNet', 'AlexNet', 'VGG11', 'ResNet18']
        model_eval_pool = ['VGG11', 'ResNet18']
    elif eval_mode == 'W': # ablation study on network width
        model_eval_pool = ['ConvNetW32', 'ConvNetW64', 'ConvNetW128', 'ConvNetW256']
    elif eval_mode == 'D': # ablation study on network depth
        model_eval_pool = ['ConvNetD1', 'ConvNetD2', 'ConvNetD3', 'ConvNetD4']
    elif eval_mode == 'A': # ablation study on network activation function
        model_eval_pool = ['ConvNetAS', 'ConvNetAR', 'ConvNetAL']
    elif eval_mode == 'P': # ablation study on network pooling layer
        model_eval_pool = ['ConvNetNP', 'ConvNetMP', 'ConvNetAP']
    elif eval_mode == 'N': # ablation study on network normalization layer
        model_eval_pool = ['ConvNetNN', 'ConvNetBN', 'ConvNetLN', 'ConvNetIN', 'ConvNetGN']
    elif eval_mode == 'S': # itself
        if 'BN' in model:
            print('Attention: Here I will replace BN with IN in evaluation, as the synthetic set is too small to measure BN hyper-parameters.')
        model_eval_pool = [model[:model.index('BN')]] if 'BN' in model else [model]
    elif eval_mode == 'SS':  # itself
        model_eval_pool = [model]
    else:
        model_eval_pool = [model_eval]
    return model_eval_pool

class ParamDiffAug():
    def __init__(self):
        self.aug_mode = 'S' #'multiple or single'
        self.prob_flip = 0.5
        self.ratio_scale = 1.2
        self.ratio_rotate = 15.0
        self.ratio_crop_pad = 0.125
        self.ratio_cutout = 0.5 # the size would be 0.5x0.5
        self.brightness = 1.0
        self.saturation = 2.0
        self.contrast = 0.5

def set_seed_DiffAug(param):
    if param.latestseed == -1:
        return
    else:
        torch.random.manual_seed(param.latestseed)
        param.latestseed += 1

def DiffAugment(x, strategy='', seed = -1, param = None):
    if strategy == 'None' or strategy == 'none' or strategy == '':
        return x

    if seed == -1:
        param.Siamese = False
    else:
        param.Siamese = True

    param.latestseed = seed

    if strategy:
        if param.aug_mode == 'M': # original
            for p in strategy.split('_'):
                for f in AUGMENT_FNS[p]:
                    x = f(x, param)
        elif param.aug_mode == 'S':
            pbties = strategy.split('_')
            set_seed_DiffAug(param)
            p = pbties[torch.randint(0, len(pbties), size=(1,)).item()]
            for f in AUGMENT_FNS[p]:
                x = f(x, param)
        else:
            exit('unknown augmentation mode: %s'%param.aug_mode)
        x = x.contiguous()
    return x

# We implement the following differentiable augmentation strategies based on the code provided in https://github.com/mit-han-lab/data-efficient-gans.
def rand_scale(x, param):
    # x>1, max scale
    # sx, sy: (0, +oo), 1: orignial size, 0.5: enlarge 2 times
    ratio = param.ratio_scale
    set_seed_DiffAug(param)
    sx = torch.rand(x.shape[0]) * (ratio - 1.0/ratio) + 1.0/ratio
    set_seed_DiffAug(param)
    sy = torch.rand(x.shape[0]) * (ratio - 1.0/ratio) + 1.0/ratio
    theta = [[[sx[i], 0,  0],
            [0,  sy[i], 0],] for i in range(x.shape[0])]
    theta = torch.tensor(theta, dtype=torch.float)
    if param.Siamese: # Siamese augmentation:
        theta[:] = theta[0]
    grid = F.affine_grid(theta, x.shape).to(x.device)
    x = F.grid_sample(x, grid)
    return x

def rand_rotate(x, param): # [-180, 180], 90: anticlockwise 90 degree
    ratio = param.ratio_rotate
    set_seed_DiffAug(param)
    theta = (torch.rand(x.shape[0]) - 0.5) * 2 * ratio / 180 * float(np.pi)
    theta = [[[torch.cos(theta[i]), torch.sin(-theta[i]), 0],
        [torch.sin(theta[i]), torch.cos(theta[i]),  0],]  for i in range(x.shape[0])]
    theta = torch.tensor(theta, dtype=torch.float)
    if param.Siamese: # Siamese augmentation:
        theta[:] = theta[0]
    grid = F.affine_grid(theta, x.shape).to(x.device)
    x = F.grid_sample(x, grid)
    return x

def rand_flip(x, param):
    prob = param.prob_flip
    set_seed_DiffAug(param)
    randf = torch.rand(x.size(0), 1, 1, 1, device=x.device)
    if param.Siamese: # Siamese augmentation:
        randf[:] = randf[0]
    return torch.where(randf < prob, x.flip(3), x)

def rand_brightness(x, param):
    ratio = param.brightness
    set_seed_DiffAug(param)
    randb = torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device)
    if param.Siamese:  # Siamese augmentation:
        randb[:] = randb[0]
    x = x + (randb - 0.5)*ratio
    return x

def rand_saturation(x, param):
    ratio = param.saturation
    x_mean = x.mean(dim=1, keepdim=True)
    set_seed_DiffAug(param)
    rands = torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device)
    if param.Siamese:  # Siamese augmentation:
        rands[:] = rands[0]
    x = (x - x_mean) * (rands * ratio) + x_mean
    return x

def rand_contrast(x, param):
    ratio = param.contrast
    x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
    set_seed_DiffAug(param)
    randc = torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device)
    if param.Siamese:  # Siamese augmentation:
        randc[:] = randc[0]
    x = (x - x_mean) * (randc + ratio) + x_mean
    return x

def rand_crop(x, param):
    # The image is padded on its surrounding and then cropped.
    ratio = param.ratio_crop_pad
    shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
    set_seed_DiffAug(param)
    translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)
    set_seed_DiffAug(param)
    translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device)
    if param.Siamese:  # Siamese augmentation:
        translation_x[:] = translation_x[0]
        translation_y[:] = translation_y[0]
    grid_batch, grid_x, grid_y = torch.meshgrid(
        torch.arange(x.size(0), dtype=torch.long, device=x.device),
        torch.arange(x.size(2), dtype=torch.long, device=x.device),
        torch.arange(x.size(3), dtype=torch.long, device=x.device),
    )
    grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1)
    grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1)
    x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])
    x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2)
    return x

def rand_cutout(x, param):
    ratio = param.ratio_cutout
    cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
    set_seed_DiffAug(param)
    offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device)
    set_seed_DiffAug(param)
    offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device)
    if param.Siamese:  # Siamese augmentation:
        offset_x[:] = offset_x[0]
        offset_y[:] = offset_y[0]
    grid_batch, grid_x, grid_y = torch.meshgrid(
        torch.arange(x.size(0), dtype=torch.long, device=x.device),
        torch.arange(cutout_size[0], dtype=torch.long, device=x.device),
        torch.arange(cutout_size[1], dtype=torch.long, device=x.device),
    )
    grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1)
    grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1)
    mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)
    mask[grid_batch, grid_x, grid_y] = 0
    x = x * mask.unsqueeze(1)
    return x

AUGMENT_FNS = {
    'color': [rand_brightness, rand_saturation, rand_contrast],
    'crop': [rand_crop],
    'cutout': [rand_cutout],
    'flip': [rand_flip],
    'scale': [rand_scale],
    'rotate': [rand_rotate],
}

class CIFAR10_BADNETS(VisionDataset):
    base_folder = 'cifar-10-batches-py'
    url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
    filename = "cifar-10-python.tar.gz"
    tgz_md5 = 'c58f30108f718f92721af3b95e74349a'
    train_list = [
        ['data_batch_1', 'c99cafc152244af753f735de768cd75f'],
        ['data_batch_2', 'd4bba439e000b95fd0a9bffe97cbabec'],
        ['data_batch_3', '54ebc095f3ab1f0389bbae665268c751'],
        ['data_batch_4', '634d18415352ddfa80567beed471001a'],
        ['data_batch_5', '482c414d41f54cd18b22e5b47cb7c3cb'],
    ]

    test_list = [
        ['test_batch', '40351d587109b95175f43aff81a1287e'],
    ]
    meta = {
        'filename': 'batches.meta',
        'key': 'label_names',
        'md5': '5ff9c542aee3614f3951f8cda6e48888',
    }

    def __init__(
            self,
            root: str,
            train: bool = True,
            transform: Optional[Callable] = None,
            target_transform: Optional[Callable] = None,
            download: bool = False,
            trigger_label: int = 0,
            portion: float =0.1,
            backdoor_size: int = 2,
            backdoor: bool = True,
            clean_test: bool = True,
    ) -> None:

        super(CIFAR10_BADNETS, self).__init__(root, transform=transform,
                                      target_transform=target_transform)

        self.train = train  # training set or test set

        self.trigger_label = trigger_label
        self.portion = portion
        self.backdoor = backdoor
        self.clean_test = clean_test
        self.backdoor_size = backdoor_size

        if download:
            self.download()

        if not self._check_integrity():
            raise RuntimeError('Dataset not found or corrupted.' +
                               ' You can use download=True to download it')

        if self.train:
            downloaded_list = self.train_list
        else:
            downloaded_list = self.test_list

        self.data: Any = []
        self.targets = []

        # now load the picked numpy arrays
        for file_name, checksum in downloaded_list:
            file_path = os.path.join(self.root, self.base_folder, file_name)
            with open(file_path, 'rb') as f:
                entry = pickle.load(f, encoding='latin1')
                self.data.append(entry['data'])
                if 'labels' in entry:
                    self.targets.extend(entry['labels'])
                else:
                    self.targets.extend(entry['fine_labels'])

        self.data = np.vstack(self.data).reshape(-1, 3, 32, 32)
        self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC

        self._load_meta()

        self.targets = np.array(self.targets)

        if self.backdoor:
            if not self.train:
                if self.clean_test:
                    self.portion = 0
                else:
                    self.portion = 1

            self._add_trigger()

        ''''
        self.bad_data, self.bad_targets = self._add_trigger()

        self.total_data = np.concatenate((self.data,self.bad_data),0)
        self.total_targets = np.concatenate((self.targets,self.bad_targets),0)
        '''

    def _add_trigger(self):
        '''
        Based on Vera Xinyue Shen Badnets https://github.com/verazuo/badnets-pytorch
        '''
        perm = np.random.permutation(len(self.data))[0: int(len(self.data) * self.portion)]
        width, height, _ = self.data.shape[1:]
        # self.data[perm, width-3, height-3, :] = 255
        # self.data[perm, width-3, height-2, :] = 255
        # self.data[perm, width-2, height-3, :] = 255
        # self.data[perm, width-2, height-2, :] = 255

        # assert self.backdoor_size == 4

        self.data[perm, width-self.backdoor_size-1:width-1, height-self.backdoor_size-1:height-1, :] = 255
        self.targets[perm] = self.trigger_label

        '''

        new_data = self.data[perm]
        new_targets = self.targets[perm]

        new_data[:, width-3, height-3, :] = 255
        new_data[:, width-3, height-2, :] = 255
        new_data[:, width-2, height-3, :] = 255
        new_data[:, width-2, height-2, :] = 255

        new_targets[:] = self.trigger_label

        '''

        # logging.info("Injecting Over: %d Bad Imgs" % len(perm))
        # return new_data, new_targets

    def _load_meta(self) -> None:
        path = os.path.join(self.root, self.base_folder, self.meta['filename'])
        if not check_integrity(path, self.meta['md5']):
            raise RuntimeError('Dataset metadata file not found or corrupted.' +
                               ' You can use download=True to download it')
        with open(path, 'rb') as infile:
            data = pickle.load(infile, encoding='latin1')
            self.classes = data[self.meta['key']]
        self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)}

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], self.targets[index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self) -> int:
        return len(self.data)

    def _check_integrity(self) -> bool:
        root = self.root
        for fentry in (self.train_list + self.test_list):
            filename, md5 = fentry[0], fentry[1]
            fpath = os.path.join(root, self.base_folder, filename)
            if not check_integrity(fpath, md5):
                return False
        return True

    def download(self) -> None:
        if self._check_integrity():
            print('Files already downloaded and verified')
            return
        download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5)

    def extra_repr(self) -> str:
        return "Split: {}".format("Train" if self.train is True else "Test")

class SVHN_BADNETS(VisionDataset):
    split_list = {
        "train": [
            "http://ufldl.stanford.edu/housenumbers/train_32x32.mat",
            "train_32x32.mat",
            "e26dedcc434d2e4c54c9b2d4a06d8373",
        ],
        "test": [
            "http://ufldl.stanford.edu/housenumbers/test_32x32.mat",
            "test_32x32.mat",
            "eb5a983be6a315427106f1b164d9cef3",
        ],
        "extra": [
            "http://ufldl.stanford.edu/housenumbers/extra_32x32.mat",
            "extra_32x32.mat",
            "a93ce644f1a588dc4d68dda5feec44a7",
        ],
    }

    def __init__(
        self,
        root: str,
        split: str = "train",
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
        trigger_label: int = 0,
        portion: float =0.1,
        backdoor_size: int = 2,
        backdoor: bool = True,
        clean_test: bool = True,
    ) -> None:
        super().__init__(root, transform=transform, target_transform=target_transform)
        self.split = verify_str_arg(split, "split", tuple(self.split_list.keys()))

        self.url = self.split_list[split][0]
        self.filename = self.split_list[split][1]
        self.file_md5 = self.split_list[split][2]
        self.trigger_label = trigger_label
        self.portion = portion
        self.backdoor_size = backdoor_size
        self.backdoor = backdoor
        self.clean_test = clean_test

        if download:
            self.download()

        if not self._check_integrity():
            raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it")

        # import here rather than at top of file because this is
        # an optional dependency for torchvision
        import scipy.io as sio

        # reading(loading) mat file as array
        loaded_mat = sio.loadmat(os.path.join(self.root, self.filename))

        self.data = loaded_mat["X"]
        # loading from the .mat file gives an np array of type np.uint8
        # converting to np.int64, so that we have a LongTensor after
        # the conversion from the numpy array
        # the squeeze is needed to obtain a 1D tensor
        self.labels = loaded_mat["y"].astype(np.int64).squeeze()

        # the svhn dataset assigns the class label "10" to the digit 0
        # this makes it inconsistent with several loss functions
        # which expect the class labels to be in the range [0, C-1]
        np.place(self.labels, self.labels == 10, 0)
        self.data = np.transpose(self.data, (3, 2, 0, 1))

        if self.backdoor:
            if self.split != "train":
                if self.clean_test:
                    self.portion = 0
                else:
                    self.portion = 1

            self._add_trigger()

    def _add_trigger(self):
        '''
        Based on Vera Xinyue Shen Badnets https://github.com/verazuo/badnets-pytorch
        '''
        perm = np.random.permutation(len(self.data))[0: int(len(self.data) * self.portion)]
        width, height, _ = self.data.shape[1:]
        # self.data[perm, width-3, height-3, :] = 255
        # self.data[perm, width-3, height-2, :] = 255
        # self.data[perm, width-2, height-3, :] = 255
        # self.data[perm, width-2, height-2, :] = 255

        # assert self.backdoor_size == 4

        self.data[perm, width-self.backdoor_size-1:width-1, height-self.backdoor_size-1:height-1, :] = 255
        self.labels[perm] = self.trigger_label

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Indexlogging.info

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], int(self.labels[index])

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(np.transpose(img, (1, 2, 0)))

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self) -> int:
        return len(self.data)

    def _check_integrity(self) -> bool:
        root = self.root
        md5 = self.split_list[self.split][2]
        fpath = os.path.join(root, self.filename)
        return check_integrity(fpath, md5)

    def download(self) -> None:
        md5 = self.split_list[self.split][2]
        download_url(self.url, self.root, self.filename, md5)

    def extra_repr(self) -> str:
        return "Split: {split}".format(**self.__dict__)

def get_activation(name, activation):
    def hook(model, input, output):
        activation[name] = output
    return hook

def get_middle_output(net, x, layer):
    temp = []
    for name, _ in net.named_parameters():
        if "weight" in name:
            temp.append(name)

    if -layer > len(temp):
        raise IndexError('layer is out of range')

    name = temp[layer].split('.')
    var = eval('net.' + name[0])
    out = {}
    var[int(name[1])].register_forward_hook(get_activation(str(layer), out))

    _ = net(x)

    return out[str(layer)]

def select_neuron(net, layer, topk):
    logits_weights = None

    count = -1
    for name, param in reversed(list(net.named_parameters())):
        if 'bias' in name:
            continue

        logits_weights = param
        if count == layer:
            break

        count -= 1

    weights = torch.abs(logits_weights.cpu().detach())
    sum_weights = torch.sum(weights, axis=1)
    # max_connection_position = torch.argmax(sum_weights)
    _, max_connection_position = torch.topk(sum_weights, topk)

    return max_connection_position

def update_trigger(net, init_trigger, layer, device, mask, topk, alpha):
    net.eval()
    key_to_maximize = select_neuron(net, layer, topk)
    optimizer = torch.optim.Adam([init_trigger], lr=0.08, betas=[0.9, 0.99])
    criterion = torch.nn.MSELoss().to(device)

    init_output = 0
    cost_threshold = 0.5

    for i in range(1000):
        optimizer.zero_grad()
        # output = model.forward_with_param(state.init_trigger, weights)
        output = get_middle_output(net, init_trigger, layer)

        output = output[:, key_to_maximize]
        if i == 0:
            init_output = output.detach()
        loss = criterion(output, alpha*init_output)
        if loss.item() < cost_threshold:
            break

        loss.backward()
        init_trigger.grad.data.mul_(mask)

        optimizer.step()

    return init_trigger

def update_inv_trigger(net, init_trigger, layer, device, std, black):
    net.eval()
    key_to_maximize = select_neuron(net, layer)
    optimizer = torch.optim.Adam([init_trigger], lr=0.08, betas=[0.9, 0.99])
    criterion = torch.nn.MSELoss().to(device)

    c = 1
    tao = 1/ min(std)
    tao_best = float('inf')

    EARLY_STOP_THRESHOLD = 1.0
    EARLY_STOP_PATIENCE = 200
    early_stop_counter = 0
    early_stop_reg_best = tao_best

    for i in range(100):
        optimizer.zero_grad()
        output = get_middle_output(net, init_trigger, layer)

        output = output[:, key_to_maximize]
        l_inf = torch.abs(init_trigger - black)
        loss = c * criterion(output, 10.*output) + l_inf[l_inf > tao].sum()
        # logging.info(loss.item())

        if i % 500 == 0:
            c *= 0.8

        if l_inf.max().item() < tao:
            tao *= 0.9

        tao_best = min(tao_best, tao)

        if tao_best < float('inf'):
            if tao_best >= EARLY_STOP_THRESHOLD * early_stop_reg_best:
                early_stop_counter += 1
            else:
                early_stop_counter = 0
        if early_stop_counter >= EARLY_STOP_PATIENCE:
            break
        early_stop_reg_best = min(tao_best, early_stop_reg_best)

        loss.backward()
        optimizer.step()

    return init_trigger

def save_trigger_img(args, expr_dir):
    from torchvision.utils import save_image
    img = args.init_trigger[0]
    save_data_path = os.path.join(expr_dir, 'trigger.png')
    save_image(img, save_data_path)

class STL10_BADNETS(VisionDataset):

    base_folder = "stl10_binary"
    url = "http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz"
    filename = "stl10_binary.tar.gz"
    tgz_md5 = "91f7769df0f17e558f3565bffb0c7dfb"
    class_names_file = "class_names.txt"
    folds_list_file = "fold_indices.txt"
    train_list = [
        ["train_X.bin", "918c2871b30a85fa023e0c44e0bee87f"],
        ["train_y.bin", "5a34089d4802c674881badbb80307741"],
        ["unlabeled_X.bin", "5242ba1fed5e4be9e1e742405eb56ca4"],
    ]

    test_list = [["test_X.bin", "7f263ba9f9e0b06b93213547f721ac82"], ["test_y.bin", "36f9794fa4beb8a2c72628de14fa638e"]]
    splits = ("train", "train+unlabeled", "unlabeled", "test")

    def __init__(
        self,
        root: str,
        split: str = "train",
        folds: Optional[int] = None,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
        trigger_label: int = 0,
        portion: float =0.1,
        backdoor_size: int = 2,
        backdoor: bool = True,
        clean_test: bool = True,
    ) -> None:
        super().__init__(root, transform=transform, target_transform=target_transform)
        self.split = verify_str_arg(split, "split", self.splits)
        self.folds = self._verify_folds(folds)

        self.trigger_label = trigger_label
        self.portion = portion
        self.backdoor = backdoor
        self.clean_test = clean_test
        self.backdoor_size = backdoor_size

        if download:
            self.download()
        elif not self._check_integrity():
            raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it")

        # now load the picked numpy arrays
        self.labels: Optional[np.ndarray]
        if self.split == "train":
            self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0])
            self.labels = cast(np.ndarray, self.labels)
            self.__load_folds(folds)

        elif self.split == "train+unlabeled":
            self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0])
            self.labels = cast(np.ndarray, self.labels)
            self.__load_folds(folds)
            unlabeled_data, _ = self.__loadfile(self.train_list[2][0])
            self.data = np.concatenate((self.data, unlabeled_data))
            self.labels = np.concatenate((self.labels, np.asarray([-1] * unlabeled_data.shape[0])))

        elif self.split == "unlabeled":
            self.data, _ = self.__loadfile(self.train_list[2][0])
            self.labels = np.asarray([-1] * self.data.shape[0])
        else:  # self.split == 'test':
            self.data, self.labels = self.__loadfile(self.test_list[0][0], self.test_list[1][0])

        class_file = os.path.join(self.root, self.base_folder, self.class_names_file)
        if os.path.isfile(class_file):
            with open(class_file) as f:
                self.classes = f.read().splitlines()

        if self.backdoor:
            if self.split != "train":
                if self.clean_test:
                    self.portion = 0
                else:
                    self.portion = 1
            self.perm = np.random.permutation(len(self.data))[0: int(len(self.data) * self.portion)]

    # def _add_trigger(self):
    #     '''
    #     Based on Vera Xinyue Shen Badnets https://github.com/verazuo/badnets-pytorch
    #     '''

    #     _, width, height = self.data.shape[1:]
    #     # self.data[perm, width-3, height-3, :] = 255
    #     # self.data[perm, width-3, height-2, :] = 255
    #     # self.data[perm, width-2, height-3, :] = 255
    #     # self.data[perm, width-2, height-2, :] = 255

    #     # assert self.backdoor_size == 4

    #     self.data[perm, :, width-self.backdoor_size-1:width-1, height-self.backdoor_size-1:height-1] = 255
    #     self.labels[perm] = self.trigger_label

    def _verify_folds(self, folds: Optional[int]) -> Optional[int]:
        if folds is None:
            return folds
        elif isinstance(folds, int):
            if folds in range(10):
                return folds
            msg = "Value for argument folds should be in the range [0, 10), but got {}."
            raise ValueError(msg.format(folds))
        else:
            msg = "Expected type None or int for argument folds, but got type {}."
            raise ValueError(msg.format(type(folds)))

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        target: Optional[int]
        if self.labels is not None:
            img, target = self.data[index], int(self.labels[index])
        else:
            img, target = self.data[index], None

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(np.transpose(img, (1, 2, 0)))
        img = FF.resize(img, (32,32))
        if self.backdoor:
            if index in self.perm:
                img = np.asarray(img)
                width, height, _ = img.shape
                img[width-self.backdoor_size-1:width-1, height-self.backdoor_size-1:height-1, :] = 255
                img = Image.fromarray(img)
                target = self.trigger_label

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self) -> int:
        return self.data.shape[0]

    def __loadfile(self, data_file: str, labels_file: Optional[str] = None) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        labels = None
        if labels_file:
            path_to_labels = os.path.join(self.root, self.base_folder, labels_file)
            with open(path_to_labels, "rb") as f:
                labels = np.fromfile(f, dtype=np.uint8) - 1  # 0-based

        path_to_data = os.path.join(self.root, self.base_folder, data_file)
        with open(path_to_data, "rb") as f:
            # read whole file in uint8 chunks
            everything = np.fromfile(f, dtype=np.uint8)
            images = np.reshape(everything, (-1, 3, 96, 96))
            images = np.transpose(images, (0, 1, 3, 2))

        return images, labels

    def _check_integrity(self) -> bool:
        root = self.root
        for fentry in self.train_list + self.test_list:
            filename, md5 = fentry[0], fentry[1]
            fpath = os.path.join(root, self.base_folder, filename)
            if not check_integrity(fpath, md5):
                return False
        return True

    def download(self) -> None:
        if self._check_integrity():
            print("Files already downloaded and verified")
            return
        download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5)
        self._check_integrity()

    def extra_repr(self) -> str:
        return "Split: {split}".format(**self.__dict__)

    def __load_folds(self, folds: Optional[int]) -> None:
        # loads one of the folds if specified
        if folds is None:
            return
        path_to_folds = os.path.join(self.root, self.base_folder, self.folds_list_file)
        with open(path_to_folds) as f:
            str_idx = f.read().splitlines()[folds]
            list_idx = np.fromstring(str_idx, dtype=np.int64, sep=" ")
            self.data = self.data[list_idx, :, :, :]
            if self.labels is not None:
                self.labels = self.labels[list_idx]

class CIFAR100_BADNETS(CIFAR10_BADNETS):
    """`CIFAR100 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.

    This is a subclass of the `CIFAR10` Dataset.
    """

    base_folder = "cifar-100-python"
    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
    filename = "cifar-100-python.tar.gz"
    tgz_md5 = "eb9058c3a382ffc7106e4002c42a8d85"
    train_list = [
        ["train", "16019d7e3df5f24257cddd939b257f8d"],
    ]

    test_list = [
        ["test", "f0ef6b0ae62326f3e7ffdfab6717acfc"],
    ]
    meta = {
        "filename": "meta",
        "key": "fine_label_names",
        "md5": "7973b15100ade9c7d40fb424638fde48",
    }
class MNIST_BADNETS(VisionDataset):
    mirrors = [
        'http://yann.lecun.com/exdb/mnist/',
        'https://ossci-datasets.s3.amazonaws.com/mnist/',
    ]

    resources = [
        ("train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"),
        ("train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"),
        ("t10k-images-idx3-ubyte.gz", "9fb629c4189551a2d022fa330f9573f3"),
        ("t10k-labels-idx1-ubyte.gz", "ec29112dd5afa0611ce80d1b7f02629c")
    ]

    training_file = 'training.pt'
    test_file = 'test.pt'
    classes = ['0 - zero', '1 - one', '2 - two', '3 - three', '4 - four',
               '5 - five', '6 - six', '7 - seven', '8 - eight', '9 - nine']

    @property
    def train_labels(self):
        warnings.warn("train_labels has been renamed targets")
        return self.targets

    @property
    def test_labels(self):
        warnings.warn("test_labels has been renamed targets")
        return self.targets

    @property
    def train_data(self):
        warnings.warn("train_data has been renamed data")
        return self.data

    @property
    def test_data(self):
        warnings.warn("test_data has been renamed data")
        return self.data

    def __init__(
            self,
            root: str,
            train: bool = True,
            transform: Optional[Callable] = None,
            target_transform: Optional[Callable] = None,
            download: bool = False,
            trigger_label: int = 0,
            portion: float =0.01,
            backdoor_size: int = 2,
            backdoor: bool = True,
            clean_test: bool = True,
    ) -> None:
        super(MNIST_BADNETS, self).__init__(root, transform=transform,
                                    target_transform=target_transform)
        self.train = train  # training set or test set

        self.trigger_label = trigger_label
        self.portion = portion
        self.backdoor_size = backdoor_size
        self.backdoor = backdoor
        self.clean_test = clean_test

        if self._check_legacy_exist():
            self.data, self.targets = self._load_legacy_data()
            return

        if download:
            self.download()

        if not self._check_exists():
            raise RuntimeError('Dataset not found.' +
                               ' You can use download=True to download it')

        self.data, self.targets = self._load_data()
        if self.backdoor:
            if not self.train:
                if self.clean_test:
                    self.portion = 0
                else:
                    self.portion = 1

            self.perm = np.random.permutation(len(self.data))[0: int(len(self.data) * self.portion)]

    # def _add_trigger(self):
    #     perm = np.random.permutation(len(self.data))[0: int(len(self.data) * self.portion)]

    #     width, height = self.data.shape[1:]

    #     self.data[perm, width-self.backdoor_size-1:width-1, height-self.backdoor_size-1:height-1] = 255
    #     self.targets[perm] = self.trigger_label

    def _check_legacy_exist(self):
        processed_folder_exists = os.path.exists(self.processed_folder)
        if not processed_folder_exists:
            return False

        return all(
            check_integrity(os.path.join(self.processed_folder, file)) for file in (self.training_file, self.test_file)
        )

    def _load_legacy_data(self):
        # This is for BC only. We no longer cache the data in a custom binary, but simply read from the raw data
        # directly.
        data_file = self.training_file if self.train else self.test_file
        return torch.load(os.path.join(self.processed_folder, data_file))

    def _load_data(self):
        image_file = f"{'train' if self.train else 't10k'}-images-idx3-ubyte"
        data = datasets.mnist.read_image_file(os.path.join(self.raw_folder, image_file))

        label_file = f"{'train' if self.train else 't10k'}-labels-idx1-ubyte"
        targets = datasets.mnist.read_label_file(os.path.join(self.raw_folder, label_file))

        return data, targets

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """

        img, target = self.data[index], int(self.targets[index])

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img.numpy(), mode='L')
        img = FF.resize(img, (32,32))
        if self.backdoor:
            if index in self.perm:
                img = np.asarray(img)
                width, height = img.shape
                img[width-self.backdoor_size-1:width-1, height-self.backdoor_size-1:height-1] = 255
                img = Image.fromarray(img)
                target = self.trigger_label

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self) -> int:
        return len(self.data)

    @property
    def raw_folder(self) -> str:
        return os.path.join(self.root, self.__class__.__name__, 'raw')

    @property
    def processed_folder(self) -> str:
        return os.path.join(self.root, self.__class__.__name__, 'processed')

    @property
    def class_to_idx(self) -> Dict[str, int]:
        return {_class: i for i, _class in enumerate(self.classes)}

    def _check_exists(self) -> bool:
        return all(
            check_integrity(os.path.join(self.raw_folder, os.path.splitext(os.path.basename(url))[0]))
            for url, _ in self.resources
        )

    def download(self) -> None:
        """Download the MNIST data if it doesn't exist already."""

        if self._check_exists():
            return

        os.makedirs(self.raw_folder, exist_ok=True)

        # download files
        for filename, md5 in self.resources:
            for mirror in self.mirrors:
                url = "{}{}".format(mirror, filename)
                try:
                    print("Downloading {}".format(url))
                    download_and_extract_archive(
                        url, download_root=self.raw_folder,
                        filename=filename,
                        md5=md5
                    )
                except URLError as error:
                    print(
                        "Failed to download (trying next):\n{}".format(error)
                    )
                    continue
                finally:
                    print()
                break
            else:
                raise RuntimeError("Error downloading {}".format(filename))

    def extra_repr(self) -> str:
        return "Split: {}".format("Train" if self.train is True else "Test")

class FashionMNIST_BADNETS(MNIST_BADNETS):
    """`Fashion-MNIST <https://github.com/zalandoresearch/fashion-mnist>`_ Dataset.

    Args:
        root (string): Root directory of dataset where ``FashionMNIST/raw/train-images-idx3-ubyte``
            and  ``FashionMNIST/raw/t10k-images-idx3-ubyte`` exist.
        train (bool, optional): If True, creates dataset from ``train-images-idx3-ubyte``,
            otherwise from ``t10k-images-idx3-ubyte``.
        download (bool, optional): If True, downloads the dataset from the internet and
            puts it in root directory. If dataset is already downloaded, it is not
            downloaded again.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
    """

    mirrors = ["http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/"]

    resources = [
        ("train-images-idx3-ubyte.gz", "8d4fb7e6c68d591d4c3dfef9ec88bf0d"),
        ("train-labels-idx1-ubyte.gz", "25c81989df183df01b3e8a0aad5dffbe"),
        ("t10k-images-idx3-ubyte.gz", "bef4ecab320f06d8554ea6380940ec79"),
        ("t10k-labels-idx1-ubyte.gz", "bb300cfdad3c16e7a12a480ee83cd310"),
    ]
    classes = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]


"""
Менеджер контекста
"""

fh = None
try:
    fh = open('file.txt')
    for line in fh:
        print(line)
except Exception as e:
    print(e)
finally:
    if fh:
        fh.close()

# равнозначная запись
try:
    with open('filename.txt') as fh:
        for line in fh:
            print(line)
except Exception as e:
    print(e)

# Для того чтобы создать свой менеджер контекста
# из класса - определите спецметоды __enter__(self) и __exit__(self, exc_type, exc_value, tb_value)
# из функции - воспользуйтесь декоратором `contextlib.contextmanager`

# быстрое задание - напишите менеджер контекста, позволяющий делать несколько действий со списком атомарно


class Solution:
    def computeArea(self, ax1: int, ay1: int, ax2: int, ay2: int, bx1: int, by1: int, bx2: int, by2: int) -> int:

        R1_area = (ax2-ax1) * (ay2-ay1)
        R2_area = (bx2-bx1) * (by2-by1)
        Common_area = 0
        if min(ax2,bx2) > max (ax1, bx1) and min (ay2,by2) > max(ay1, by1):
            Common_area = (min(ax2,bx2) - max (ax1, bx1)) * (min (ay2,by2) - max(ay1, by1))

        Total_area = R1_area + R2_area - Common_area
        return Total_area

Additional Algorithms/LC 223 Compute Area of  Overlapping Rectangles.py


""" 05_Collisions_v2 by Sun Woo Yi
This version will be carried on from 05_Collisions_v1_testing_2
This version will show a collision being detected between two objects
When the objects collide the game will quit automatically
26/05/2023
"""

import pygame

# Initialize Pygame
pygame.init()

# Set the dimensions of the window
window_width = 261
window_height = 377

# Create the window
screen = pygame.display.set_mode((window_width, window_height))

class GreenBox(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.Surface((50, 50))
        self.image.fill((0, 255, 0))
        self.rect = self.image.get_rect()
        self.rect.x = x
        self.rect.y = y

class BlueBox(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.Surface((50, 50))
        self.image.fill((0, 0, 255))
        self.rect = self.image.get_rect()
        self.rect.x = x
        self.rect.y = y

# Set up the green box
green_box_x = window_width/2 - 50/2
green_box_y = 377 - 50
green_box = GreenBox(green_box_x, green_box_y)
all_sprites = pygame.sprite.Group(green_box)

# Set up the blue box
blue_box_x = window_width/2 - 50/2
blue_box_y = 0
blue_box = BlueBox(blue_box_x, blue_box_y)
all_sprites.add(blue_box)

# Game loop
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Move the boxes
    green_box.rect.y -= 1
    blue_box.rect.y += 1

    # Check for collisions
    if pygame.sprite.collide_rect(green_box, blue_box):
        quit()

    # Draw the boxes
    screen.fill((255, 255, 255))
    all_sprites.draw(screen)
    pygame.display.update()

pygame.quit()


class Solution:
    def longestPalindrome(self, s: str) -> str:
        """
        Time complexity : O(n^3). Assume that n is the length of the input string,
        there are a total of C(n, 2) = n(n-1)/2 substrings
        (excluding the trivial solution where a character itself is a palindrome).
        Since verifying each substring takes O(n) time, the run time complexity is O(n^3).

        Space complexity : O(1).
        """
        longest_palindrome = ""
        for i in range(len(s)):
            for j in range(1, len(s)+1):
                subs = s[i:j]
                if len(subs) > len(longest_palindrome) and subs == subs[::-1]:
                    longest_palindrome = s[i:j]
        return longest_palindrome

    def longestPalindrome2(self, s: str) -> str:
        """
        Time complexity : O(n^2).
        Since expanding a palindrome around its center could take O(n) time, the overall complexity is O(n^2).

        Space complexity : O(1).
        """
        if len(s) <= 1:
            return s

        def expand_from_center(left, right):
            while left >= 0 and right < len(s) and s[left] == s[right]:
                left -= 1
                right += 1
            return s[left + 1:right]

        max_str = s[0]

        for i in range(len(s) - 1):
            odd = expand_from_center(i, i)
            even = expand_from_center(i, i + 1)

            if len(odd) > len(max_str):
                max_str = odd
            if len(even) > len(max_str):
                max_str = even

        return max_str

test1 = Solution().longestPalindrome("a")
print(test1)


from functools import lru_cache
from typing import List

class Solution:
    def minScoreTriangulation(self, values: List[int]) -> int:
        @lru_cache(None)
        def dp(i, j):
            if i + 2 > j:
                return 0
            if i + 2 == j:
                return values[i] * values[i + 1] * values[j]
            return min((values[i] * values[k] * values[j] + dp(i, k) + dp(k, j)) for k in range(i + 1, j))

        return dp(0, len(values) - 1)

if __name__  == '__main__':
    values = [1,2,3]
    values = [3,7,4,5]
    values = [1,3,1,4,1,5]
    rtn = Solution().minScoreTriangulation(values)
    print(rtn)

1039. 多边形三角剖分的最低得分.py


#  -*-coding:utf8-*-

################################################################################
#
#
#
################################################################################
"""
模块用法说明: 应用各UIAWindow的名称

Authors: turinblueice
Date: 2016/7/28
"""

class WindowNames(object):

    LOGIN_MAIN = 'name_login_background.png'   # 登录页的名称
    IN_MAIN = 'name_paizhao'  # in主页名称,in主页底部tab存在拍照按钮固定不变
    IN_CENTER = 'name_我的'  # in主页-中心页

    LOGIN_FRIEND_RECOMMEND = 'name_推荐好友'  # 登录的好友推荐页

    USER_INFO = 'name_编辑资料'  # 编辑资料名称
    CROPPER = 'name_选取'  # 裁剪图片名称
    USER_SETTING = 'name_设置'  # 用户设置页的activity名称

    PASTER_MALL = 'name_INPasterMallResource.bundle/pastermallbg.jpg'  # 我的贴纸
    PASTER_DIY_GUIDE = 'name_INPasterMallResource.bundle/pasterdiy_addguide'  # 创建贴纸的引导
    CUSTOM_PASTER_CROP = 'name_选择区域'   # 自定义贴纸裁剪选择页
    CUSTOM_PASTER_EDITOR = 'name_自定义贴纸'  #自定义贴纸搭配页
    SELECT_PHOTO_GUIDE = 'name_INPasterMallResource.bundle/pasterdiy_selectimagetip'  # 图片选择引导页
    CHOOSE_PART_PASTER = 'name_选择想要的部分'  # 图片编辑筛选页

    CAMERA = 'name_INPublishResource.bundle/camera_opening_above.png'  # 照片编辑/拍照页,滑动切换滤镜说明所在页
    PUBLISH_CORE = 'name_INPublishResource.bundle/tag_local'  # 图片加工发布页面
    PUBLISH = 'name_发布'  # 图片发布页

    PHOTO_STORY_GALLERY = 'name_相机胶卷'  # 故事集/图片选择页,顶部"相机胶卷"文字不变
    PHOTO_TAKING = 'name_PhotoCapture'  # 拍照页
    PHOTO_ALBUM_PICKER = 'name_照片'  # 照片集上传页的标志/照片集选取页

    STORY_EDIT = 'name_编辑故事'  # 故事编辑页
    STORY_DETAIL = 'name_INAlbumDetailWhiteBack'  # 故事预览/详情页
    STORY_SETTING = 'name_故事集设置'  # 故事集设置页
    STORY_SHARE = 'name_发布成功!'  # 故事发布成功分享页

    ADD_FRIEND = '.usercenter.activity.UserCenterFriendsAddActivity'  # 添加好友页
    USER_CENTER_FRIEND = 'name_我的好友'  # 用户中心的我的好友页

    DIARY_INFO = '.diary.other.v260.DiaryOtherActivity'  # 日记主页
    PHOTO_ALBUM_CORE = '.photo.PhotoCoreActivity'  # 图片集详情页

    TOPIC_DETAIL = '.module.tag.activity.TagActivityV253'  # 话题详情页
    FRIEND_PHOTO_DETAIL = '.friend.activity.FriendPhotoDetailActivity'  # 好友图片详情页


from __future__ import print_function
import os, sys, tempfile, shutil, tarfile
import log, argdb
from urllib.request import urlretrieve
from urllib import parse as urlparse_local
import subprocess
import socket
from shutil import which  # just to break compatibility with python2

# Fix parsing for nonstandard schemes
urlparse_local.uses_netloc.extend(['bk', 'ssh', 'svn'])

class Package:

  packages = []    # list of packages processed and added so far

  def __init__(self,argdb,log):
    self.installable     = False  # an already installed package can be picked --with-xxx-dir
    self.downloadable    = False  # package can be downloaded and installed with --download-xxx
    self.downloadpackage = 0
    self.packagetype     = ''     # can be 'gnu', 'cmake', 'source_c', or empty
    self.packagedir      = ''
    self.packagelibs     = ''
    self.packageincludes = ''
    self.packageurl      = ''
    self.buildflags      = ''
    self.log             = log
    self.supportsscalar  = ['real', 'complex']
    self.supportssingle  = False
    self.supports64bint  = False
    self.fortran         = False
    self.hasheaders      = False
    self.requested       = False
    self.havepackage     = False

  def RunCommand(self,instr):
    try:
      self.log.write('- '*35+'\nRunning command:\n'+instr+'\n'+'- '*35)
    except AttributeError: pass
    try:
      output = subprocess.check_output(instr,shell=True,stderr=subprocess.STDOUT)
      result = 0
    except subprocess.CalledProcessError as ex:
      output = ex.output
      result = ex.returncode
    output = output.decode(encoding='UTF-8',errors='replace').rstrip()
    try:
      self.log.write('Output:\n'+output+'\n'+'- '*35)
    except AttributeError: pass
    return (result,output)

  def ProcessArgs(self,argdb,petscpackages=''):
    if hasattr(self,'petscdepend') and self.petscdepend in petscpackages:
      self.requested = True
    if self.installable and not hasattr(self,'petscdepend'):
      string,found = argdb.PopPath('with-'+self.packagename+'-dir',exist=True)
      if found:
        self.requested = True
        self.packagedir = string
      string,found = argdb.PopString('with-'+self.packagename+'-lib')
      if found:
        if self.packagedir:
          self.log.Exit('Specify either "--with-'+self.packagename+'-dir" or "--with-'+self.packagename+'-lib%s", but not both!' % (' --with-'+self.packagename+'-include' if self.hasheaders else ''))
        self.requested = True
        self.packagelibs = string
      if self.hasheaders:
        string,found = argdb.PopString('with-'+self.packagename+'-include')
        if found:
          self.requested = True
          self.packageincludes = string
    if self.installable:
      value,found = argdb.PopBool('with-'+self.packagename)
      if found:
        self.requested = value
    if self.downloadable:
      flagsfound = False
      if self.packagetype == 'gnu':
        string,flagsfound = argdb.PopString('download-'+self.packagename+'-configure-arguments')
      elif self.packagetype == 'cmake':
        string,flagsfound = argdb.PopString('download-'+self.packagename+'-cmake-arguments')
      elif self.packagetype == 'source_c':
        string,flagsfound = argdb.PopString('download-'+self.packagename+'-cflags')
      url,flag,found = argdb.PopUrl('download-'+self.packagename)
      if found:
        if self.requested:
          self.log.Exit('Cannot request both download and install simultaneously')
        self.requested = True
        self.download = True
        self.packageurl = url
        self.downloadpackage = flag
      if flagsfound:
        if not hasattr(self,'download') or not self.download:
          if self.packagetype == 'gnu':
            self.log.Exit('--download-'+self.packagename+'-configure-arguments must be used together with --download-'+self.packagename)
          elif self.packagetype == 'cmake':
            self.log.Exit('--download-'+self.packagename+'-cmake-arguments must be used together with --download-'+self.packagename)
          elif self.packagetype == 'source_c':
            self.log.Exit('--download-'+self.packagename+'-cflags must be used together with --download-'+self.packagename)
        self.buildflags = string

  def Process(self,slepcconf,slepcvars,slepcrules,slepc,petsc,archdir=''):
    self.make = petsc.make
    if petsc.buildsharedlib:
      self.slflag = petsc.cc_linker_slflag
    if self.requested:
      name = self.packagename.upper()
      if self.downloadpackage:
        if hasattr(self,'version') and self.packageurl=='':
          self.log.NewSection('Installing '+name+' version '+self.version+'...')
        else:
          self.log.NewSection('Installing '+name+'...')
        self.Precondition(slepc,petsc)
        self.DownloadAndInstall(slepcconf,slepcvars,slepc,petsc,archdir,slepc.prefixdir)
      elif self.installable:
        self.log.NewSection('Checking '+name+'...')
        self.Precondition(slepc,petsc)
        if petsc.buildsharedlib:
          self.packagelibs = self.DistilLibList(self.packagelibs,petsc)
        self.Check(slepcconf,slepcvars,petsc,archdir)
        if not self.havepackage: self.log.setLastFailed()
      try:
        self.LoadVersion(slepcconf)
        self.log.write('Version number for '+name+' is '+self.iversion)
      except AttributeError:
        pass
      if self.havepackage: Package.packages.append(self)
    else: # not requested
      if hasattr(self,'SkipInstall'):
        self.SkipInstall(slepcrules)

  def Precondition(self,slepc,petsc):
    package = self.packagename.upper()
    if petsc.scalar == 'complex':
      if 'complex' not in self.supportsscalar:
        self.log.Exit(package+' does not support complex scalars')
    elif petsc.scalar == 'real':
      if 'real' not in self.supportsscalar:
        self.log.Exit(package+' is supported only with complex scalars')
    if petsc.precision == 'single':
      if not self.supportssingle:
        self.log.Exit(package+' is supported only in double precision')
    elif petsc.precision != 'double':
      self.log.Exit('Precision '+petsc.precision+' is not supported for external packages')
    if petsc.ind64 and not self.supports64bint:
      self.log.Exit(package+' cannot be used with 64-bit integers')
    if self.downloadpackage and self.fortran and not hasattr(petsc,'fc'):
      self.log.Exit('Option --download-'+self.packagename+' requires a Fortran compiler')

  def Require(self,packagename):
    for p in Package.packages:
      if p.packagename.upper() == packagename.upper():
        return p
    self.log.Exit('The package '+self.packagename.upper()+' requires configuring also with '+packagename.upper()+". Run configure --help for details on how to install it")

  def DistilLibList(self,packagelibs,petsc):
    libs = []
    for l in packagelibs.split():
      if l.endswith(petsc.sl_linker_suffix):
        filename = os.path.basename(l)
        libs.append('-L'+l.rstrip(filename))
        libs.append(self.slflag+l.rstrip(filename))
        libs.append('-l'+filename.lstrip('lib').rstrip('.'+petsc.sl_linker_suffix))
      else:
        libs.append(l)
    newldflags = []
    newlibs = []
    dupflags = ['-L',self.slflag]
    for j in libs:
      # remove duplicate -L, -Wl,-rpath options - and only consecutive -l options
      if j in newldflags and any([j.startswith(flg) for flg in dupflags]): continue
      if newlibs and j == newlibs[-1]: continue
      if list(filter(j.startswith,['-l'])) or list(filter(j.endswith,['.lib','.a','.so','.o'])) or j in ['-Wl,-Bstatic','-Wl,-Bdynamic','-Wl,--start-group','-Wl,--end-group']:
        newlibs.append(j)
      else:
        newldflags.append(j)
    liblist = ' '.join(newldflags + newlibs)
    return liblist

  def GetArchiveName(self):
    '''Return name of archive after downloading'''
    if self.packageurl=='':
      archivename = self.archive
    else:
      parsed = urlparse_local.urlparse(self.packageurl)
      archivename = os.path.basename(parsed[2])
    if archivename[0] == 'v':
      archivename = archivename[1:]
    try:
      if archivename[0].isdigit() or int(archivename.split('.')[0],16):
        archivename = self.packagename+'-'+archivename
    except: pass
    return archivename

  def GetDirectoryName(self):
    '''Return name of the directory after extracting the tarball'''
    dirname = self.GetArchiveName()
    for suffix in ('.tar.gz','.tgz'):
      if dirname.endswith(suffix):
        dirname = dirname[:-len(suffix)]
    return dirname

  def MissingTarball(self,downloaddir):
    '''Check if tarball is missing in downloaddir'''
    if self.downloadable and hasattr(self,'download') and self.download:
      localFile = os.path.join(downloaddir,self.GetArchiveName())
      if not os.path.exists(localFile):
        url = self.packageurl
        if url=='':
          url = self.url
        return self.packagename+': '+url+' --> '+localFile

  def Download(self,externdir,downloaddir):
    # Quick return: check if source is already available
    if os.path.exists(os.path.join(externdir,self.GetDirectoryName())):
      self.log.write('Using '+os.path.join(externdir,self.GetDirectoryName()))
      return os.path.join(externdir,self.GetDirectoryName())

    if downloaddir:
      # Get tarball from download dir
      localFile = os.path.join(downloaddir,self.GetArchiveName())
      if not os.path.exists(localFile):
        self.log.Exit('Could not find file '+self.GetArchiveName()+' under '+downloaddir)
      url = localFile
      filename = os.path.basename(url)
    else:
      # Download tarball
      url = self.packageurl
      if url=='':
        url = self.url
      if os.path.exists(url):
        url = 'file:'+url
      filename = os.path.basename(urlparse_local.urlparse(url)[2])
      localFile = os.path.join(externdir,self.GetArchiveName())
      self.log.write('Downloading '+url+' to '+localFile)

      if os.path.exists(localFile):
        os.remove(localFile)
      try:
        sav_timeout = socket.getdefaulttimeout()
        socket.setdefaulttimeout(30)
        urlretrieve(url, localFile)
        socket.setdefaulttimeout(sav_timeout)
      except Exception as e:
        socket.setdefaulttimeout(sav_timeout)
        failureMessage = '''\
Unable to download package %s from: %s
* If URL specified manually - perhaps there is a typo?
* If your network is disconnected - please reconnect and rerun ./configure
* Or perhaps you have a firewall blocking the download
* You can run with --with-packages-download-dir=/adirectory and ./configure will instruct you what packages to download manually
* or you can download the above URL manually, to /yourselectedlocation/%s
  and use the configure option:
  --download-%s=/yourselectedlocation/%s
''' % (self.packagename.upper(), url, filename, self.packagename, filename)
        self.log.Exit(failureMessage)

    # Uncompress tarball
    extractdir = os.path.join(externdir,self.GetDirectoryName())
    self.log.write('Uncompressing '+localFile+' to directory '+extractdir)
    if os.path.exists(extractdir):
      for root, dirs, files in os.walk(extractdir, topdown=False):
        for name in files:
          os.remove(os.path.join(root,name))
        for name in dirs:
          os.rmdir(os.path.join(root,name))
    failureMessage = '''\
Downloaded package %s from: %s is not a tarball.
[or installed python cannot process compressed files]
* If you are behind a firewall - please fix your proxy and rerun ./configure
  For example at LANL you may need to set the environmental variable http_proxy (or HTTP_PROXY?) to  http://proxyout.lanl.gov
* You can run with --with-packages-download-dir=/adirectory and ./configure will instruct you what packages to download manually
* or you can download the above URL manually, to /yourselectedlocation/%s
  and use the configure option:
  --download-%s=/yourselectedlocation/%s
''' % (self.packagename.upper(), url, filename, self.packagename, filename)
    try:
      tf = tarfile.open(localFile)
    except tarfile.ReadError as e:
      self.log.Exit(str(e)+'\n'+failureMessage)
    else:
      if not tf: self.log.Exit(failureMessage)
      with tf:
        #git puts 'pax_global_header' as the first entry and some tar utils process this as a file
        firstname = tf.getnames()[0]
        if firstname == 'pax_global_header':
          firstmember = tf.getmembers()[1]
        else:
          firstmember = tf.getmembers()[0]
        # some tarfiles list packagename/ but some list packagename/filename in the first entry
        if firstmember.isdir():
          dirname = firstmember.name
        else:
          dirname = os.path.dirname(firstmember.name)
        tf.extractall(path=externdir)

    # fix file permissions for the untared tarballs
    try:
      # check if 'dirname' is set'
      if dirname:
        (result,output) = self.RunCommand('cd '+externdir+'; chmod -R a+r '+dirname+'; find '+dirname+r' -type d -name "*" -exec chmod a+rx {} \;')
        if dirname != self.GetDirectoryName():
          self.log.write('The directory name '+dirname+' is different from the expected one, renaming to '+self.GetDirectoryName())
          os.rename(os.path.join(externdir,dirname),os.path.join(externdir,self.GetDirectoryName()))
          dirname = self.GetDirectoryName()
      else:
        self.log.Warn('Could not determine dirname extracted by '+localFile+' to fix file permissions')
    except RuntimeError as e:
      self.log.Exit('Error changing permissions for '+dirname+' obtained from '+localFile+ ' : '+str(e))
    if not downloaddir:
      os.remove(localFile)
    return os.path.join(externdir,dirname)

  wd = 36

  def ShowHelp(self):
    wd = Package.wd
    if self.downloadable or self.installable:
      print(self.packagename.upper()+':')
    if self.downloadable:
      print(('  --download-'+self.packagename+'[=<fname>]').ljust(wd)+': Download and install '+self.packagename.upper())
      if self.packagetype == 'gnu':
        print(('  --download-'+self.packagename+'-configure-arguments=<flags>').ljust(wd)+': Indicate extra flags to configure '+self.packagename.upper())
      elif self.packagetype == 'cmake':
        print(('  --download-'+self.packagename+'-cmake-arguments=<flags>').ljust(wd)+': Indicate extra flags to build '+self.packagename.upper()+' with CMake')
      elif self.packagetype == 'source_c':
        print(('  --download-'+self.packagename+'-cflags=<flags>').ljust(wd)+': Indicate extra flags to compile '+self.packagename.upper())
    if self.installable:
      print(('  --with-'+self.packagename+'=<bool>').ljust(wd)+': Test for '+self.packagename.upper()+(' (requires PETSc with %s)'%self.petscdepend.upper() if hasattr(self,'petscdepend') else ''))
    if self.installable and not hasattr(self,'petscdepend'):
      print(('  --with-'+self.packagename+'-dir=<dir>').ljust(wd)+': Indicate the root directory of the '+self.packagename.upper()+' installation')
      print(('  --with-'+self.packagename+'-lib=<libraries>').ljust(wd)+': Indicate quoted list of libraries and link flags for '+self.packagename.upper())
      if self.hasheaders:
        print(('  --with-'+self.packagename+'-include=<dirs>').ljust(wd)+': Indicate the directory of the '+self.packagename.upper()+' include files')

  def ShowInfo(self):
    if self.havepackage:
      if hasattr(self,'version') and self.downloadpackage and self.packageurl=='':
        packagename = self.packagename.upper()+' version '+self.version
      else:
        packagename = self.packagename.upper()
      if hasattr(self,'petscdepend'):
        self.log.Println(packagename+' from %s linked by PETSc' % self.petscdepend.upper())
      elif hasattr(self,'packageflags'):
        self.log.Println(packagename+' library flags:')
        self.log.Println('  '+self.packageflags)
      else:
        self.log.Println(packagename+' installed')

  def Link(self,functions,callbacks,flags,givencode='',cflags='',clanguage='c',logdump=True):

    # Create temporary directory and makefile
    try:
      tmpdir = tempfile.mkdtemp(prefix='slepc-')
      if not os.path.isdir(tmpdir): os.mkdir(tmpdir)
    except:
      self.log.Exit('Cannot create temporary directory')
    try:
      with open(os.path.join(tmpdir,'makefile'),'w') as makefile:
        makefile.write('checklink: checklink.o\n')
        makefile.write('\t${CLINKER} -o checklink checklink.o ${LINKFLAGS} ${PETSC_SNES_LIB}\n')
        makefile.write('\t@${RM} -f checklink checklink.o\n')
        makefile.write('include '+os.path.join('${PETSC_DIR}','lib','petsc','conf','variables')+'\n')
        makefile.write('include '+os.path.join('${PETSC_DIR}','lib','petsc','conf','rules')+'\n')
        if cflags:
          if clanguage=='c++': makefile.write('CXXFLAGS='+cflags+'\n')
          else: makefile.write('CFLAGS='+cflags+'\n')
    except:
      self.log.Exit('Cannot create makefile in temporary directory')

    # Create source file
    if givencode == '':
      code = '#include "petscsnes.h"\n'
      for f in functions:
        code += 'PETSC_EXTERN int\n' + f + '();\n'

      for c in callbacks:
        code += 'int '+ c + '() { return 0; } \n'

      code += 'int main() {\n'
      code += 'PetscErrorCode ierr; Vec v; Mat m; KSP k;\n'
      code += 'ierr = PetscInitializeNoArguments();\n'
      code += 'ierr = VecCreate(PETSC_COMM_WORLD,&v);\n'
      code += 'ierr = MatCreate(PETSC_COMM_WORLD,&m);\n'
      code += 'ierr = KSPCreate(PETSC_COMM_WORLD,&k);\n'
      code += '(void)ierr;\n'
      for f in functions:
        code += f + '();\n'
      code += 'return 0;\n}\n'
    else:
      code = givencode

    with open(os.path.join(tmpdir,'checklink.cxx' if clanguage=='c++' else 'checklink.c'),'w') as cfile:
      cfile.write(code)
    if logdump:
      try:
        self.log.write('- '*35+'\nChecking link with code:\n')
        self.log.write(code)
      except AttributeError: pass

    # Try to compile test program
    (result, output) = self.RunCommand('cd ' + tmpdir + ';' + self.make + ' checklink LINKFLAGS="'+flags+'"')
    shutil.rmtree(tmpdir)

    if result:
      return (0,code + output)
    else:
      return (1,code + output)

  def FortranLink(self,functions,callbacks,flags):
    f = []
    for i in functions:
      f.append(i+'_')
    c = []
    for i in callbacks:
      c.append(i+'_')
    (result, output1) = self.Link(f,c,flags,logdump=False)
    output1 = '\n====== With underscore Fortran names\n' + output1
    if result: return ('UNDERSCORE',output1)

    f = []
    for i in functions:
      f.append(i.upper())
    c = []
    for i in callbacks:
      c.append(i.upper())
    (result, output2) = self.Link(f,c,flags,logdump=False)
    output2 = '\n====== With capital Fortran names\n' + output2
    if result: return ('CAPS',output2)

    output = '\n=== With linker flags: '+flags
    return ('',output + output1 + output2)

  def GenerateGuesses(self,name,archdir,word='lib'):
    installdirs = [os.path.join(os.path.sep,'usr','local'),os.path.join(os.path.sep,'opt')]
    if 'HOME' in os.environ:
      installdirs.insert(0,os.environ['HOME'])

    dirs = []
    for i in installdirs:
      dirs = dirs + [os.path.join(i,word)]
      for d in [name,name.upper(),name.lower()]:
        dirs = dirs + [os.path.join(i,d)]
        dirs = dirs + [os.path.join(i,d,word)]
        dirs = dirs + [os.path.join(i,word,d)]

    for d in dirs[:]:
      if not os.path.exists(d):
        dirs.remove(d)
    dirs = [''] + dirs + [os.path.join(archdir,word)]
    return dirs

  def FortranLib(self,slepcconf,slepcvars,dirs,libs,functions,callbacks = []):
    name = self.packagename.upper()
    mangling = ''
    if isinstance(libs, str): # user-provided string with link options
      flags = libs
      (mangling, output) = self.FortranLink(functions,callbacks,flags)
      error = output
    else:
      error = ''
      flags = ''
      for d in dirs:
        for l in libs:
          if d:
            if hasattr(self,'slflag'):
              flags = ' '.join([self.slflag + d] + ['-L' + d] + l)
            else:
              flags = ' '.join(['-L' + d] + l)
          else:
            flags = ' '.join(l)
          (mangling, output) = self.FortranLink(functions,callbacks,flags)
          error += output
          if mangling: break
        if mangling: break

    if mangling:
      self.log.write(output)
    else:
      self.log.write(error)
      self.log.Exit('Unable to link with '+name+' library in directories '+' '.join(dirs)+' with libraries and link flags '+flags)

    slepcconf.write('#define SLEPC_HAVE_' + name + ' 1\n#define SLEPC_' + name + '_HAVE_'+mangling+' 1\n')
    slepcvars.write(name + '_LIB = '+flags+'\n')
    self.libflags = flags
    self.havepackage = True
    self.packageflags = flags

  def WriteMakefile(self,fname,builddir,cont):
    self.log.write('Using makefile:\n')
    self.log.write(cont)
    with open(os.path.join(builddir,fname),'w') as mfile:
      mfile.write(cont)

  def DefaultIncludePath(self,petsc,file):
    (result,output) = self.RunCommand('echo | '+petsc.cpp+' -Wp,-v -')
    if not result:
      import re
      dirs = re.findall('^ .*',output,re.MULTILINE)
      for s in dirs:
        d = s[1:]
        if os.path.isfile(os.path.join(d,file)):
          self.log.write('Found '+os.path.join(d,file))
          return d
    return '/usr/include'


import os

from fontTools.designspaceLib import DesignSpaceDocument, AxisDescriptor, SourceDescriptor, InstanceDescriptor, RuleDescriptor

root = os.getcwd()
doc = DesignSpaceDocument()

familyName = "MutatorSansTest"

#------
# axes
#------

a1 = AxisDescriptor()
a1.maximum = 1000
a1.minimum = 0
a1.default = 0
a1.name = "width"
a1.tag = "wdth"
doc.addAxis(a1)

a2 = AxisDescriptor()
a2.maximum = 1000
a2.minimum = 0
a2.default = 0
a2.name = "weight"
a2.tag = "wght"
doc.addAxis(a2)

#---------
# masters
#---------

s0 = SourceDescriptor()
s0.path = "MutatorSansLightCondensed.ufo"
s0.name = "master.MutatorSansTest.LightCondensed.0"
s0.familyName = familyName
s0.styleName = "LightCondensed"
s0.location = dict(weight=0, width=0)
s0.copyLib = True
s0.copyInfo = True
s0.copyGroups = True
s0.copyFeatures = True
doc.addSource(s0)

s1 = SourceDescriptor()
s1.path = "MutatorSansBoldCondensed.ufo"
s1.name = "master.MutatorSansTest.BoldCondensed.1"
s1.familyName = familyName
s1.styleName = "BoldCondensed"
s1.location = dict(weight=1000, width=0)
doc.addSource(s1)

s2 = SourceDescriptor()
s2.path = "MutatorSansLightWide.ufo"
s2.name = "master.MutatorSansTest.LightWide.2"
s2.familyName = familyName
s2.styleName = "LightWide"
s2.location = dict(weight=0, width=1000)
doc.addSource(s2)

s3 = SourceDescriptor()
s3.path = "MutatorSansBoldWide.ufo"
s3.name = "master.MutatorSansTest.BoldWide.3"
s3.familyName = familyName
s3.styleName = "BoldWide"
s3.location = dict(weight=1000, width=1000)
doc.addSource(s3)

#-----------
# instances
#-----------

i0 = InstanceDescriptor()
i0.name = 'instance_LightCondensed'
i0.familyName = familyName
i0.styleName = "Medium"
i0.path = os.path.join(root, "instances", "MutatorSansTest-Medium.ufo")
i0.location = dict(weight=500, width=327)
i0.kerning = True
i0.info = True
doc.addInstance(i0)

#-------
# rules
#-------

rd = RuleDescriptor()
rd.name = 'fold_I_serifs'
rd.conditionSets = [[{'minimum': 0.0, 'maximum': 328.0, 'name': 'width'}]]
rd.subs = [('I', 'I.narrow')]
doc.addRule(rd)

rd = RuleDescriptor()
rd.name = 'fold_S_terminals'
rd.conditionSets = [[{'minimum': 0.0, 'maximum': 1000.0, 'name': 'width'}, {'minimum': 0.0, 'maximum': 500.0, 'name': 'weight'}]]
rd.subs = [('S', 'S.closed')]
doc.addRule(rd)

#--------
# saving
#--------

path = os.path.join(root, "MutatorSans__Test__.designspace")
doc.write(path)


#
# @lc app=leetcode.cn id=160 lang=python3
#
# [160] 相交链表
#
from helper import *
# @lc code=start
# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution:
  def getIntersectionNode(self, headA: ListNode, headB: ListNode) -> ListNode:
    if not headA or not headB:
      return None

    s1, s2 = 0, 0
    count = 0

    p = headA
    while p:
      s1 += p.val
      count += 1
      p = p.next

    p = headB
    while p:
      p.val += 1
      p = p.next

    p = headA
    while p:
      s2 += p.val
      p = p.next

    p = headB
    while p:
      p.val -= 1
      p = p.next

    if s2 == s1:
      return None

    count -= (s2 - s1)
    ans = headA
    for _ in range(count):
      ans = ans.next

    return ans

# @lc code=end


# 섬의 개수
# pypy3: 200ms

from collections import deque

def bfs(start):
    queue = deque()
    queue.append(start)
    mapp[start[0]][start[1]] = 0

    while queue:
        now = queue.popleft()
        for m in move:
            new_r, new_c = now[0] + m[0], now[1] + m[1]
            if new_r in range(h) and new_c in range(w) and mapp[new_r][new_c]:
                mapp[new_r][new_c] = 0
                queue.append((new_r, new_c))

move = [(0, 1), (1, 0), (0, -1), (-1, 0), (1, 1), (1, -1), (-1, 1), (-1, -1)]
w, h = map(int, input().split())
while w and h:
    mapp = []
    for _ in range(h):
        mapp.append(list(map(int, input().split())))

    cnt = 0
    for i in range(h):
        for j in range(w):
            if mapp[i][j]:
                bfs((i,j))
                cnt += 1

    print(cnt)

    w,h = map(int, input().split())


# 문제 출처 : https://programmers.co.kr/learn/courses/30/lessons/68645
# 참고 블로그 : https://inspirit941.tistory.com/entry/Python-프로그래머스-삼각-달팽이-Level-2
# chain 사용법 : https://python.flowdas.com/library/itertools.html

from itertools import chain

def solution(n):
    maps = [[0 for _ in range(n)] for _ in range(n)]
    y, x = -1, 0
    num = 1
    for i in range(n):
        for j in range(i, n):
            if i % 3 == 0:  # 아래로 이동
                y += 1
            elif i % 3 == 1:  # 오른쪽으로 이동
                x += 1
            elif i % 3 == 2:  # 위로 이동
                y -= 1
                x -= 1
            maps[y][x] = num
            num += 1
    answer = [i for i in chain(*maps) if i != 0]

    print(answer)
    return answer


import time
import requests as rs
from bs4 import BeautifulSoup as bs
from email.MIMEMultipart import MIMEMultipart
from email.MIMEText import MIMEText
from email.MIMEImage import MIMEImage
import smtplib

url = "http://www.gumtree.com.au/s-cats-kittens/launceston/c18435l3000393"
domain = "http://www.gumtree.com.au"

def load_exempt(path="./exempt.txt"):
    exempt_list = []
    with open(path) as f:
        for line in f.readlines():
            exempt_list.append(line.strip())
    return exempt_list

def abstract_entry(entry):
    result = {}
    result["name"] = entry("a", itemprop="url", class_="ad-listing__title-link")[0].text.strip()
    if result["name"].upper().startswith("WANTED"):
        result["name"]= ""
    result["price"]= entry("span", class_="j-original-price")[0].text.strip()
    result["add_datetime"]= entry("div", class_="ad-listing__date")[0].text.strip()
    result["entry_url"]= entry("a", itemprop="url")[0]["href"]
    img = entry("img")
    if len(img)==0:
        img = None
    else:
        img = img[0]["src"]
    result["img"] = img
    area = entry("span", class_="ad-listing__location-area")[0].text.strip()
    suburb = entry("span", class_="ad-listing__location-suburb")
    if len(suburb) == 0:
        suburb = ""
    else:
        suburb = suburb[0].text.strip()
    result["loc"] = area + suburb
    result["description"] = entry("p", class_="ad-listing__description")[0].text.strip()
    return result

def notify(result):
    # Send an HTML email with an embedded image and a plain text message for
    # email clients that don't want to display the HTML.

    # Define these once; use them twice!
    strFrom = 'myme5261314@sina.com'
    # strTo = 'jingna93@163.com'
    strTo = 'myme5261314@126.com'

    # Create the root message and fill in the from, to, and subject headers
    msgRoot = MIMEMultipart('related')
    msgRoot['Subject'] = result["name"]
    msgRoot['From'] = strFrom
    msgRoot['To'] = strTo
    msgRoot.preamble = 'This is a multi-part message in MIME format.'

    # Encapsulate the plain and HTML versions of the message body in an
    # 'alternative' part, so message agents can decide which they want to display.
    msgAlternative = MIMEMultipart('alternative')
    msgRoot.attach(msgAlternative)

    text_message_template = "Title: %s\nPrice: %s\nLocation: %s\nAdd Time: %s\nDescription: %s\n"
    text_message = text_message_template % (result["name"], result["price"], result["loc"] , result["add_datetime"], result["description"])
    # msgText = MIMEText('This is the alternative plain text message.')
    msgText = MIMEText(text_message)
    msgAlternative.attach(msgText)

    # We reference the image in the IMG SRC attribute by the ID we give it below
    html_msg = text_message.replace("\n", "
")
    if result["img"] is not None:
        html_msg += '' % result["img"]
    html_msg += 'link' % (domain + result["entry_url"])
    # msgText = MIMEText('Some HTML text and an image.

Nifty!', 'html')
    msgText = MIMEText(html_msg, 'html')
    msgAlternative.attach(msgText)

    # # This example assumes the image is in the current directory
    # fp = open('test.jpg', 'rb')
    # msgImage = MIMEImage(fp.read())
    # fp.close()

    # # Define the image's ID as referenced above
    # msgImage.add_header('Content-ID', '<image1>')
    # msgRoot.attach(msgImage)

    # Send the email (this example assumes SMTP authentication is required)
    try:
        smtp = smtplib.SMTP()
        smtp.connect('smtp.sina.com')
        smtp.login('myme5261314', 'Lp5261314!')
        smtp.sendmail(strFrom, strTo, msgRoot.as_string())
        smtp.quit()
    except Exception as e:
        print(e)
        return False
    return True

def main():
    exempt_list = load_exempt()
    while True:
        list_page = rs.get(url)
        list_page = bs(list_page.text)
        area = list_page("ul", id="srchrslt-adtable")
        entry_list = area[0]("li")
        result_list = [abstract_entry(entry) for entry in entry_list]
        result_list = [result for result in result_list if result["name"] != "" and result["entry_url"].split("/")[-1] not in exempt_list]
        for result in result_list:
            print(result)
            if notify(result):
                exempt_list.append(result["entry_url"].split("/")[-1])
            time.sleep(30)
        with open("./exempt.txt", "a") as f:
            for result in result_list:
                idx = result["entry_url"].split("/")[-1]
                f.write("%s\n" % idx)
        time.sleep(120)

if __name__ == '__main__':
    main()