import os
import re
from PIL import Image
from datetime import datetime
from django.contrib.contenttypes.models import ContentType
from django.db import models
from django.utils.timezone import now
from django.utils.html import format_html
from django.utils.text import slugify
from filebrowser.fields import FileBrowseField
from taggit.managers import TaggableManager
from streamfield.fields import StreamField
from streamblocks.models import (IndexedParagraph, CaptionedImage, Gallery,
    LandscapeGallery, DownloadableFile, LinkableList, BoxedText, EventUpgrade)
from .choices import *
from users.models import User, Member

def date_directory_path(instance, filename):
    if instance.date:
        now = instance.date
    else:
        now = datetime.now()
    year = now.strftime("%Y")
    month = now.strftime("%m")
    day = now.strftime("%d")
    return 'uploads/{0}/{1}/{2}/{3}'.format(year, month, day, filename)

def generate_unique_slug(klass, field):
    """
    return unique slug if origin slug exists.
    eg: `foo-bar` => `foo-bar-1`

    :param `klass` is Class model.
    :param `field` is specific field for title.
    Thanks to djangosnippets.org!
    """
    origin_slug = slugify(field)
    unique_slug = origin_slug
    numb = 1
    while klass.objects.filter(slug=unique_slug).exists():
        unique_slug = '%s-%d' % (origin_slug, numb)
        numb += 1
    return unique_slug

class Location(models.Model):
    fb_image = FileBrowseField("Immagine", max_length=200,
        directory="locations/", extensions=[".jpg", ".png", ".jpeg", ".gif",
        ".tif", ".tiff"], null=True, blank=True)
    title = models.CharField('Titolo',
        help_text='Il nome del luogo',
        max_length = 50)
    slug = models.SlugField(max_length=50, unique=True)
    address = models.CharField('Indirizzo', max_length = 200,
        help_text = 'Via/Piazza, civico, CAP, Città',)
    gmap_link = models.URLField('Link di Google Map',
        blank= True, null=True,
        help_text="Dal menu di Google Maps seleziona 'Condividi/link', \
                   copia il link e incollalo qui",
    )
    gmap_embed = models.TextField('Incorpora Google Map',
        blank= True, null=True, max_length=500,
        help_text="Dal menu di Google Maps seleziona 'Condividi/incorpora', \
                   copia il link e incollalo qui",
    )
    body = models.TextField('Descrizione', blank= True, null=True,)
    website = models.URLField('Sito internet',
        blank= True, null=True, )
    email = models.EmailField(blank = True, null = True,
        verbose_name = 'Email',)
    phone = models.CharField(max_length = 50,
        blank = True, null = True, verbose_name = 'Telefono/i',)

    def save(self, *args, **kwargs):
        if not self.gmap_embed.startswith('http'):
            # thanks to geeksforgeeks.com! findall returns a list
            list = re.findall('http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\), ]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', self.gmap_embed)
            if list:
                self.gmap_embed = list[0]
        if not self.slug:  # create
            self.slug = generate_unique_slug(Location, self.title)
        super(Location, self).save(*args, **kwargs)

    def get_gmap_link(self):
        if self.gmap_link:
            link = format_html('Mappa',
                self.gmap_link)
        else:
            link = '-'
        return link
    get_gmap_link.short_description = 'Link di Google Maps'

    def __str__(self):
        return self.title

    class Meta:
        verbose_name = 'Luogo'
        verbose_name_plural = 'Luoghi'
        ordering = ('id', )

def update_indexed_paragraphs(stream_list, type, id):
    for block in stream_list:
        if block['model_name'] == 'IndexedParagraph':
            par = IndexedParagraph.objects.get(id = block['id'])
            par.parent_type = type
            par.parent_id = id
            par.save()

class Event(models.Model):
    fb_image = FileBrowseField("Immagine", max_length=200, directory="events/",
        extensions=[".jpg", ".png", ".jpeg", ".gif", ".tif", ".tiff"],
        null=True, blank=True)
    carousel = StreamField(model_list=[ LandscapeGallery, ],
        null=True, blank=True, verbose_name="Galleria",
        help_text="Una sola galleria, per favore, larghezza minima immagini 2048px")
    title = models.CharField('Titolo',
        help_text="Il titolo dell'evento",
        max_length = 50)
    slug = models.SlugField(max_length=50, editable=False, null=True)
    date = models.DateTimeField('Quando', default = now)
    last_updated = models.DateTimeField(editable=False, null=True)
    location = models.ForeignKey(Location, on_delete=models.SET_NULL,
        null = True, verbose_name = 'Dove', )
    intro = models.CharField('Introduzione',
        default = 'Un altro appuntamento con RP!', max_length = 100)
    stream = StreamField(model_list=[ IndexedParagraph, CaptionedImage,
        Gallery, DownloadableFile, LinkableList, BoxedText, ],
        verbose_name="Lancio")
    upgrade_stream = StreamField(model_list=[ EventUpgrade, IndexedParagraph,
        DownloadableFile, ],
        verbose_name="Aggiornamenti")
    chron_stream = StreamField(model_list=[ IndexedParagraph, CaptionedImage,
        Gallery, DownloadableFile, LinkableList, BoxedText],
        verbose_name="Cronaca")
    restr_stream = StreamField(model_list=[ IndexedParagraph, CaptionedImage,
        Gallery, DownloadableFile, LinkableList, BoxedText],
        verbose_name="Area riservata",
        help_text="Inserisci qui materiale riservato ai soci",)
    manager = models.ForeignKey(User, on_delete=models.SET_NULL,
        blank= True, null=True, verbose_name = 'Responsabile')
    tags = TaggableManager(verbose_name="Categorie",
        help_text="Lista di categorie separate da virgole",
        through=None, blank=True)
    notice = models.CharField(max_length = 4, choices = NOTICE,
        blank = True, null = True, verbose_name = 'Notifica via email',
        help_text = """Non invia in automatico, per farlo seleziona l'Evento
            dalla Lista degli Eventi, imposta l'azione 'Invia notifica' e fai
            clic su 'Vai'.
            """)

    def get_badge_color(self):
        if self.date.date() > datetime.today().date():
            return 'success'
        elif self.date.date() < datetime.today().date():
            return 'secondary'
        else:
            return 'warning'

    def get_image(self):
        if self.fb_image:
            return self.fb_image
        elif self.location.fb_image:
            return self.location.fb_image
        return

    def get_tags(self):
        return list(self.tags.names())

    def get_upgrades(self):
        return EventUpgrade.objects.filter(event_id=self.id)

    def get_chronicle(self):
        if self.date.date() < datetime.today().date():
            return True
        return False

    def get_uploads(self):
        return UserUpload.objects.filter(event_id=self.id)

    def get_path(self):
        return '/calendario/' + self.date.strftime("%Y/%m/%d") + '/' + self.slug

    def save(self, *args, **kwargs):
        if not self.slug:
            self.slug = generate_unique_slug(Event, self.title)
        if not self.notice:
            self.notice = 'SPAM'
        self.last_updated = now()
        super(Event, self).save(*args, **kwargs)
        #update parent_type end parent_id in IndexedParagraph streamblocks
        type = ContentType.objects.get(app_label='pagine', model='event').id
        id = self.id
        stream_list = self.stream.from_json()
        update_indexed_paragraphs(stream_list, type, id)
        stream_list = self.upgrade_stream.from_json()
        update_indexed_paragraphs(stream_list, type, id)
        stream_list = self.chron_stream.from_json()
        update_indexed_paragraphs(stream_list, type, id)
        stream_list = self.restr_stream.from_json()
        update_indexed_paragraphs(stream_list, type, id)

    def __str__(self):
        return self.title

    class Meta:
        verbose_name = 'Evento'
        verbose_name_plural = 'Eventi'
        ordering = ('-date', )

class EventUpgrade(models.Model):
    event = models.ForeignKey(Event, on_delete = models.CASCADE,
        null = True, related_name='event_upgrades')
    title = models.CharField('Titolo',
        help_text="Il titolo dell'aggiornamento",
        max_length = 50)
    date = models.DateTimeField('Data', default = now)
    body = models.TextField('Aggiornamento',
        help_text = "Accetta tag HTML.", )

    def __str__(self):
        return self.title

    class Meta:
        verbose_name = 'Aggiornamento'
        verbose_name_plural = 'Aggiornamenti'
        ordering = ('-date', )

class Blog(models.Model):
    fb_image = FileBrowseField("Immagine", max_length=200, directory="blogs/",
        extensions=[".jpg", ".png", ".jpeg", ".gif", ".tif", ".tiff"],
        null=True, blank=True)
    carousel = StreamField(model_list=[ LandscapeGallery, ],
        null=True, blank=True, verbose_name="Galleria",
        help_text="Una sola galleria, per favore, larghezza minima immagini 2048px")
    title = models.CharField('Titolo',
        help_text="Il titolo dell'articolo",
        max_length = 50)
    slug = models.SlugField(max_length=50, editable=False, null=True)
    date = models.DateTimeField('Data', default = now)
    last_updated = models.DateTimeField(editable=False, null=True)
    intro = models.CharField('Introduzione',
        default = 'Un altro articolo di approfondimento da RP!', max_length = 100)
    stream = StreamField( model_list=[ IndexedParagraph, CaptionedImage,
            Gallery, DownloadableFile, LinkableList, BoxedText],
            verbose_name="Testo" )
    author = models.ForeignKey(User, on_delete=models.SET_NULL,
        blank= True, null=True, verbose_name = 'Autore')
    tags = TaggableManager(verbose_name="Categorie",
        help_text="Lista di categorie separate da virgole",
        through=None, blank=True)

    def get_path(self):
        return '/articoli/' + self.slug

    def get_uploads(self):
        return UserUpload.objects.filter(post_id=self.id)

    def get_tags(self):
        return list(self.tags.names())

    def save(self, *args, **kwargs):
        if not self.slug:
            self.slug = generate_unique_slug(Blog, self.title)
        self.last_updated = now
        super(Blog, self).save(*args, **kwargs)
        #update parent_type end parent_id in IndexedParagraph streamblocks
        type = ContentType.objects.get(app_label='pagine', model='blog').id
        id = self.id
        stream_list = self.stream.from_json()
        update_indexed_paragraphs(stream_list, type, id)

    def __str__(self):
        return self.title

    class Meta:
        verbose_name = 'Articolo'
        verbose_name_plural = 'Articoli'
        ordering = ('-date', )

class UserUpload(models.Model):
    event = models.ForeignKey(Event, on_delete = models.CASCADE,
        null = True, related_name='event_uploads')
    post = models.ForeignKey(Blog, on_delete = models.CASCADE,
        null = True, related_name='blog_uploads')
    user = models.ForeignKey(User, on_delete=models.SET_NULL, null = True,
        verbose_name = 'Utente')
    date = models.DateTimeField('Data', default = now, )
    image = models.ImageField('Immagine', blank = True, null = True,
        upload_to = date_directory_path,)
    body = models.TextField('Testo', help_text = "Scrivi qualcosa.", )

    def __str__(self):
        return 'Contributo - ' + str(self.id)

    class Meta:
        verbose_name = 'Contributo'
        verbose_name_plural = 'Contributi'
        ordering = ('-id', )


from collections import defaultdict

puzzle = open('puzzle', 'r').read().splitlines()
grid = defaultdict(str)

for y in range(len(puzzle)):
	for x in range(len(puzzle[y])):
		grid[(y, x)] = puzzle[y][x]

def surrounding(cords, grid):
	ret = []
	for y in range(cords[0]-1, cords[0]+2):
		for x in range(cords[1]-1, cords[1]+2):
			if (y, x) != cords:
				ret.append(grid[(y, x)])
	return ret

minutes = 1000000000
seen = []
seen.append(grid.copy())
cycle = 0

for m in range(minutes):
	old_grid = grid.copy()
	for y in range(len(puzzle)):
		for x in range(len(puzzle[y])):
			surr = surrounding((y, x), old_grid)
			this = grid[(y, x)]
			if this == '.' and surr.count('|') >= 3:
				grid[(y, x)] = '|'
			elif this == '|' and surr.count('#') >= 3:
				grid[(y, x)] = '#'
			elif this == '#' and (surr.count('#') == 0 or surr.count('|') == 0):
				grid[(y, x)] = '.'

	if grid in seen:
		cycle = minutes % (m-seen.index(grid))
		break
	seen.append(grid.copy())

for m in range(cycle):
	old_grid = grid.copy()
	for y in range(len(puzzle)):
		for x in range(len(puzzle[y])):
			surr = surrounding((y, x), old_grid)
			this = grid[(y, x)]
			if this == '.' and surr.count('|') >= 3:
				grid[(y, x)] = '|'
			elif this == '|' and surr.count('#') >= 3:
				grid[(y, x)] = '#'
			elif this == '#' and (surr.count('#') == 0 or surr.count('|') == 0):
				grid[(y, x)] = '.'

print(list(grid.values()).count('|') * list(grid.values()).count('#'))


# !/usr/bin/env python
# -*- coding: utf-8 -*-

"""
@author: lishuang
@description: 基于能力描述的薪资预测
数据集：抓取了4512个职位的能力描述，薪资
Step1，数据加载
Step2，可视化，使用Networkx
Step3，提取文本特征 TFIDF
Step4，回归分析，使用KNN回归，朴素贝叶斯回归，训练能力和薪资匹配模型
Step5，基于指定的能力关键词，预测薪资
"""
import random
import re

import jieba
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.neighbors import KNeighborsRegressor

# 显示所有列
pd.set_option('display.max_columns', None)

# 用来正常显示负号
plt.rcParams['axes.unicode_minus'] = False

def handle_job_string(row):
    """
    处理每一行数据

    :param row:
    :return:
    """
    job_string = ''
    for idx, element in enumerate(row.split('\n')):
        if len(element.split()) == 2:
            idx, value = element.split()
            if idx == 0:
                continue
            job_string += value
    return job_string

def predict_by_label(test_string, vectorizer, model):
    """
    通过能力标签预测薪资

    :param test_string:
    :param vectorizer:
    :param model:
    :return:
    """
    test_words = list(jieba.cut(test_string))
    test_vec = vectorizer.transform(test_words)
    predict_value = model.predict(test_vec)
    return predict_value[0]

def main():
    # 1. 数据加载
    data = pd.read_excel('./data/jobs_4k.xls')
    print(data.head())

    # 2. 观察职位和技能之间的关联度，并可视化
    position_names = data['positionName'].tolist()
    skill_labels = data['skillLabels'].tolist()

    # 构建职位技能字典
    position_to_skill_dict = dict()
    for p, s in zip(position_names, skill_labels):
        if position_to_skill_dict.get(p) is None:
            position_to_skill_dict[p] = eval(s)
        else:
            position_to_skill_dict[p] += eval(s)
    print(position_to_skill_dict)

    # 随机选择k个工作岗位
    sample_nodes = random.sample(position_names, k=5)
    print(f'随机选择k个工作岗位: \n{sample_nodes}')

    # 将职位信息和能力描述结合到一起
    sample_nodes_connections = sample_nodes
    for p, s in position_to_skill_dict.items():
        if p in sample_nodes:
            sample_nodes_connections += s
    print(f'将职位和能力放在一起: \n{sample_nodes_connections}')

    # 绘制图像
    G = nx.Graph(position_to_skill_dict)

    # 抽取原始G中的节点作为子图
    sample_graph = G.subgraph(sample_nodes_connections)
    plt.figure(figsize=(16, 8))
    pos = nx.spring_layout(sample_graph, k=1)
    nx.draw(sample_graph, pos, with_labels=True)
    # plt.show()

    # 使用PageRank计算节点（技能）影响力
    pr = nx.pagerank(G, alpha=0.9)
    ranked_position_skill = sorted([(name, value) for name, value in pr.items()], key=lambda _: _[1], reverse=True)
    print(f'排序后的职位和技能: \n{ranked_position_skill}')

    # 3. 构造特征，使用TF-IDF提取文本特征
    data_X = data.drop(['salary'], axis=1).copy()
    targets = data['salary'].tolist()

    # 将所有的特征放在一起
    data_X['merged'] = data_X.apply(lambda _: ''.join(str(_)), axis=1)

    # 分词处理
    cut_X = list()
    for idx, row in enumerate(data_X['merged']):
        job_string = handle_job_string(row)
        cut_X.append(' '.join(list(jieba.cut(''.join(re.findall('\\w+', job_string))))))
    print(f'分词处理后的数据: \n{cut_X[0]}')

    # 使用 TF-IDF 提取文本特征
    vectorizer = TfidfVectorizer()
    X = vectorizer.fit_transform(cut_X)
    print(f'提取文本特征后的数据: \n{X[0]}')

    # 求平均值 薪资 10k - 15k => 12.5k
    target_numerical = [np.mean(list(map(float, re.findall('\\d+', target)))) for target in targets]
    Y = target_numerical

    # 4. 回归分析，使用KNN回归，训练能力和薪资匹配模型
    model = KNeighborsRegressor(n_neighbors=2)
    model.fit(X, Y)
    print(f'KNN模型评分为: {model.score}')

    # 5. 基于指定的能力关键词，预测薪资
    test_string_list = [
        '测试 北京 3年 专科',
        '测试 北京 4年 专科',
        '算法 北京 4年 本科',
        'UI 北京 4年 本科',
        '广州Java本科3年掌握大数据',
        '沈阳Java硕士3年掌握大数据',
        '沈阳Java本科3年掌握大数据',
        '北京算法硕士3年掌握图像识别'
    ]

    for test_string in test_string_list:
        print(f'职位和能力信息: {test_string}, 预测的薪资为: {predict_by_label(test_string, vectorizer, model)}')

if __name__ == '__main__':
    main()


import pandas as pd
import pathlib
from model import toa

datadir = pathlib.Path(__file__).parents[0].joinpath('data')

def test_toa_per_region():
    # we use a modified version of the silvopasture TLA data as mock TLA values
    sp_land_dist_all = [331.702828, 181.9634517, 88.98630743, 130.15193962, 201.18287123,
                        0.,  # new ABNJ region
                        933.98739798, 37.60589239, 7.02032552, 43.64118779, 88.61837725]
    index = pd.Index(
        ['OECD90', 'Eastern Europe', 'Asia (Sans Japan)', 'Middle East and Africa', 'Latin America', 'ABNJ', 'Global',
         'China', 'India', 'EU', 'USA'])
    ocean_dist = pd.DataFrame(sp_land_dist_all, columns=['All'], index=index)
    expected = pd.read_csv(datadir.joinpath('sp_tla_with_abnj.csv'), index_col=0)
    result = toa.toa_per_region(ocean_dist=ocean_dist)
    pd.testing.assert_frame_equal(expected, result, check_dtype=False)


#!/usr/bin/env python3

from dialog.parser import Parser
from dialog.scope import Scope
from dialog.returns import Returns
import dialog.link_parser

import multiprocessing
import json

from dialog.interpreter import Dialog

class Instance(Dialog):
    """
    Dialog interperter connected with websockets.
    """
    def start_server(self, to_user, to_ds):
        """
        Interprets dialog
        """
        # occupation = multiprocessing.Event()
        # listener_queue = multiprocessing.Queue(maxsize=0)
        # recognizer_queue = multiprocessing.Queue(maxsize=0)
        # speaker_queue = multiprocessing.Queue(maxsize=0)
        # speaker = multiprocessing.Process(
        #     target=speech.speaker,
        #     args=(occupation, speaker_queue, ))
        # recognizer = multiprocessing.Process(
        #     target=speech.recognizer,
        #     args=(recognizer_queue, listener_queue, ))
        # listener = multiprocessing.Process(
        #     target=speech.listener,
        #     args=(occupation, recognizer_queue, ))
        # speaker.start()
        # recognizer.start() # IN CASE OF SPEECH RECOGNITION
        # listener.start()   # IN CASE OF SPEECH RECOGNITION
        # occupation.set()
        # print("======")
        # for state in self.expected:
        #     print("\t%s" % (state))
        # print("======")
        while True:
            # process routines answers
            answers = self.returns.get_returns()
            for answer in answers:
                tosay, questions = answer.accept()
                to_user.put(json.dumps({'type': 'phrase', 'text': tosay}))
                self._extend_expected(questions)
            # process input
            if not to_ds.empty():
                input_origin = to_ds.get()
                # input_phrase = listener_queue.get() # IN CASE OF SPEECH RECOGNITION
                input_phrase = link_parser.parse(input_origin)
                states_probability = []
                for state in self.expected:
                    # print(state, state.compare(input_phrase))
                    states_probability.append((state, state.compare(input_phrase)))
                states_probability = sorted(states_probability, key=lambda x: x[1], reverse=True)
                # print("======")
                # for state in states_probability:
                #     print("%.2f\t%s" % (state[1], state[0]))
                # print("======")
                state = states_probability[0][0]
                if states_probability[0][1] < 0.2:
                    # print("Bot> ???")
                    to_user.put(json.dumps({
                        'type': 'origin',
                        'text': input_origin
                    }))
                    to_user.put(json.dumps({'type': 'unknown'}))
                    #TODO: save unknown request, with state!!!
                else:
                    to_user.put(json.dumps({
                        'type': 'interpretation',
                        'origin': input_origin,
                        'phrase': str(state),
                        'similarity': states_probability[0][1]
                    }))
                    tosay, questions = state.accept(input_phrase)
                    for answer in tosay:
                        if answer != "":
                            # speaker_queue.put(answer)
                            # print("Bot> "+answer)
                            to_user.put(json.dumps({'type': 'phrase', 'text': answer}))
                    self._extend_expected(questions)

from main_slave import export

py_filename = "main_slave.py"
ddl_filename = "examples/demo.dlg"

import tornado.httpserver
import tornado.websocket
import tornado.ioloop
import tornado.web

import time, os.path

class WSHandler(tornado.websocket.WebSocketHandler):
    def check_origin(self, origin):
        return True

    def open(self):
        # probably start handlers
        pass

        # print("starting dialog system")
        # self.write_message("Starting dialog system")
        dialog = Instance(export)
        # self.write_message("Using dialog description from %s" \
        #     % (dialog_description))
        dialog.load(ddl_filename)
        self.to_user = multiprocessing.Queue(maxsize=0)
        self.to_ds = multiprocessing.Queue(maxsize=0)
        self.dialog_system = multiprocessing.Process(
            target=dialog.start_server,
            args=(self.to_user, self.to_ds, ))
        self.dialog_system.start()
        self.producer_process = multiprocessing.Process(
            target=self.producer,
            args=())
        self.producer_process.start()

    def on_message(self, message):
        message = json.loads(message)
        if message['type'] == 'phrase':
            print(">", message)
            self.to_ds.put(message['text'])
        elif message['type'] == 'sources':
            global py_file
            global ddl_file
            ddl_modified = int(os.path.getmtime(ddl_filename))
            py_modified = int(os.path.getmtime(py_filename))
            ddl_content = open(ddl_filename, "rb").read().decode("utf-8")
            py_content = open(py_filename, "rb").read().decode("utf-8")
            response = {
                "type": "sources",
                "modified": max(py_modified, ddl_modified) * 1000,
                "description": {
                    "filename": ddl_filename,
                    "content":  ddl_content,
                    "modified": ddl_modified * 1000
                },
                "code": {
                    "filename": py_filename,
                    "content":  py_content,
                    "modified": py_modified * 1000
                }
            }
            self.to_user.put(json.dumps(response))
        elif message['type'] == 'interpretation':
            print(">", message)
            # print(message['origin'])
            #TODO: save reports

    def on_close(self):
        self.producer_process.terminate()
        self.dialog_system.terminate()

    def producer(self):
        while True:
            message = self.to_user.get()
            time.sleep(0.5)
            print("<", message)
            self.write_message(message)

# class MainHandler(tornado.web.RequestHandler):
#     def get(self):
#         self.render('index.html')

class Application(tornado.web.Application):
    def __init__(self):
        handlers = (
            # (r'/', MainHandler),
            (r'/api/?', WSHandler),
        )
        tornado.web.Application.__init__(self, handlers)

application = Application()
http_server = tornado.httpserver.HTTPServer(application)
http_server.listen(8888)
tornado.ioloop.IOLoop.instance().start()


from django import forms
from django.contrib import admin

from .models import Category, Comment, Genre, Review, Title, User

EMPTY = '-пусто-'

class UserAdminForm(forms.ModelForm):
    def __init__(self, *args, **kwargs):
        super(UserAdminForm, self).__init__(*args, **kwargs)
        self.fields['password'].widget = forms.PasswordInput()

    class Meta:
        model = User
        fields = '__all__'

@admin.register(User)
class UserAdmin(admin.ModelAdmin):
    list_display = (
        'username', 'email', 'first_name', 'last_name', 'role',
        'is_superuser', 'is_staff', 'date_joined',
    )
    empty_value_display = EMPTY
    form = UserAdminForm

@admin.register(Title)
class TitleAdmin(admin.ModelAdmin):
    list_display = ('pk', 'name', 'year', 'description', 'category')
    list_editable = ('category',)
    search_fields = ('name', 'description')
    empty_value_display = EMPTY

@admin.register(Category)
class CategoryAdmin(admin.ModelAdmin):
    list_display = ('pk', 'name', 'slug')
    empty_value_display = EMPTY

@admin.register(Genre)
class GenreAdmin(admin.ModelAdmin):
    list_display = ('pk', 'name', 'slug')
    empty_value_display = EMPTY

@admin.register(Review)
class ReviewAdmin(admin.ModelAdmin):
    list_display = ('pk', 'title', 'text', 'author', 'pub_date')
    empty_value_display = EMPTY

@admin.register(Comment)
class CommentAdmin(admin.ModelAdmin):
    list_display = ('pk', 'text', 'author', 'review', 'pub_date')
    empty_value_display = EMPTY
    search_fields = ('text',)
    list_filter = ('pub_date',)


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"]


import copy
import fnmatch
import logging
import re
from collections import namedtuple
from enum import Enum
from typing import Dict, Iterable, Optional, Union

import yaml
from meltano.core.behavior import NameEq
from meltano.core.behavior.canonical import Canonical
from meltano.core.behavior.hookable import HookObject
from meltano.core.setting_definition import SettingDefinition
from meltano.core.utils import NotFound, compact, find_named, flatten

logger = logging.getLogger(__name__)

class VariantNotFoundError(Exception):
    def __init__(self, plugin: "PluginDefinition", variant_name):
        self.plugin = plugin
        self.variant_name = variant_name

    def __str__(self):
        return "{type} '{name}' variant '{variant}' is not known to Meltano. Variants: {variant_labels}".format(
            type=self.plugin.type.descriptor.capitalize(),
            name=self.plugin.name,
            variant=self.variant_name,
            variant_labels=self.plugin.variant_labels,
        )

class YAMLEnum(str, Enum):
    def __str__(self):
        return self.value

    @staticmethod
    def yaml_representer(dumper, obj):
        return dumper.represent_scalar("tag:yaml.org,2002:str", str(obj))

yaml.add_multi_representer(YAMLEnum, YAMLEnum.yaml_representer)

class PluginType(YAMLEnum):
    EXTRACTORS = "extractors"
    LOADERS = "loaders"
    TRANSFORMS = "transforms"
    MODELS = "models"
    DASHBOARDS = "dashboards"
    ORCHESTRATORS = "orchestrators"
    TRANSFORMERS = "transformers"
    FILES = "files"

    def __str__(self):
        return self.value

    @property
    def descriptor(self):
        if self is self.__class__.FILES:
            return "file bundle"

        return self.singular

    @property
    def singular(self):
        """Makes it singular for `meltano add PLUGIN_TYPE`"""
        return self.value[:-1]

    @property
    def verb(self):
        if self is self.__class__.TRANSFORMS:
            return self.singular

        return self.value[:-3]

    @classmethod
    def value_exists(cls, value):
        return value in cls._value2member_map_

    @classmethod
    def cli_arguments(cls):
        args = [type.singular for type in cls]
        args.extend([type for type in cls])
        return args

    @classmethod
    def from_cli_argument(cls, value):
        if not value.endswith("s"):
            value += "s"

        return cls(value)

class PluginRef(Canonical):
    def __init__(self, plugin_type: Union[str, PluginType], name: str, **kwargs):
        self._type = (
            plugin_type
            if isinstance(plugin_type, PluginType)
            else PluginType(plugin_type)
        )

        super().__init__(name=name, **kwargs)

    @property
    def type(self):
        return self._type

    def __eq__(self, other):
        return self.name == other.name and self.type == other.type

    def __hash__(self):
        return hash((self.type, self.name))

    def set_presentation_attrs(self, extras):
        self.update(
            hidden=extras.pop("hidden", None),
            label=extras.pop("label", None),
            logo_url=extras.pop("logo_url", None),
            description=extras.pop("description", None),
        )

class Variant(NameEq, Canonical):
    ORIGINAL_NAME = "original"
    DEFAULT_NAME = "default"

    def __init__(
        self,
        name: str = None,
        original: Optional[bool] = None,
        deprecated: Optional[bool] = None,
        docs: Optional[str] = None,
        repo: Optional[str] = None,
        pip_url: Optional[str] = None,
        executable: Optional[str] = None,
        capabilities: Optional[list] = [],
        settings_group_validation: Optional[list] = [],
        settings: Optional[list] = [],
        **extras,
    ):
        super().__init__(
            name=name,
            original=original,
            deprecated=deprecated,
            docs=docs,
            repo=repo,
            pip_url=pip_url,
            executable=executable,
            capabilities=list(capabilities),
            settings_group_validation=list(settings_group_validation),
            settings=list(map(SettingDefinition.parse, settings)),
            extras=extras,
        )

class PluginDefinition(PluginRef):
    def __init__(
        self,
        plugin_type: PluginType,
        name: str,
        namespace: str,
        variant: Optional[str] = None,
        variants: Optional[list] = [],
        **extras,
    ):
        super().__init__(plugin_type, name)

        self._defaults["label"] = lambda p: p.name

        def default_logo_url(p):
            short_name = re.sub(r"^(tap|target)-", "", p.name)
            return f"/static/logos/{short_name}-logo.png"

        self._defaults["logo_url"] = default_logo_url

        if not variants:
            variant = Variant(variant, **extras)

            # Any properties considered "extra" by the variant should be
            # considered extras of the plugin definition.
            extras = variant.extras
            variant.extras = {}

            variants = [variant]

        # Attributes will be listed in meltano.yml in this order:
        self.namespace = namespace
        self.set_presentation_attrs(extras)
        self.extras = extras
        self.variants = list(map(Variant.parse, variants))

    def __iter__(self):
        for k, v in super().__iter__():
            if k == "variants" and len(v) == 1:
                # If there is only a single variant, its properties can be
                # nested in the plugin definition
                for variant_k, variant_v in v[0]:
                    if variant_k == "name":
                        variant_k = "variant"

                    yield (variant_k, variant_v)
            else:
                yield (k, v)

    def get_variant(self, variant_name: str) -> Variant:
        try:
            return find_named(self.variants, variant_name)
        except NotFound as err:
            raise VariantNotFoundError(self, variant_name) from err

    def find_variant(self, variant_or_name: Union[str, Variant] = None):
        if isinstance(variant_or_name, Variant):
            return variant_or_name

        if variant_or_name is None or variant_or_name == Variant.DEFAULT_NAME:
            return self.variants[0]

        if variant_or_name == Variant.ORIGINAL_NAME:
            try:
                return next(v for v in self.variants if v.original)
            except StopIteration:
                return self.variants[0]

        return self.get_variant(variant_or_name)

    def variant_label(self, variant):
        """Return label for specified variant."""
        variant = self.find_variant(variant)

        label = variant.name or Variant.ORIGINAL_NAME
        if variant == self.variants[0]:
            label = f"{label} (default)"
        elif variant.deprecated:
            label = f"{label} (deprecated)"

        return label

    @property
    def variant_labels(self):
        """Return labels for supported variants."""
        return ", ".join([self.variant_label(variant) for variant in self.variants])

class BasePlugin(HookObject):
    EXTRA_SETTINGS = []

    def __init__(self, plugin_def: PluginDefinition, variant: Variant):
        super().__init__()

        self._plugin_def = plugin_def
        self._variant = variant

    def __eq__(self, other):
        return (
            self._plugin_def == other._plugin_def  # noqa: WPS437
            and self._variant == other._variant  # noqa: WPS437
        )

    def __hash__(self):
        return hash((self._plugin_def, self._variant))

    def __iter__(self):
        yield from self._plugin_def

    def __getattr__(self, attr):
        try:
            return getattr(self._plugin_def, attr)
        except AttributeError:
            return getattr(self._variant, attr)

    @property
    def variant(self):
        return self._variant.name

    @property
    def executable(self):
        return self._variant.executable or self._plugin_def.name

    @property
    def extras(self):
        return {**self._plugin_def.extras, **self._variant.extras}

    @property
    def extra_settings(self):
        defaults = {f"_{k}": v for k, v in self.extras.items()}

        existing_settings = []
        for setting in self.EXTRA_SETTINGS:
            default_value = defaults.get(setting.name)
            if default_value is not None:
                setting = setting.with_attrs(value=default_value)

            existing_settings.append(setting)

        # Create setting definitions for unknown defaults,
        # including flattened keys of default nested object items
        existing_settings.extend(
            SettingDefinition.from_missing(
                existing_settings, defaults, custom=False, default=True
            )
        )

        return existing_settings

    def env_prefixes(self, for_writing=False):
        """Return environment variable prefixes to use for settings."""
        return [self.name, self.namespace]

    def is_installable(self):
        return self.pip_url is not None

    def is_invokable(self):
        return self.is_installable()

    def is_configurable(self):
        return True

    def should_add_to_file(self):
        return True

    @property
    def runner(self):
        return None

    def exec_args(self, files: Dict):
        return []

    @property
    def config_files(self):
        """Return a list of stubbed files created for this plugin."""
        return dict()

    @property
    def output_files(self):
        return dict()

    def process_config(self, config):
        return config