from PyQt5 import QtCore, QtGui, QtWidgets
import os

directory = os.getcwd()

class Cryptor:
		text = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', ' ', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.', '/', ':', ';', '<', '=', '>', '?', '@', '[', '\\', ']', '^', '_', '`', '{', '|', '}', '~','\n']
		key = ["Contact developer for keys"]
		def enc(inpt):
			return ''.join([Cryptor.key[Cryptor.text.index(tol)] for tol in inpt if tol in Cryptor.text])
		def dec(inpt):
			de = [inpt[d:d+5] for d in range(0, len(inpt), 5)]
			return ''.join([Cryptor.text[Cryptor.key.index(dr)] for dr in de if dr in Cryptor.key])

class Ui_MainWindow(object):
    def setupUi(self, MainWindow):
        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(411, 408)
        font = QtGui.QFont()
        font.setFamily("Arial")
        font.setPointSize(9)
        MainWindow.setFont(font)
        icon = QtGui.QIcon()
        icon.addPixmap(QtGui.QPixmap(f"{directory}/logotk.ico"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
        MainWindow.setWindowIcon(icon)
        self.data = 'Encrypt'
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")
# =============================================================================================================================
        self.messagebox = QtWidgets.QTextEdit(self.centralwidget)
        self.messagebox.setGeometry(QtCore.QRect(10, 70, 391, 131))
        self.messagebox.setObjectName("messagebox")
# =============================================================================================================================
        self.outputbox = QtWidgets.QTextEdit(self.centralwidget)
        self.outputbox.setGeometry(QtCore.QRect(10, 240, 291, 131))
        self.outputbox.setObjectName("outputbox")
# =============================================================================================================================
        self.label = QtWidgets.QLabel(self.centralwidget)
        self.label.setGeometry(QtCore.QRect(10, 30, 261, 41))
        self.label.setObjectName("label")
# =============================================================================================================================
        self.label_2 = QtWidgets.QLabel(self.centralwidget)
        self.label_2.setGeometry(QtCore.QRect(70, 200, 201, 41))
        self.label_2.setObjectName("label_2")
# =============================================================================================================================
        self.Choose = QtWidgets.QComboBox(self.centralwidget)
        self.Choose.setGeometry(QtCore.QRect(280, 40, 121, 22))
        self.Choose.setObjectName("Choose")
        self.Choose.addItem('Encrypt')
        self.Choose.addItem('Decrypt')
        self.Choose.activated[str].connect(self.onPressed)
# =============================================================================================================================
        self.crypt = QtWidgets.QPushButton(self.centralwidget)  # Cryptor Button
        self.crypt.setGeometry(QtCore.QRect(310, 330, 91, 31))
        self.crypt.setObjectName("crypt")
        self.crypt.clicked.connect(self.onChanged)
# =============================================================================================================================
        self.label_3 = QtWidgets.QLabel(self.centralwidget)
        self.label_3.setGeometry(QtCore.QRect(80, 0, 241, 41))
        self.label_3.setAutoFillBackground(False)
        self.label_3.setObjectName("label_3")
# =============================================================================================================================
        self.widget = QtWidgets.QWidget(self.centralwidget)
        self.widget.setGeometry(QtCore.QRect(310, 240, 91, 75))
        self.widget.setObjectName("widget")
        self.verticalLayout = QtWidgets.QVBoxLayout(self.widget)
        self.verticalLayout.setContentsMargins(0, 0, 0, 0)
        self.verticalLayout.setObjectName("verticalLayout")
        self.label_4 = QtWidgets.QLabel(self.widget)
        self.label_4.setObjectName("label_4")
        self.verticalLayout.addWidget(self.label_4)
# =============================================================================================================================
        self.pushButton = QtWidgets.QPushButton(self.widget)   # Clear Button
        self.pushButton.setObjectName("pushButton")
        self.pushButton.clicked.connect(self.clear)
        self.verticalLayout.addWidget(self.pushButton)
        MainWindow.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 411, 21))
        self.menubar.setObjectName("menubar")
        MainWindow.setMenuBar(self.menubar)
        self.statusbar = QtWidgets.QStatusBar(MainWindow)
        self.statusbar.setObjectName("statusbar")
        MainWindow.setStatusBar(self.statusbar)

        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "Cryptor"))
        self.label.setText(_translate("MainWindow", "
Enter Your Message Below
"))
        self.label_2.setText(_translate("MainWindow", "
Generated Output
"))
        self.crypt.setText(_translate("MainWindow", "C R Y P T"))
        self.label_3.setText(_translate("MainWindow", "
GhosTmaN EnCryption
"))
        self.label_4.setText(_translate("MainWindow", "
Click here to

clear Screen
"))
        self.pushButton.setText(_translate("MainWindow", "Clear"))

    def onPressed(self,text):
        self.data = text

    def onChanged(self):
        text = self.messagebox.toPlainText()
        if self.data == "Encrypt":
            self.outputbox.setPlainText(Cryptor.enc(text))
        elif self.data == "Decrypt":
            self.outputbox.setPlainText(Cryptor.dec(text))
    def clear(self):
        self.messagebox.setPlainText('')
        self.outputbox.setPlainText('')

if __name__ == "__main__":
    import sys
    app = QtWidgets.QApplication(sys.argv)
    MainWindow = QtWidgets.QMainWindow()
    ui = Ui_MainWindow()
    ui.setupUi(MainWindow)
    MainWindow.show()
    sys.exit(app.exec_())


from elasticsearch import AsyncElasticsearch

class DocumentsIndex:
    """
    The documents index.
    """

    def __init__(self):
        self.name = 'documents'
        self.mapping = {
            'properties': {
                'iD': {'type': 'integer', 'coerce': False},
                'text': {'type': 'text'}
            }
        }

    async def create(self, es: AsyncElasticsearch):
        """
        Creates a documents index.
        """
        return await es.indices.create(index=self.name, mappings=self.mapping)


#10815 숫자 카드
"""
문제 설명
숫자 카드는 정수 하나가 적혀져 있는 카드이고, N개를 가지고 있다.
정수 M개가 주어졌을 때, 이 수가 적혀있는 숫자 카드를 갖고 있는지 아닌지
구하는 프로그램
입력
첫째 줄에는 갖고 있는 숫자 카드의 개수 N개가 주어짐.(1<=N<=500,000)
둘째 줄에는 숫자 카드에 적혀있는 정수가 주어진다.(-10,000,000<=수<=10,000,000)
*두 숫자 카드에 같은 수가 적혀있는 경우는 없다.
셋째 줄에는 M이 주어진다. (1<=M<=500,000)
넷째 줄에는 갖고 있는 숫자 카드인지 아닌지를 구해야 할 M개의 정수가 주어지며,
이 수는 공백으로 구분되어져 있다. (-10,000,000<=수<=10,000,000)
출력
첫째 줄에 입력으로 주어진 M개의 수에 대해서,
각 수가 적힌 숫자 카드를 갖고 있으면 1, 아니면 0을 공백으로 구분해 출력
"""
import sys

N = int(input())
num = list(map(int,sys.stdin.readline().split()))
M = int(input())
check = list(map(int,sys.stdin.readline().split()))

num.sort()      #이진탐색 하기 위해 정렬

def bin_search(a,key,left,right):
    while start<=end:
        mid = (left+right)//2
        if a[mid] == key:
            return mid
        elif a[mid]>key:
            right = mid -1
        else:
            left = mid +1
    return None

for i in range(M):
    if bin_search(num,check[i],0,N-1) is not None:
        print(1,end=' ')
    else:
        print(0,end=' ')

#시간복잡도 줄이기 -> set(집합함수) 자료구조 hash 사용

import sys

N = int(input())
card = set(map(int,sys.stdin.readline().split()))
M = int(input())
check = list(map(int,sys.stdin.readline().split()))

for i in check:
    if i in card:
        print(1,end=' ')
    else:
        print(0,end=' ')


import itertools
def permutation(s):
    if(not s):
        return None
    if(len(s) == 1):
        return list(s)
    l = list(s)
    l.sort()
    result = []
    for i in range(len(l)):
        if(i > 0 and l[i] == l[i-1]):
            continue
        part2 = permutation(''.join(l[:i]) +''.join(l[i+1:]))
        for j in part2:
            result.append(l[i]+j)
    return result

def permutation2(s):
    result = [''.join(x) for x in itertools.permutations(s, len(s))]
    return list(result)
if __name__ == '__main__':
    s = 'abc'
    result = permutation(s)
    for i in result:
        print(i)
    result = permutation2(s)
    for i in result:
        print(i)


import os
import yara
import uuid

from loguru import logger

from vigil.schema import YaraMatch
from vigil.schema import ScanModel
from vigil.schema import BaseScanner

from vigil.registry import Registration

@Registration.scanner(name='yara', requires_config=True)
class YaraScanner(BaseScanner):
    def __init__(self, rules_dir: str):
        self.name = 'scanner:yara'
        self.rules_dir = rules_dir
        self.compiled_rules = None

        if not os.path.exists(self.rules_dir):
            logger.error(f'Directory not found: {self.rules_dir}')
            raise Exception

        if not os.path.isdir(self.rules_dir):
            logger.error(f'Path is not a valid directory: {self.rules_dir}')
            raise Exception

        self.load_rules()
        logger.success('Loaded scanner')

    def load_rules(self) -> bool:
        """Compile all YARA rules in a directory and store in memory"""
        logger.info(f'Loading rules from directory: {self.rules_dir}')
        rules = os.listdir(self.rules_dir)

        if len(rules) == 0:
            return

        yara_paths = {}
        for _file in rules:
            if self.is_yara_file(_file):
                yara_paths[_file] = os.path.join(self.rules_dir, _file)

        try:
            self.compiled_rules = yara.compile(filepaths=yara_paths)
        except Exception as err:
            logger.error(f'YARA compilation error: {err}')
            raise err

    def is_yara_file(self, file_path: str) -> bool:
        """Check if file is rule by extension"""
        if file_path.lower().endswith('.yara') or file_path.lower().endswith('.yar'):
            return True
        return False

    def analyze(self, scan_obj: ScanModel, scan_id: uuid.uuid4) -> ScanModel:
        """Run scan against input data and return list of YaraMatchs"""
        logger.info(f'Performing scan; id="{scan_id}"')

        if scan_obj.prompt.strip() == '':
            logger.error(f'No input data; id="{scan_id}"')
            return scan_obj

        try:
            matches = self.compiled_rules.match(data=scan_obj.prompt)
        except Exception as err:
            logger.error(f'Failed to perform yara scan; id="{scan_id}" error="{err}"')
            return scan_obj

        for match in matches:
            m = YaraMatch(rule_name=match.rule, tags=match.tags, category=match.meta.get('category', None))
            logger.warning(f'Matched rule rule="{m.rule_name} tags="{m.tags}" category="{m.category}"')
            scan_obj.results.append(m)

        if len(scan_obj.results) == 0:
            logger.info(f'No matches found; id="{scan_id}"')

        return scan_obj


# -*- coding: utf-8 -*-
"""
Created on Wed Jun 23 19:47:03 2021

@author: DELL
"""

import numpy as np
# import matplotlib.pyplot as plt
# from mpl_toolkits.mplot3d import Axes3D
from my_spectrum_analyzer import img_spectrum_analyzer
import torch
from torchvision import transforms #datasets, models,
from tqdm import tqdm
import os
import sys
import torch.nn as nn
import time
from models.cifar.allconv import AllConvNet
from models.resnet import resnet50
from models.vgg import vgg16_bn
from models.cifar.allconv import AllConvNet
import torchvision.models as models
from third_party.ResNeXt_DenseNet.models.densenet import densenet
from third_party.ResNeXt_DenseNet.models.resnext import resnext29
from third_party.WideResNet_pytorch.wideresnet import WideResNet
from art.attacks.evasion import FastGradientMethod,DeepFool
from art.attacks.evasion import CarliniL2Method,CarliniLInfMethod
from art.attacks.evasion import ProjectedGradientDescent
# from art.attacks.evasion import UniversalPerturbation
from art.estimators.classification import PyTorchClassifier
import json
sys.path.append("..")
# from train_code.my_img_transformer import img_transformer
# from mpl_toolkits.mplot3d import Axes3D
# from sklearn.manifold import TSNE
sys.path.append('../common_code')
from load_cifar_data import load_CIFAR_batch,load_CIFAR_train

class img_spectrum_labeler:

    # 解释器初始化
    def __init__(self,dataset):
        if 'imagenet'==dataset:
            self.img_size=224
            mean_now=[0.485, 0.456, 0.406]
            std_now=[0.229, 0.224, 0.225]
            self.num_classes=1000

        elif 'cifar-10'==dataset:
            self.img_size=32
            mean_now=[0.5] * 3
            std_now=[0.5] * 3
            self.num_classes=10
        else:
            print('ERROR DATASET')

        self.trans=transforms.Compose([transforms.Normalize(mean=mean_now, std=std_now)])
        self.s_analyzer=img_spectrum_analyzer(self.img_size).batch_get_spectrum_feature#batch_get_spectrum_energy

    def select_attack(self, fmodel, attack_idx, eps_idx):
        attack_names=['FGSM_L2_IDP','PGD_L2_IDP','CW_L2_IDP','Deepfool_L2_IDP','FGSM_Linf_IDP','PGD_Linf_IDP','CW_Linf_IDP']
        eps_L2=[0.1,1.0,10.0,100.0]
        eps_Linf=[0.01,0.1,1.0,10.0]

        att_method=attack_names[attack_idx]
        if 'L2' in att_method:
            eps=float(eps_L2[eps_idx%len(eps_L2)])
        else:
            eps=float(eps_Linf[eps_idx%len(eps_Linf)])

        if att_method   == 'FGSM_L2_IDP':
            attack = FastGradientMethod(estimator=fmodel,eps=eps,norm=2)
        elif att_method == 'PGD_L2_IDP':
            attack = ProjectedGradientDescent(estimator=fmodel,eps=eps,norm=2,batch_size=512,verbose=False)
        elif att_method == 'CW_L2_IDP':
            attack = CarliniL2Method(classifier=fmodel,batch_size=512,verbose=False)
        elif att_method == 'Deepfool_L2_IDP':
            attack = DeepFool(classifier=fmodel,batch_size=512,verbose=False)

        elif att_method == 'FGSM_Linf_IDP':
            attack = FastGradientMethod(estimator=fmodel,eps=eps,norm=np.inf)
        elif att_method == 'PGD_Linf_IDP':
            attack = ProjectedGradientDescent(estimator=fmodel,eps=eps,norm=np.inf,batch_size=512,verbose=False)
        elif att_method == 'CW_Linf_IDP':
            attack = CarliniLInfMethod(classifier=fmodel,eps=eps,batch_size=512,verbose=False)

        else:
            raise Exception('Wrong Attack Mode: {} !!!'.format(att_method))
        return attack, eps

    def get_energy_label(self, model, imgs_in, labels_in, is_adv):
        assert imgs_in.shape[-2]==imgs_in.shape[-1]
        spectrum  = self.s_analyzer(imgs_in)
        labels_out = is_adv * np.ones(labels_in.shape)
        return spectrum,labels_out.reshape((-1,1))

if __name__=='__main__':
    '''
    settings
    '''
    # 配置解释器参数
    if len(sys.argv)!=4:
        print('Manual Mode !!!')
        model_type    = 'allconv'
        data          = 'train'
        device        = 2
    else:
        print('Terminal Mode !!!')
        model_type  = sys.argv[1]
        data        = sys.argv[2]
        device      = int(sys.argv[3])

    '''
    加载cifar-10图像
    '''
    os.environ['CUDA_VISIBLE_DEVICES']=str(device)
    dir_cifar     = '../../../../../media/ubuntu204/F/Dataset/Dataset_tar/cifar-10-batches-py'
    if 'test'==data:
        images,labels = load_CIFAR_batch(os.path.join(dir_cifar,'test_batch'))
    elif 'train'==data:
        images,labels = load_CIFAR_train(dir_cifar)
    else:
        print('Wrong data mode !!!')

    '''
    加载模型
    '''
    dataset='cifar-10'
    dir_model = '../models/cifar_vanilla_'+model_type+'.pth.tar'
    if model_type == 'resnet50_imagenet':
        model = models.resnet50(pretrained=True).eval()
        model = torch.nn.DataParallel(model).cuda()
        dataset = 'imagenet'
    elif model_type == 'vgg16_imagenet':
        model = models.vgg16(pretrained=True).eval()
        model = torch.nn.DataParallel(model).cuda()
        dataset = 'imagenet'
    elif model_type == 'alexnet_imagenet':
        model = models.alexnet(pretrained=True).eval()
        model = torch.nn.DataParallel(model).cuda()
        dataset = 'imagenet'
    elif model_type == 'resnet50':
        model = resnet50().eval()
        model = torch.nn.DataParallel(model).cuda()
        checkpoint = torch.load(dir_model)
        model.load_state_dict(checkpoint["state_dict"],True)
    elif model_type == 'vgg16':
        model = vgg16_bn().eval()
        model = torch.nn.DataParallel(model).cuda()
        checkpoint = torch.load(dir_model)
        model.load_state_dict(checkpoint["state_dict"],True)
    elif model_type == 'allconv':
        model = AllConvNet(10).eval()
        model = torch.nn.DataParallel(model).cuda()
        checkpoint = torch.load(dir_model)
        model.load_state_dict(checkpoint["state_dict"],True)
    else:
        raise Exception('Wrong model name: {} !!!'.format(model_type))
    model.eval()

    if 'cifar-10'==dataset:
        mean   = np.array((0.5,0.5,0.5),dtype=np.float32)
        std    = np.array((0.5,0.5,0.5),dtype=np.float32)
        nb_classes = 10
        input_shape=(3,32,32)
        rober_np= np.array([[0,1],[0,2],[0,3],
                        [1,1],[1,2],[1,3],
                        [8,1],[8,2],[8,3],
                        [9,1],[9,2],[9,3],
                        [4,1],[4,2],[4,3],
                        [5,1],[5,2],[5,3],
                        [12,1],[12,2],[12,3],
                        [13,1],[13,2],[13,3],
                       ]) # 被标为rober的设置
    elif 'imagenet'==dataset:
        mean   = np.array((0.485, 0.456, 0.406),dtype=np.float32)
        std    = np.array((0.229, 0.224, 0.225),dtype=np.float32)
        nb_classes = 1000
        input_shape=(3,224,224)
        rober_np= np.array([[0,1],[0,2],[0,3],
                        [1,1],[1,2],[1,3],
                        [8,1],[8,2],[8,3],
                        [9,1],[9,2],[9,3],
                        [4,1],[4,2],[4,3],
                        [5,1],[5,2],[5,3],
                        [12,1],[12,2],[12,3],
                        [13,1],[13,2],[13,3],
                       ]) # 被标为rober的设置
    else:
        raise Exception('Wrong dataset type: {} !!!'.format(dataset))

    fmodel = PyTorchClassifier(model = model,nb_classes=nb_classes,clip_values=(0,1),
                               input_shape=input_shape,loss = nn.CrossEntropyLoss(),
                               preprocessing=(mean, std))

    '''
    读取数据
    '''
    labeler = img_spectrum_labeler(dataset)
    # fft_transformer = img_transformer(8,0,6)

    batch           = 10
    batch_num       = int(len(labels)/batch)
    spectrums_list  = []
    labels_list     = []
    start_time = time.time()
    for i in tqdm(range(batch_num)):
        '''
        攻击与防护
        '''
        images_batch = images[batch*i:batch*(i+1),...].transpose(0,3,1,2)
        labels_batch = labels[batch*i:batch*(i+1),...]

        # 标为对抗样本
        attack_name=np.random.randint(1)
        attack_eps=np.random.randint(4)

        attack,eps=labeler.select_attack(fmodel,attack_name, attack_eps)
        images_adv = attack.generate(x=images_batch,y=labels_batch)

        # 根据列表标为有害或无害样本
        flag_rober=0
        for m in range(rober_np.shape[0]):
            if (attack_name==rober_np[m,0]) and (attack_eps==rober_np[m,1]):
                flag_rober=1
                break

        spectrums_save,labels_save=labeler.get_energy_label(model, images_adv, labels_batch, flag_rober)
        spectrums_list.append(spectrums_save)
        labels_list.append(labels_save)

    sub_dir='spectrum_label/'+model_type
    saved_dir_path  = '../saved_tests/img_attack/'+sub_dir
    if not os.path.exists(saved_dir_path):
        os.makedirs(saved_dir_path)
    spectrums_np=np.vstack(spectrums_list)
    labels_np=np.vstack(labels_list)

    np.save(os.path.join(saved_dir_path,'spectrums_'+data+'.npy'), spectrums_np)
    np.save(os.path.join(saved_dir_path,'labels_'+data+'.npy'), labels_np)
    end_time=time.time()
    print(("Time %f s")%(end_time-start_time))


'''
Into to deep learning Koopman operators
===
Author: Nicholas Geneva (MIT Liscense)
url: https://nicholasgeneva.com/blog/
github: https://github.com/NickGeneva/blog-code
===
'''
import torch
import numpy as np
import os

import matplotlib as mpl
# mpl.use('agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib.legend_handler import HandlerBase

class HandlerColormap(HandlerBase):
    def __init__(self, cmap, num_stripes=8, **kw):
        HandlerBase.__init__(self, **kw)
        self.cmap = cmap
        self.num_stripes = num_stripes
    def create_artists(self, legend, orig_handle,
                       xdescent, ydescent, width, height, fontsize, trans):
        stripes = []
        for i in range(self.num_stripes):
            s = Rectangle([xdescent + i * width / self.num_stripes, ydescent],
                          width / self.num_stripes,
                          height,
                          fc=self.cmap((2 * i + 1) / (2 * self.num_stripes)),
                          transform=trans)
            stripes.append(s)
        return stripes

def plotDuffingPrediction(args, yPred, yTar, bidx=0, epoch=0):
    '''
    Plots the prediction of the model for the duffing equation in the space
    of the two state variables. A color map is used to denote time
    progression.
    Args:
        args (argparse): object with programs arguements
        yPred (torch.Tensor): [mb x t] tensor of a batch of model predictions
        yTar (torch.Tensor): [mb x t] tensor of a batch of targets
        bidx (int): integer specifying batch case to plot
    '''
    yPred = yPred.cpu()
    yTar = yTar.cpu()
    plt.close('all')

    # Use color map to indicate time
    cmaps = [plt.get_cmap("spring"), plt.get_cmap("summer")]
    predColors = cmaps[0](np.linspace(0,1,yPred.size(1)))
    tarColors = cmaps[1](np.linspace(0,1,yPred.size(1)))

    for i in range(0,yPred.size(1)-1):
        plt.plot(yPred[0,i:i+2,0], yPred[0,i:i+2,1], '--*', color=predColors[i])
        plt.plot(yTar[0,i:i+2,0], yTar[0,i:i+2,1], color=tarColors[i])

    plt.xlim([-2,2])
    plt.ylim([-2,2])
    cmap_handles = [Rectangle((0, 0), 1, 1) for _ in cmaps]
    handler_map = dict(zip(cmap_handles,
                        [HandlerColormap(cm, num_stripes=8) for cm in cmaps]))

    # Create custom legend with color map rectangels
    plt.legend(handles=cmap_handles, labels=['Prediction','Target'], handler_map=handler_map, loc='upper right', framealpha=0.95)

    file_dir = args.pred_dir
    # If director does not exist create it
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)
    file_name = file_dir+"/{:s}Pred{:d}_{:d}".format(args.exp_name, bidx, epoch)
    plt.savefig(file_name+".png", bbox_inches='tight')

def plotEignValues(args, model, epoch=0):
    '''
    Plots the eigen values of the learned Koopman operator
    Args:
        args (argparse): object with programs arguements
        model (pytorch module): pytorch model with koopman operator kMatrix
    '''
    # Get koopman operator from model
    kMatrix = model.getKoopmanMatrix().detach().cpu().numpy()

    try:
        w, v = np.linalg.eig(kMatrix)
    except:
        print('issue computing eigs')
        return

    plt.close('all')
    plt.scatter(np.real(np.log(np.abs(w))), np.imag(w))

    file_dir = args.pred_dir+'/eig'
    # If director does not exist create it
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)
    file_name = file_dir+"/{:s}Eig_{:d}".format(args.exp_name, epoch)
    plt.savefig(file_name+".png", bbox_inches='tight')

def plotEignVectors(args, model, n=3, epoch=0):
    '''
    Plots the eigen vectors of the learned Koopman operator
    Args:
        args (argparse): object with programs arguements
        model (pytorch module): pytorch model with koopman operator kMatrix
    '''
    # Get koopman operator from model
    kMatrix = model.kMatrix.detach().cpu().numpy()
    w, v = np.linalg.eig(kMatrix)
    idx = w.argsort()[::-1]
    w = w[idx]
    v = v[:,idx]

    data = np.load('./duffing-test.npy')
    yPred = torch.Tensor(data[:,:,:-1])

    plt.close("all")
    fig, ax = plt.subplots(1, n, figsize=(4*n, 4))

    for i in range(n):
        ev_data = np.zeros(yPred.size(0))
        x_data = np.zeros(yPred.size(0))
        xdot_data = np.zeros(yPred.size(0))
        for j in range(yPred.size(0)):

            xin = yPred[j,:args.stride].to(args.device)
            xin = xin.view(-1).unsqueeze(0)

            g0,_= model(xin)
            g0 = g0.squeeze(0).detach().cpu().numpy()

            ev_data[j] = np.real(np.dot(g0,v[i]))
            x_data[j] = xin[0,-2].cpu().numpy()
            xdot_data[j] = xin[0,-1].cpu().numpy()

        ax[i].set_title('{:.02f} + {:.02f}i'.format(w[i].real, w[i].imag))
        ax[i].tricontourf(x_data, xdot_data, ev_data, cmap="RdBu_r")

    file_dir = args.pred_dir+'/eigVects'
    # If director does not exist create it
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)
    file_name = file_dir+"/{:s}EigVects_{:d}".format(args.exp_name, epoch)
    plt.savefig(file_name+".png", bbox_inches='tight')


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

"""
@author: chenhao
@time: 2023/7/13 15:17
"""
import math
import random
import re

from icetk import icetk
from snippets import *
from tqdm import tqdm

def try_load(p):
    try:
        ds = jload_lines(p)
        return ds
    except Exception:
        ds = jload(p)
        return ds

def is_valid(item, show=False, no_gpt=True):
    input = item.get("input")
    target = item.get("target")
    if not input or not target:
        return False

    input_list = icetk.encode(item["input"])
    target_list = icetk.encode(item["target"])
    rs = True
    if not 4 <= len(input_list) <= 512:
        rs = False
    if not 8 <= len(target_list) <= 2048:
        rs = False
    input_unk_num = len([x for x in input_list if x == 20000])
    target_unk_num = len([x for x in target_list if x == 20000])
    if input_unk_num >= 2 or target_unk_num >= 2:
        rs = False
    if no_gpt:
        pt = "gpt"
        if re.search(pt, item["input"], re.I) or re.search(pt, item["target"], re.I):
            rs = False

    if not rs and show:
        print(f"item: {item} is not valid")
    return rs

def parse_alpaca_zh(item):
    rs = dict(input=item["instruction"].strip(), target=item["output"].strip(), source="alpaca_gpt4_data_zh")
    return rs

def parse_tigerbot(item):
    rs = dict(input=(item["instruction"] + "\n" + item["input"]).strip(), target=item["output"].strip(),
              source="tigerbot")
    return rs

def parse_belle(item):
    rs = dict(input=(item["instruction"] + "\n" + item["input"]).strip(), target=item["output"].strip(),
              source="belle")
    return rs

def is_literature_generation(item):
    pt = "故事|诗|小说|文案|菜谱|报告|总结|剧本|社论|文章|道歉信|感激信|消化|发言稿|作文|菜谱|对联|教程|评论|短文"
    return re.search(pt, item["input"]) is not None

def extract_word_num(text, pattern="至少(\d+)字|不少于(\d+)字"):
    for e in re.findall(pattern, text):
        for m in e:
            try:
                # print(m)
                rs = int(m.strip())
                return rs
            except Exception as e:
                pass
    return None

def match_word_limit(item):

    tgt_len = len(item["target"])
    ipt = item["input"]
    min_num = extract_word_num(ipt, "[至最]少(\d+)字|不少于(\d+)字")
    if min_num and tgt_len >= min_num:
        return True
    max_num = extract_word_num(ipt, "[至最]多(\d+)字|不多于(\d+)字")
    if max_num and tgt_len <= max_num:
        return True
    around_num = extract_word_num(ipt, "大约(\d+)字|(\d+)字左右")
    if around_num and tgt_len * 0.85 <= around_num <= tgt_len * 1.15:
        return True
    return False

# def parse_share_gpt(item):
#     rs = []
#     pre = ""
#     rd = 1
#     for idx, ele in enumerate(item["conversations"]):
#         if ele["from"] == "human":
#             pre = pre + f"##第 {rd} 轮##\n\n问:{ele['value']}\n\n答:"
#         else:
#             rs.append(dict(input=pre.strip(), target=ele["value"], source="share-gpt"))
#             pre = pre + ele["value"] + "\n\n"
#             rd += 1
#     return rs

def is_chinese(strs, pct=0.4):
    ch_num = 0

    for _char in strs:
        if '\u4e00' <= _char <= '\u9fa5':
            ch_num += 1
    return ch_num / len(strs) >= pct

def parse_guanaco_non_chat(item):
    source = "guanaco_non_chat"
    if not item["output"].strip():
        return None

    if is_chinese(item["output"]):
        source += "_cn"

    rs = dict(input=(item["instruction"] + "\n" + item["input"]).strip(), target=item["output"].strip(),
              source=source)
    return rs

def parse_guanaco_chat(item):
    ipt = (item["instruction"] + item["input"]).strip()
    sentences = re.split(r"User:|Assistant:|System:", ipt)
    sentences = [e.strip() for e in sentences if e.strip()]
    # print(sentences)
    # sentences
    sentences += [item["output"].strip()]
    if not len(sentences) % 2 == 0:
        print(f"error in {len(sentences)} sentences")
        for s in sentences:
            print(s)
        return []

    pre = ""
    rs = []
    for idx in range(0, len(sentences), 2):
        rd = idx // 2 + 1
        pmt = pre + f"##第 {rd} 轮##\n\n问:{sentences[idx]}\n\n答:"
        if idx + 1 >= len(sentences):
            break
        ans = sentences[idx + 1].strip()
        if not ans:
            break
        source = "guanaco_chat"
        if is_chinese(ans):
            source += "_cn"
        rs.append(dict(input=pmt.strip(), target=ans, source=source))
        pre = pmt + ans + "\n\n"
    return rs

def _get_sentences_multiturn_chat(instruction):
    cur_sentence = ""
    rs = []
    tag = None
    sentences = re.split("(Assistant:)|(Human:)",instruction, flags=re.I)

    for s in sentences:
        if not s:
            continue
        if s in ["Assistant:", "Human:"]:
            if not tag or tag != s:
                tag = s
                if cur_sentence:
                    rs.append(cur_sentence.strip())
                    cur_sentence = ""
        else:
            cur_sentence += s

    if cur_sentence:
        rs.append(cur_sentence.strip())
    return rs

def parse_belle_multiturn_chat(item):
    ipt = (item["instruction"]).strip()
    sentences = _get_sentences_multiturn_chat(ipt)
    sentences += [item["output"].strip()]
    if not len(sentences) % 2 == 0:
        print(f"error in {len(sentences)} sentences")
        for idx, s in enumerate(sentences):
            print((idx,s))
        return []

    pre = ""
    rd = 1
    rs = []

    for idx, ele in enumerate(sentences):
        if idx %2 == 0:
            pre = pre + f"##第 {rd} 轮##\n\n问:{ele}\n\n答:"
        else:
            rs.append(dict(input=pre.strip(), target=ele, source="multiturn_chat"))
            rd += 1
            pre = pre + ele+ "\n\n"
    return rs

def parse_belle_chat(item):
    ipt = (item["instruction"] + item["input"]).strip()
    sentences = re.split(r"Human:|Assistant:", ipt)
    sentences = [e.strip() for e in sentences if e.strip()]
    # print(sentences)
    # sentences
    sentences += [item["output"].strip()]
    if not len(sentences) % 2 == 0:
        # print(f"error in {len(sentences)} sentences")
        # for idx, s in enumerate(sentences):
        #     print((idx,s))
        return []

    pre = ""
    rs = []
    for idx in range(0, len(sentences), 2):
        rd = idx // 2 + 1
        pmt = pre + f"##第 {rd} 轮##\n\n问:{sentences[idx]}\n\n答:"
        if idx + 1 >= len(sentences):
            break
        ans = sentences[idx + 1].strip()
        if not ans:
            break
        source = "belle_chat"
        rs.append(dict(input=pmt.strip(), target=ans, source=source))
        pre = pmt + ans + "\n\n"
    return rs

def contain_gpt(t):
    t = t.upper()
    return "GPT" in t or "OPENAI" in t

def parse_tsinghua_human_label(item):

    rs = []
    source = "tsinghua_human_label"

    history = item["history"]+[dict(prompt=item["prompt"], response=item["response"])]
    # print(history)
    # print(len(history))
    if len(history) > 1:
        source += "_multichat"
        pre = ""
        for idx, ele in enumerate(item["history"]):
            pre += f"##第 {idx+1} 轮##\n\n问:{ele['prompt']}\n\n答:"
            ans = ele["response"]
            rs.append(dict(input=pre, target=ans, source= source))
            pre += ans + "\n\n"
    else:
        source += "_qa"
        rs = [dict(input=history[0]["prompt"], target=history[0]["response"], source=source)]
    return rs

def parse_share_gpt(item):
    rs = []
    cov = item["conversations"]
    if len(cov) == 2:
        rs = [dict(input=cov[0]["value"], target=cov[1]["value"], source="share-gpt")]
    else:
        pre = ""
        sentences = [e["value"] for e in cov]
        for idx in range(0, len(sentences), 2):
            rd = idx // 2 + 1
            pmt = pre + f"##第 {rd} 轮##\n\n问:{sentences[idx]}\n\n答:"
            if idx + 1 >= len(sentences):
                break
            ans = sentences[idx + 1].strip()
            if not ans:
                break
            source = "share-gpt_chat"
            rs.append(dict(input=pmt.strip(), target=ans, source=source))
            pre = pmt + ans + "\n\n"
    return rs

def process_moss_qa(item):
    # item
    sentences = re.split(r"<eoh>|<eoa>", item["plain_text"])
    sentences = [e[e.index(":") + 1:].strip() for e in sentences if ":" in e]
    # sentences
    assert len(sentences) % 2 == 1
    pre = ""
    rs = []
    for idx in range(0, len(sentences), 2):
        rd = idx // 2 + 1
        pmt = pre + f"##第 {rd} 轮##\n\n问:{sentences[idx]}\n\n答:"
        ans = sentences[idx + 1]
        rs.append(dict(input=pmt, target=ans, source="moss_qa"))
        pre = pmt + ans + "\n\n"

    return rs

def process_moss_multi_chat(item, do_multi_chat=True):
    # item
    chats = item["chat"]
    # chats
    rs = []
    pre =""

    for idx, (t,ele) in enumerate(chats.items()):
        # ele
        human = ele["Human"].replace("<|Human|>: ","").replace("<eoh>","").strip()
        moss = ele["MOSS"].replace("<|MOSS|>: ","").replace("<eom>","").strip()
        human, moss
        if not do_multi_chat:
            rs.append(dict(input=human, target=moss, source="moss_qa"))
        else:
            q = \
    f'''##第 {idx+1} 轮##
    问:{human}

    答:
    '''
            a = moss
            pre = pre + "\n\n" + q
            rs.append(dict(input=pre.strip(), target=a.strip(), source="moss_multi_chat"))
            pre += a
    return rs

def parse_data(data, func):
    rs = []
    for idx, item in enumerate(tqdm(data[:])):
        try:
            tmp = func(item)
            if tmp:
                rs.append(tmp)
            else:
                pass
        except:
            pass
    print(f"origin rs len:{len(rs)}")
    if isinstance(rs[0], list):
        rs = flat(rs)
    print(f"parsed rs len:{len(rs)}")
    print(f"sample rs item :{random.choice(rs)}")
    return rs

def view_data(data, sample_num=10):
    sample = random.sample(data, min(len(data), sample_num))
    for item in sample:
        print(item["input"])
        print("*" * 30)
        print(item["target"])
        print("*" * 30)
        print(item["source"])
        print("*" * 30)
        print("\n\n")

def equal_sample(sources, tgt_num):
    rs = []
    t_num = int(math.ceil(tgt_num / len(sources)))
    for s in sources:
        # print(len(s))
        # print(t_num)
        print(f"sample {t_num} from {len(s)}")
        assert len(s) > t_num
        rs.extend(random.sample(s, t_num))
    return rs[:tgt_num]

if __name__ == "__main__":
    print(match_word_limit(dict(input="写一篇609999字左右的作文")))


#! /usr/bin/env python2
import sys
import binascii
import struct
import getopt
import re

def print_usage():
    sys.stderr.write('Usage:\n')
    sys.stderr.write('\tconvertleds.py [-d delay] input.led [output.l3d]\n')

options, remainder = getopt.getopt(sys.argv[1:], 'd:o:h', ['delay=', 'outdir=', 'help'])

delay_arg = None
delay_file = None
default_delay = 50
delay_regex = "^delay:(.+)\n"

outdir = "."

for opt, arg in options:
    if opt in ('-d', '--delay'):
        delay_arg = int(arg)
    if opt in ('-o', '--outdir'):
        outdir = arg
    if opt in ('-h', '--help'):
        print_usage()
        sys.exit(0)

if len(remainder) == 0:
    sys.stderr.write('missing input file\n')
    print_usage()
    sys.exit(1)

elif len(remainder) > 2:
    sys.stderr.write('too many arguments\n')
    print_usage()
    sys.exit(1)

filename = remainder[0]

if len(remainder) == 1:
    outfilename = re.sub('\.[^.]*$','.l3d',filename)
    outfilename = outdir + "/" + re.sub('.*/','',outfilename)

else:
    outfilename = remainder[1]

with open(filename) as fp:
    contents = fp.read()

    delay_match = re.search(delay_regex, contents)
    if delay_match:
        delay_file = delay_match.group(1)
        contents = re.sub(delay_regex, '', contents)

    # first use argument if available, then use delay from file if available, then fallback to default
    delay = delay_arg if delay_arg else int(delay_file) if delay_file else default_delay

    lines = len(filter(lambda x: x.strip(), contents.splitlines()))

    if lines > 50:
        sys.stderr.write('ERROR: failed to convert %s\n' % filename)
        sys.stderr.write('currently only animations with a maximum of 50 frames are supported!\n')
        sys.exit(2)

    elif lines == 1:
        delay = 65535 # for static files

    contents = contents.replace(chr(10), '')
    contents = contents.replace(' ', '')
    with open(outfilename, 'w') as fpW:
        fpW.write(struct.pack('>H', delay))
        fpW.write(binascii.unhexlify(contents))


from keras import backend as keras_back
from keras.layers import Input
import tensorflow as tf
import keras

import efficientnet.keras as efn
from keras.applications.resnet50 import ResNet50
# Todo - Testar a resnet 50 V2
from tensorflow.keras.applications.resnet_v2 import ResNet50V2
from tensorflow.keras.applications.resnet_v2 import ResNet101V2
from tensorflow.keras.applications.resnet_v2 import ResNet152V2
from keras.applications import InceptionV3, InceptionResNetV2, Xception
from keras.layers import (Layer,
                          Input, Dense, Dropout, Lambda, BatchNormalization,
                          Flatten, Add, Concatenate,
                          Activation, LeakyReLU,
                          Conv1D, Conv2D, Conv3D,
                          Conv2DTranspose, Conv3DTranspose,
                          SeparableConv1D, SeparableConv2D,
                          MaxPooling1D, MaxPooling2D, MaxPooling3D,
                          GlobalAveragePooling1D, GlobalAveragePooling2D,
                          GlobalAveragePooling3D)

from utils import utils
from icecream import ic

def get_model_effnet(x_data, weights, effnet_version):
    keras_back.set_image_data_format('channels_last')
    img_shape = (x_data.shape[1], x_data.shape[2], x_data.shape[3])
    ic("Input Shape Matrix: ", img_shape)
    # ic("X_data Shape: 1- {}, 2- {}, 3- {}".format(x_data.shape[1], x_data.shape[2], x_data.shape[3]))
    img_input = Input(shape=img_shape)

    ic('\n ** Utilizando a Rede EfficientNet ', effnet_version)
    ic('\n ** Pesos carregados: ', weights)

    if effnet_version == 'B0':
        effnet = efn.EfficientNetB0(include_top=False, input_tensor=img_input, weights=weights, pooling=None,
                                    input_shape=img_shape)

    elif effnet_version == 'B1':
        effnet = efn.EfficientNetB1(include_top=False, input_tensor=img_input, weights=weights, pooling=None,
                                    input_shape=img_shape)

    elif effnet_version == 'B2':
        effnet = efn.EfficientNetB2(include_top=False, input_tensor=img_input, weights=weights, pooling=None,
                                    input_shape=img_shape)

    elif effnet_version == 'B3':
        effnet = efn.EfficientNetB3(include_top=False, input_tensor=img_input, weights=weights, pooling=None,
                                    input_shape=img_shape)

    elif effnet_version == 'B4':
        effnet = efn.EfficientNetB4(include_top=False, input_tensor=img_input, weights=weights, pooling=None,
                                    input_shape=img_shape)

    elif effnet_version == 'B5':
        effnet = efn.EfficientNetB5(include_top=False, input_tensor=img_input, weights=weights, pooling=None,
                                    input_shape=img_shape)

    elif effnet_version == 'B6':
        effnet = efn.EfficientNetB6(include_top=False, input_tensor=img_input, weights=weights, pooling=None,
                                    input_shape=img_shape)

    else:
        effnet = efn.EfficientNetB7(include_top=False, input_tensor=img_input, weights=weights, pooling=None,
                                    input_shape=img_shape)

    # flat = tf.keras.layers.Flatten()(res_net.output)
    avg_pool = GlobalAveragePooling2D()(effnet.output)
    # dropout = tf.keras.layers.Dropout(0.5)(avg_pool)
    y_hat = Dense(2, activation="sigmoid")(avg_pool)
    model = keras.models.Model(effnet.input, y_hat)

    return model

def compile_model(name_file_rede, model, opt, fold, version):
    ic('\n Compilando rede: ', name_file_rede)
    model.compile(loss='binary_crossentropy', optimizer=opt, metrics=['accuracy'])

    ic('\n ** Plotting model and callbacks...')
    name_weights = 'Train_model_weights_%s_{epoch:02d}_%s.h5' % (name_file_rede, version)
    csv_name = 'training_{}_fold_{}_ver_{}.csv'.format(name_file_rede, fold, version)
    callbacks = utils.get_callbacks(name_weights=name_weights, patience_lr=10, name_csv=csv_name)

    return callbacks

def fit_model(model, generator, x_data_cv, batch_size, num_epochs, x_val_cv, y_val_cv, callbacks):
    history = model.fit_generator(
        generator,
        steps_per_epoch=len(x_data_cv) / batch_size,
        epochs=num_epochs,
        verbose=1,
        validation_data=(x_val_cv, y_val_cv),
        validation_steps=len(x_val_cv) / batch_size,
        callbacks=callbacks)

    return history

def get_model_resnet(x_data, weights, resnet_depth):
    keras_back.set_image_data_format('channels_last')
    img_shape = (x_data.shape[1], x_data.shape[2], x_data.shape[3])
    ic('Input Shape Matrix: ', img_shape)
    # ic('X_data Shape: 1- {}, 2- {}, 3- {}'.format(x_data.shape[1], x_data.shape[2], x_data.shape[3]))
    img_input = Input(shape=img_shape)

    ic('\n ** Utilizando a Rede Resnet', resnet_depth)
    ic('\n ** Pesos carregados: ', weights)

    if resnet_depth == 50:

        res_net = ResNet50(include_top=False, weights=weights, input_tensor=img_input, input_shape=img_shape,
                           pooling=None)
        avg_pool = GlobalAveragePooling2D()(res_net.output)
        # dropout = tf.keras.layers.Dropout(0.5)(flat)
        # y_hat = tf.keras.layers.Dense(2, activation='sigmoid')(avg_pool)
        y_hat = Dense(2, activation='softmax')(avg_pool)
        model = keras.models.Model(res_net.input, y_hat)

    elif resnet_depth == 101:
        res_net = ResNet101V2(include_top=False, weights=weights, input_tensor=img_input, input_shape=img_shape,
                              pooling=None)
        avg_pool = GlobalAveragePooling2D()(res_net.output)
        # dropout = tf.keras.layers.Dropout(0.5)(flat)
        # y_hat = tf.keras.layers.Dense(2, activation='sigmoid')(avg_pool)
        y_hat = Dense(2, activation='softmax')(avg_pool)
        model = keras.models.Model(res_net.input, y_hat)

    else:
        res_net = ResNet152V2(include_top=False, weights=weights, input_tensor=img_input, input_shape=img_shape,
                              pooling=None)
        avg_pool = GlobalAveragePooling2D()(res_net.output)
        # dropout = tf.keras.layers.Dropout(0.5)(flat)
        # y_hat = tf.keras.layers.Dense(2, activation='sigmoid')(avg_pool)
        y_hat = Dense(2, activation='softmax')(avg_pool)
        model = keras.models.Model(res_net.input, y_hat)

    return model

def get_model_inception(x_data, weights, version):
    keras_back.set_image_data_format('channels_last')
    img_shape = (x_data.shape[1], x_data.shape[2], x_data.shape[3])
    ic('Input Shape Matrix: ', img_shape)
    # ic('X_data Shape: 1- {}, 2- {}, 3- {}'.format(x_data.shape[1], x_data.shape[2], x_data.shape[3]))
    img_input = Input(shape=img_shape)

    ic('\n ** Utilizando a Inception', version)
    ic('\n ** Pesos carregados: ', weights)

    if version == 'V2':
        res_net = InceptionResNetV2(include_top=False, weights=weights, input_tensor=img_input, input_shape=img_shape, pooling=None)
    elif version == 'V3':
        res_net = InceptionV3(include_top=False, weights=weights, input_tensor=img_input, input_shape=img_shape, pooling=None)

    avg_pool = GlobalAveragePooling2D()(res_net.output)
    y_hat = Dense(2, activation='softmax')(avg_pool)
    model = keras.models.Model(res_net.input, y_hat)

    return model

def get_model_xception(x_data, weights):
    keras_back.set_image_data_format('channels_last')
    img_shape = (x_data.shape[1], x_data.shape[2], x_data.shape[3])
    ic('Input Shape Matrix: ', img_shape)
    # ic('X_data Shape: 1- {}, 2- {}, 3- {}'.format(x_data.shape[1], x_data.shape[2], x_data.shape[3]))
    img_input = Input(shape=img_shape)

    ic('\n ** Utilizando a Xception')
    ic('\n ** Pesos carregados: ', weights)

    res_net = Xception(include_top=False, weights=weights, input_tensor=img_input, input_shape=img_shape, pooling=None)

    avg_pool = GlobalAveragePooling2D()(res_net.output)
    y_hat = Dense(2, activation='softmax')(avg_pool)
    model = keras.models.Model(res_net.input, y_hat)

    return model


import torchvision
from torchvision import transforms

from .libsvm import LibSVM, DOWNLOAD_LINKS

libsvm_keys = list(DOWNLOAD_LINKS.keys())

def load_data(path, dataset, load_train_set=True, download=True):
    dataset = dataset.lower()
    train_set = None

    if dataset.startswith("cifar"):  # CIFAR-10/100
        transform_train = transforms.Compose([
            transforms.ToTensor(),
        ])

        transform_test = transforms.Compose([
            transforms.ToTensor(),
        ])

        if dataset == "cifar10":
            if load_train_set:
                train_set = torchvision.datasets.CIFAR10(
                    root=path, train=True, download=download, transform=transform_train)
            test_set = torchvision.datasets.CIFAR10(
                root=path, train=False, download=download, transform=transform_test)
        elif dataset == "cifar100":
            if load_train_set:
                train_set = torchvision.datasets.CIFAR100(
                    root=path, train=True, download=download, transform=transform_train)
            test_set = torchvision.datasets.CIFAR100(
                root=path, train=False, download=download, transform=transform_test)
        else:
            raise NotImplementedError(f'{dataset} is not implemented.')
    elif dataset == 'mnist':
        transform = transforms.ToTensor()
        if load_train_set:
            train_set = torchvision.datasets.MNIST(
                root=path, train=True, download=True, transform=transform)
        test_set = torchvision.datasets.MNIST(
            root=path, train=False, download=True, transform=transform)

    elif dataset in libsvm_keys:
        test_set = LibSVM(root=path, dataset_name=dataset)
        if load_train_set:
            train_set = test_set
    else:
        raise NotImplementedError(f'{dataset} is not implemented.')

    return train_set, test_set

def get_num_classes(dataset):
    dataset = dataset.lower()
    if dataset == 'cifar10':
        num_classes = 10
    elif dataset == 'cifar100':
        num_classes = 100
    elif dataset == 'mnist':
        num_classes = 10
    elif dataset in libsvm_keys:
        num_classes = 1
    else:
        raise ValueError(f"Dataset {dataset} is not supported.")
    return num_classes


import pygame
from src.persistence import Persistence
from src.button import Button, MenuButton
from src.game import Game
from src.contextmenu import Context
from src.soldier import *
from src.tower import *
from src.core import *
from src.tile import *
from src.tileview import *
import os

pygame.init()

SCREEN_WIDTH = 1248  # 1280
SCREEN_HEIGHT = 720

screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption('Castle Royale')

simulation_starting_time = 0  # 1 second = 1000 millisecond
clock = pygame.time.Clock()
FPS = 60

hamburger = Context()

# 1 - Menu
# 2 - New Game
# 3 - Load Game
# 4 - In Game

game_state = 1

in_game_menu = ["MAIN MENU", "BUILD", "TRAIN", "MOVE", "CONTINUE"]

main_menu_scale = 5
# MENU ASSETS
menu_images = {
    "BOTTOM": pygame.image.load("assets/menu_assets/also_tabla.png"),
    "MIDDLE": pygame.image.load("assets/menu_assets/koztes_tabla.png"),
    "LOGO": pygame.image.load("assets/menu_assets/logo_tabla.png"),
    "BG": pygame.image.load("assets/menu_assets/menu_bg.png")
}
menu_images["LOGO"] = pygame.transform.scale(menu_images["LOGO"], (
menu_images["LOGO"].get_width() * main_menu_scale, menu_images["LOGO"].get_height() * main_menu_scale))

main_menu = [(1, "NEW GAME"), (4, "SAVE"), (3, "LOAD"), (2, "QUIT")]
main_menu_btn = []
menu_top_padding = 0
for i, elem in enumerate(main_menu):
    if (len(main_menu) - 1) == i:
        img = menu_images['BOTTOM']
    else:
        img = menu_images['MIDDLE']

    padding = (i * (menu_images['MIDDLE'].get_height() * main_menu_scale))

    main_menu_btn.append(MenuButton(SCREEN_WIDTH / 2 - img.get_width() * main_menu_scale / 2,
                                    menu_top_padding + menu_images["LOGO"].get_height() + padding,
                                    img.get_width() * main_menu_scale, img.get_height() * main_menu_scale, elem[1],
                                    (len(main_menu) - 1) == i))

btn_quit = Button((255, 0, 0), 18, 676, 228, 40, "MAIN MENU")
btn_build = Button((255, 0, 0), 264, 676, 228, 40, "BUILD")
btn_train = Button((255, 0, 0), 510, 676, 228, 40, "TRAIN")
btn_move = Button((255, 0, 0), 756, 676, 228, 40, "MOVE")
btn_continue = Button((255, 0, 0), 1002, 676, 228, 40, "CONTINUE")
btn_end = Button((255, 0, 0), SCREEN_WIDTH / 2 - 40, SCREEN_HEIGHT / 2 - 25, 80, 50, "MENU")

btn_back = Button((255, 0, 0), 18, 676, 228, 40, "MAIN MENU")

game = None
persistence = Persistence()

font = pygame.font.SysFont('comicsans', 20)
font2 = pygame.font.Font(os.path.join("assets", "fonts", 'arcadeclassic.ttf'), 23)
font_game_over = pygame.font.Font(os.path.join("assets", "fonts", 'arcadeclassic.ttf'), 100)
SCALE = 3

gold_coin = pygame.image.load("assets/gold_coin.png")
gold_coin = pygame.transform.scale(gold_coin, (16 * SCALE, 16 * SCALE))

game_icon = pygame.image.load("assets/tile_assets/blue_castle.png")
pygame.display.set_icon(game_icon)

sword = pygame.image.load("assets/sword.png")
sword = pygame.transform.scale(sword, (16 * 2, 16 * 2))

health_bar_left = pygame.image.load("assets/health_bar.png")
health_bar_left = pygame.transform.scale(health_bar_left, (44 * SCALE, 9 * SCALE))

health_bar_inside_left = pygame.image.load("assets/health_bar_inside.png")
health_bar_inside_left = pygame.transform.scale(health_bar_inside_left, (44 * SCALE, 9 * SCALE))

health_bar_inside_original_left = pygame.image.load("assets/health_bar_inside.png")
health_bar_inside_original_left = pygame.transform.scale(health_bar_inside_original_left, (44 * SCALE, 9 * SCALE))

health_bar_right = pygame.image.load("assets/health_bar.png")
health_bar_right = pygame.transform.flip(health_bar_right, True, False)
health_bar_right = pygame.transform.scale(health_bar_right, (44 * SCALE, 9 * SCALE))

health_bar_inside_right = pygame.image.load("assets/health_bar_inside.png")
health_bar_inside_right = pygame.transform.flip(health_bar_inside_right, True, False)
health_bar_inside_right = pygame.transform.scale(health_bar_inside_right, (44 * SCALE, 9 * SCALE))

health_bar_inside_original_right = pygame.image.load("assets/health_bar_inside.png")
health_bar_inside_original_right = pygame.transform.flip(health_bar_inside_original_right, True, False)
health_bar_inside_original_right = pygame.transform.scale(health_bar_inside_original_right, (44 * SCALE, 9 * SCALE))

buttons_background = pygame.image.load("assets/in_game_button_assets/button_background.png")
lock = pygame.image.load("assets/lock.png")
lock = pygame.transform.scale(lock, (16 * 2, 16 * 2))
lock2 = pygame.image.load("assets/lock.png")
lock2 = pygame.transform.scale(lock, (16 * 3, 16 * 3))

# Change name color based on round
name_color = (0, 0, 0)
name_color_now_playing = (255, 0, 0)

tile_view = []

load_buttons = []
def create_buttons(num_buttons):
    global load_buttons
    load_buttons = []
    for i in range(num_buttons):
        load_buttons.append(Button((255, 0, 0), (i * 228) + ((i+1) * 18), 18, 228, 40, "MAP {}".format(i)))

def fill(surface, original_surface, hue):
    if hue < 0:
        hue = 0
    w, h = surface.get_size()
    for x in range(w):
        for y in range(h):
            a = surface.get_at((x, y))[3]
            c = original_surface.get_at((x, y))
            c.hsla = (c.hsla[0] + hue, c.hsla[1], c.hsla[2], c.hsla[3])
            surface.set_at((x, y), pygame.Color(c.r, c.g, c.b, a))

def get_name_color(player):
    if game.current_player == player:
        return name_color_now_playing
    return name_color

def translate(value, left_min, left_max, right_min, right_max):
    left_span = left_max - left_min
    right_span = right_max - right_min

    value_scaled = float(value - left_min) / float(left_span)

    return right_min + (value_scaled * right_span)

def update_tile_view(map):
    if(tile_view == []):
        for i in range(0, 14):
            tile_view.append([])
            for j in range(0, 26):
                tile_view[i].append(TileView(game, i, j))
                tile_view[i][j]._tile = map[i][j]
    else:
        for i in range(0, 14):
            for j in range(0, 26):
                tile_view[i][j]._tile = map[i][j]

run = True
while run:
    clock.tick(FPS)
    # Background color for testing
    screen.fill((247, 217, 181))
    # Draw update function
    if game_state == 1:

        screen.blit(menu_images["BG"], [0, 0, SCREEN_WIDTH, SCREEN_HEIGHT])

        screen.blit(menu_images["LOGO"], [SCREEN_WIDTH / 2 - img.get_width() * main_menu_scale / 2, menu_top_padding,
                                          menu_images["LOGO"].get_width(), menu_images["LOGO"].get_height()])

        for elem in main_menu_btn:
            elem.draw(screen)

        if game is None:
            screen.blit(lock2, [410, 320, lock.get_width(), lock.get_height()])
        # draw
    elif game_state == 2:
        screen.blit(menu_images["BG"], [0, 0, SCREEN_WIDTH, SCREEN_HEIGHT])

        if game.winner:
            game_over_title = font_game_over.render("{} WINS".format(game.winner.name), True, (255, 215, 166))
            game_over_title_outlines = font_game_over.render("{} WINS".format(game.winner.name), True, (39, 28, 57))

            screen.blit(game_over_title_outlines, ((SCREEN_WIDTH / 2 - game_over_title.get_width() / 2) - 4,
                                                   (SCREEN_HEIGHT / 2 - game_over_title.get_height() / 2) + 4))
            screen.blit(game_over_title, (
            SCREEN_WIDTH / 2 - game_over_title.get_width() / 2, SCREEN_HEIGHT / 2 - game_over_title.get_height() / 2))

        btn_back.draw(screen)
        btn_back.is_over(pygame.mouse.get_pos())
        # draw
    elif game_state == 3:
        pass
    elif game_state == 5:
        screen.blit(menu_images["BG"], [0, 0, SCREEN_WIDTH, SCREEN_HEIGHT])

        for i in load_buttons:
            i.draw(screen)
            i.is_over(pygame.mouse.get_pos())

        btn_back.draw(screen)
        btn_back.is_over(pygame.mouse.get_pos())

        # draw
    elif game_state == 4:
        # Draw Tiles

        screen.blit(buttons_background, [0, 672, SCREEN_WIDTH, 200])

        if not game.is_ended:
            update_tile_view(game._map)
            for i in range(0, 14):
                for j in range(0, 26):
                    tile_view[i][j].draw(screen)
            for i in range(0, 14):
                for j in range(0, 26):
                    tile_view[i][j].draw_buildings_and_soldiers(screen)
            for i in range(0, 14):
                for j in range(0, 26):
                    tile_view[i][j].draw_context_menu_for_tiles(screen)
        else:
            text = pygame.font.SysFont('Arial', 25).render("The winner is {0}.".format(game.winner), False, (0, 0, 0))
            screen.blit(text, (SCREEN_WIDTH / 2 - text.get_width() / 2, SCREEN_HEIGHT / 2 - 60))
            btn_end.draw(screen)

        # Draw Player Information
        if not game.is_ended:
            p1_name = str(game.player_1.name)
            p2_name = str(game.player_2.name)

            player1_name = font2.render(p1_name, True, (255, 255, 255))
            player1_name_outlines = font2.render(p1_name, True, (39, 28, 57))
            player1_money = font2.render(str(game.player_1.gold), True, pygame.Color(39, 28, 57))
            player1_money_outline = font2.render(str(game.player_1.gold), True, pygame.Color(255, 206, 73))

            hud_pos_x = 10
            hud_pos_y = 70

            player1_hp = game.player_1.castle_tile.units[0].health
            fill(health_bar_inside_left, health_bar_inside_original_left, math.floor(player1_hp / 10))

            screen.blit(player1_name_outlines, (30, 15, 200, 200))
            screen.blit(player1_name, (30 + 2, 15 - 2, 200, 200))

            player1_display_hp = translate(player1_hp, 0, 1000, 160, 1000)
            screen.blit(health_bar_left, (20, 45, 200, 200))
            screen.blit(health_bar_inside_left, (20, 45, 200, 200),
                        (0, 0, (health_bar_inside_left.get_width() * player1_display_hp) / 1000, 200))

            screen.blit(gold_coin, (hud_pos_x, hud_pos_y))
            screen.blit(player1_money, (-5 + hud_pos_x + gold_coin.get_width(),
                                        hud_pos_y + (gold_coin.get_height() / 2 - player1_money.get_height() / 2)))
            screen.blit(player1_money_outline, (-5 + 3 + hud_pos_x + gold_coin.get_width(), hud_pos_y + (
                        gold_coin.get_height() / 2 - player1_money_outline.get_height() / 2) - 1))

            player2_name = font2.render(p2_name, True, (255, 255, 255))
            player2_name_outlines = font2.render(p2_name, True, (39, 28, 57))
            player2_money = font2.render(str(game.player_2.gold), True, pygame.Color(39, 28, 57))
            player2_money_outline = font2.render(str(game.player_2.gold), True, pygame.Color(255, 206, 73))

            if not game.start_simulation:
                if game.current_player == game.player_1:
                    screen.blit(sword, (30 + player1_name.get_width(), 8))
                else:
                    screen.blit(pygame.transform.flip(sword, flip_x=True, flip_y=False),
                                (SCREEN_WIDTH - 30 - sword.get_width() - player2_name.get_width(), 8))

            player2_hp = game.player_2.castle_tile.units[0].health
            fill(health_bar_inside_right, health_bar_inside_original_right, math.floor(player2_hp / 10))
            screen.blit(health_bar_right, (SCREEN_WIDTH - 20 - health_bar_right.get_width(), 45, 200, 200))

            player2_display_hp = translate(player2_hp, 0, 1000, 160, 1000)
            screen.blit(health_bar_inside_right, (SCREEN_WIDTH - 20 - health_bar_inside_right.get_width() + (
                        health_bar_inside_right.get_width() - (
                            (health_bar_inside_right.get_width() * player2_display_hp) / 1000)), 45, 200, 200), (
                        health_bar_inside_right.get_width() - (
                                    (health_bar_inside_right.get_width() * player2_display_hp) / 1000), 0,
                        health_bar_inside_right.get_width(), 200))

            screen.blit(player2_name_outlines, (SCREEN_WIDTH - 30 - player2_name_outlines.get_width(), 15, 200, 200))
            screen.blit(player2_name, (SCREEN_WIDTH - 30 - player2_name_outlines.get_width() + 2, 15 - 2, 200, 200))

            screen.blit(gold_coin, (SCREEN_WIDTH - gold_coin.get_width() - hud_pos_x, hud_pos_y))
            screen.blit(player2_money, (
            5 + SCREEN_WIDTH - gold_coin.get_width() - hud_pos_x - player2_money.get_width(),
            hud_pos_y + (gold_coin.get_height() / 2 - player2_money.get_height() / 2)))
            screen.blit(player2_money_outline, (
            5 + 3 + SCREEN_WIDTH - gold_coin.get_width() - hud_pos_x - player2_money_outline.get_width(),
            hud_pos_y + (gold_coin.get_height() / 2 - player2_money_outline.get_height() / 2) - 1))

            if not game.start_simulation:
                current_player_state = font2.render(str(game.current_player.state), True, (255, 255, 255))
                current_player_state_outline = font2.render(str(game.current_player.state), True, (39, 28, 57))
            else:
                current_player_state = font2.render("Simulation", True, (255, 255, 255))
                current_player_state_outline = font2.render("Simulation", True, (39, 28, 57))

            screen.blit(current_player_state_outline,
                        ((SCREEN_WIDTH / 2) - 2 - (current_player_state.get_width() / 2), 10 + 2))
            screen.blit(current_player_state, ((SCREEN_WIDTH / 2) - (current_player_state.get_width() / 2), 10))

            if game_state == 4 and game.winner:
                game_state = 2

        if not game.start_simulation:
            if not hamburger.opened:
                btn_quit.is_over(pygame.mouse.get_pos())
                btn_build.is_over(pygame.mouse.get_pos())
                btn_train.is_over(pygame.mouse.get_pos())
                btn_move.is_over(pygame.mouse.get_pos())
                btn_continue.is_over(pygame.mouse.get_pos())
                btn_end.is_over(pygame.mouse.get_pos())

        if not game.is_ended:
            btn_quit.draw(screen)
            btn_build.draw(screen)
            btn_train.draw(screen)
            btn_move.draw(screen)
            btn_continue.draw(screen)

        if game.start_simulation:
            screen.blit(lock, [17, 682, lock.get_width(), lock.get_height()])
            screen.blit(lock, [263, 682, lock.get_width(), lock.get_height()])
            screen.blit(lock, [509, 682, lock.get_width(), lock.get_height()])
            screen.blit(lock, [755, 682, lock.get_width(), lock.get_height()])
            screen.blit(lock, [1001, 682, lock.get_width(), lock.get_height()])
        # Draw Context Menu
        hamburger.draw(screen)
        hamburger.is_over(pygame.mouse.get_pos())

        if game.start_simulation:
            game.simulate()

    # Keyboard input update function
    for event in pygame.event.get():
        if game_state == 1:
            if event.type == pygame.MOUSEBUTTONUP:
                for i, btn in enumerate(main_menu_btn):
                    if btn.is_over(pygame.mouse.get_pos()):
                        if main_menu[i][0] == 1:
                            game = Game()
                            game.new_game(1000, "Player1", "Player2")
                            game_state = 4
                            tile_view = []
                        elif main_menu[i][0] == 2:
                            run = False
                        elif main_menu[i][0] == 3:
                            num_saves = persistence.get_num_saves()
                            create_buttons(num_saves)
                            game_state = 5
                        elif main_menu[i][0] == 4:
                            persistence.save(game)
                            game = None

        elif game_state == 2:
            if event.type == pygame.MOUSEBUTTONUP and (not game or not game.start_simulation):
                if btn_back.is_over(pygame.mouse.get_pos()):
                    game_state = 1

        elif game_state == 5:
            if event.type == pygame.MOUSEBUTTONUP and (not game or not game.start_simulation):
                if btn_back.is_over(pygame.mouse.get_pos()):
                    game_state = 1
                for index, btn in enumerate(load_buttons):
                    if btn.is_over(pygame.mouse.get_pos()):
                        load = persistence.load(index)
                        if load:
                            game = load
                            tile_view = []
                        game_state = 4

        elif game_state == 3:
            pass
            # keyboard input
        elif game_state == 4:
            if not game.start_simulation:
                if hamburger.opened:
                    if event.type == pygame.MOUSEBUTTONUP:
                        content = hamburger.is_over(pygame.mouse.get_pos())
                        for item in content:
                            if item[2] and item[1] == "Tower":
                                game.current_tile.build(game.current_player, "BasicTower")
                            if item[2] and item[1] == "Slow":
                                game.current_tile.build(game.current_player, "Slow")
                            if item[2] and item[1] == "Splash":
                                game.current_tile.build(game.current_player, "Splash")
                            if item[2] and item[1] == "Barracks":
                                print("barracks")
                            if item[2] and item[1] == "Upgrade":
                                game.current_tile.upgrade_tower()
                            if item[2] and item[1] == "Demolish":
                                game.current_tile.demolish_tower()
                            if item[2] and item[1] == "Remove":
                                game.current_tile.remove_tower_ruin()
                            if item[2] and item[1] == "Soldier":
                                game.current_player.castle_tile.train(game.current_player, "BasicSoldier")
                            if item[2] and item[1] == "Climber":
                                game.current_player.castle_tile.train(game.current_player, "Climber")
                            if item[2] and item[1] == "Suicide":
                                game.current_player.castle_tile.train(game.current_player, "Suicide")
                            if item[2] and item[1] == "Tank":
                                game.current_player.castle_tile.train(game.current_player, "Tank")
                            if item[2] and item[1] == "Select":
                                game.select_current_tile()
                            if item[2] and item[1] == "Waypoint":
                                game.add_waypoint()
                            if item[2] and item[1] == "Reset":
                                game.reset_waypoint()
                            hamburger.opened = False
                else:
                    for row in game.map:
                        for tile in row:
                            if tile.is_over(pygame.mouse.get_pos()) and event.type == pygame.MOUSEBUTTONUP:
                                if game.current_player.state == "BUILD":
                                    mouse_cords = pygame.mouse.get_pos()
                                    if tile.has_building and issubclass(type(tile.units[0]), Tower) and not tile.units[
                                        0].is_in_ruins:
                                        hamburger.change_content(["Upgrade", "Demolish"])
                                    elif tile.has_building and issubclass(type(tile.units[0]), Tower) and tile.units[
                                        0].is_in_ruins:
                                        hamburger.change_content(["Remove"])
                                    else:
                                        hamburger.change_content(["Tower", "Slow", "Splash"])
                                    hamburger.open(mouse_cords[0], mouse_cords[1])

                                elif game.current_player.state == "TRAIN":
                                    mouse_cords = pygame.mouse.get_pos()
                                    hamburger.change_content(["Soldier", "Climber", "Suicide", "Tank"])
                                    hamburger.open(mouse_cords[0], mouse_cords[1])

                                elif game.current_player.state == "MOVE":
                                    mouse_cords = pygame.mouse.get_pos()
                                    options = []
                                    if tile_can_be_selected(tile):
                                        options.append("Select")

                                    if game.selected_tile and tile_can_be_added_as_waypoint(tile, game.selected_tile):
                                        options.append("Waypoint")

                                    if game.selected_tile:
                                        options.append("Reset")
                                    hamburger.change_content(options)
                                    if len(options) > 0:
                                        hamburger.open(mouse_cords[0], mouse_cords[1])
                                game.current_tile = tile
                    if event.type == pygame.MOUSEBUTTONUP and not game.start_simulation:
                        if btn_continue.is_over(pygame.mouse.get_pos()):
                            btn_continue.is_over((0, 0))
                            game.next_round()
                            for row in game.map:
                                for tile in row:
                                    tile.selected = False
                                    tile.waypoint = False
                            if game.start_simulation:
                                simulation_starting_time = pygame.time.get_ticks()
                        elif btn_quit.is_over(pygame.mouse.get_pos()):
                            game_state = 1
                        elif btn_build.is_over(pygame.mouse.get_pos()):
                            game.current_player.state = "BUILD"
                        elif btn_train.is_over(pygame.mouse.get_pos()):
                            game.current_player.state = "TRAIN"
                        elif btn_move.is_over(pygame.mouse.get_pos()):
                            game.current_player.state = "MOVE"
                        elif btn_end.is_over(pygame.mouse.get_pos()):
                            pass
        if event.type == pygame.QUIT:
            run = False

    pygame.display.update()

pygame.quit()


from datetime import datetime

path = 'Data/log.txt'

def save_log(message: str) -> None:
    file = open(path, 'a')
    message = message.replace("['", "").replace("']", " ")
    file.write(f"[{datetime.today().strftime('%d/%m/%Y %H:%M:%S')}]: {message}\n")
    file.close()


def calc_model_cosine(self, decimate=None, mode='err'):
        """
        Calculates the cosine of the residuals with the model.

        Parameters
        ----------
            decimate : Int or None, optional
                Decimate the residuals by `decimate` pixels. If None, no
                decimation is used. Valid only with mode='svd'. Default
                is None
            mode : {'svd', 'err'}
                Which mode to use; see Notes section. Default is 'err'.

        Returns
        -------
            abs_cos : numpy.float64
                The absolute value of the model cosine.

        Notes
        -----
        The model cosine is defined in terms of the geometric view of
        curve-fitting, as a model manifold embedded in a high-dimensional
        space. The model cosine is the cosine of the residuals vector
        with its projection on the tangent space: :math:`cos(phi) = |P^T r|/|r|`
        where :math:`P^T` is the projection operator onto the model manifold
        and :math:`r` the residuals. This can be calculated two ways: By
        calculating the projection operator P directly with SVD (mode=`svd`),
        or by using the expected error if the model were linear to calculate
        a model sine first (mode=`err`). Since the SVD of a large matrix is
        slow, mode=`err` is faster.

        `decimate` allows for every nth pixel only to be counted in the
        SVD matrix of J for speed. While this is n x faster, it is
        considerably less accurate, so the default is no decimation.
        """
        if mode == 'svd':
            slicer = slice(0, None, decimate)

            #1. Calculate projection term
            u, sig, v = np.linalg.svd(self.J[:,slicer], full_matrices=False) #slow part
            # p = np.dot(v.T, v) - memory error, so term-by-term
            r = self.calc_residuals()[slicer]
            abs_r = np.sqrt((r*r).sum())

            v_r = np.dot(v,r/abs_r)
            projected = np.dot(v.T, v_r)

            abs_cos = np.sqrt((projected*projected).sum())
        elif mode == 'err':
            expected_error = self.find_expected_error(delta_params='perfect')
            model_sine_2 = expected_error / self.error  #error = distance^2
            abs_cos = np.sqrt(1 - model_sine_2)
        else:
            raise ValueError('mode must be one of `svd`, `err`')
        return abs_cos


# LIBRARIES
import pygame
from queue import LifoQueue, PriorityQueue
import time
import tkinter.filedialog as tkf
import pickle
import random
import preloadedmazes as mazes

###############################################################################################
# INITIAL SETUP
###############################################################################################

#PROGRAM INFORMATION
version = "V5.0"
lastUpdate = "12/6/2022"

#WINDOW DISPLAY
WIDTH = 800
WIN = pygame.display.set_mode((WIDTH, WIDTH))
pygame.display.set_caption("Visual Application of Pathfinding Algorithms")

#COLOR SET
RED = (255, 0, 0)
GREEN = (0, 255, 0)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
PURPLE = (128, 0, 128)
ORANGE = (255, 165 ,0)
GREY = (128, 128, 128)
TURQUOISE = (64, 224, 208)

#DEFINES EACH GRID AND DEFINES STATUS
class Node:
    def __init__(self, row, col, width, total_rows):
        self.row = row
        self.col = col
        self.x = row * width
        self.y = col * width
        self.color = WHITE
        self.neighbors = []
        self.width = width
        self.total_rows = total_rows

    def get_pos(self):
        return self.row, self.col

    def is_closed(self):
        return self.color == RED

    def is_open(self):
        return self.color == GREEN

    def is_barrier(self):
        return self.color == BLACK

    def is_start(self):
        return self.color == ORANGE

    def is_end(self):
        return self.color == TURQUOISE

    def is_path(self):
        return self.color == PURPLE

    def is_opennode(self):
        return self.color == WHITE

    def reset(self):
        self.color = WHITE

    def make_start(self):
        self.color = ORANGE

    def make_closed(self):
        self.color = RED

    def make_open(self):
        self.color = GREEN

    def make_barrier(self):
        self.color = BLACK

    def make_end(self):
        self.color = TURQUOISE

    def make_path(self):
        self.color = PURPLE

    def draw(self, win):
        pygame.draw.rect(win, self.color, (self.x, self.y, self.width, self.width))

    #DETERMINES BARRIERS ON A GLOBAL LEVEL BEFORE ALGORITHM STARTS
    def update_neighbors(self, grid):
        self.neighbors = []
        if self.row < self.total_rows - 1 and not grid[self.row + 1][self.col].is_barrier(): #DOWN
            self.neighbors.append(grid[self.row + 1][self.col])

        if self.row > 0 and not grid[self.row - 1][self.col].is_barrier(): #UP
            self.neighbors.append(grid[self.row - 1][self.col])

        if self.col < self.total_rows - 1 and not grid[self.row][self.col + 1].is_barrier(): #RIGHT
            self.neighbors.append(grid[self.row][self.col + 1])

        if self.col > 0 and not grid[self.row][self.col - 1].is_barrier(): #LEFT
            self.neighbors.append(grid[self.row][self.col - 1])

    def __lt__(self, other):
        return False

#FUNCTION TO FIND TAXI-CAB DISTANCE
def h(p1, p2):
    x1, y1 = p1
    x2, y2 = p2
    return abs(x1 - x2) + abs(y1 - y2)

#SHOWS SHORTEST PATH
def reconstruct_path(came_from, current, draw):
    while current in came_from:
        current = came_from[current]
        current.make_path()
        draw()

###############################################################################################
# ALGORITHMS FOR SEACHING
###############################################################################################

#A* ALGORITHM
def astar(draw, grid, start, end):
    count = 0
    open_set = PriorityQueue() #SET THAT RETURNS THE NODE WITH SMALLEST VALUE WHEN CALLED
    open_set.put((0, count, start))
    came_from = {} #RECORDS PATH FROM PREVIOUS NODE
    g_score = {node: float("inf") for row in grid for node in row} #DISTANCE FROM THE START
    g_score[start] = 0
    f_score = {node: float("inf") for row in grid for node in row} #DISTANCE FROM THE END
    f_score[start] = h(start.get_pos(), end.get_pos())
    open_set_hash = {start} #SEARCHABLE QUEUE SINCE PRIORITYQUEUE DOESN'T ALLOW US TO

    while not open_set.empty():
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()

        current = open_set.get()[2]
        open_set_hash.remove(current)

        #SHOWS SHORTEST PATH
        if current == end:
            reconstruct_path(came_from, end, draw)
            end.make_end()
            start.make_start()
            return True

        for neighbor in current.neighbors:
            temp_g_score = g_score[current] + 1 #ONE SPACE AWAY FROM CURRENT NODE
            #REROUTES PATH AND UPDATES VALUE IF CLOSER PATH FOUND
            if temp_g_score < g_score[neighbor]:
                came_from[neighbor] = current
                g_score[neighbor] = temp_g_score
                f_score[neighbor] = temp_g_score + h(neighbor.get_pos(), end.get_pos())
                #OPENS PATH OPTION IN NODE IF IT WAS UNEXPLORED
                if neighbor not in open_set_hash:
                    count += 1
                    open_set.put((f_score[neighbor], count, neighbor))
                    open_set_hash.add(neighbor)
                    neighbor.make_open()
        draw()
        if current != start:
            current.make_closed()

#DIJKSTRA'S ALGORITHM
def dijkstra(draw, grid, start, end):
    count = 0
    open_set = PriorityQueue() #SET THAT RETURNS THE NODE WITH SMALELST VALUE WHEN CALLED
    open_set.put((0, count, start))
    came_from = {} #RECORDS PATH FROM PREVIOUS NODE
    g_score = {node: float("inf") for row in grid for node in row} #DISTANCE FROM THE START
    g_score[start] = 0
    open_set_hash = {start} #SEARCHABLE QUEUE SINCE PRIORITYQUEUE DOESN'T ALLOW US TO

    while not open_set.empty():
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()

        current = open_set.get()[2]
        open_set_hash.remove(current)

        #SHOWS SHORTEST PATH
        if current == end:
            reconstruct_path(came_from, end, draw)
            end.make_end()
            start.make_start()
            return True

        for neighbor in current.neighbors:
            temp_g_score = g_score[current] + 1 #ONE SPACE AWAY FROM CURRENT NODE
            #REROUTES PATH AND UPDATES VALUE IF CLOSER PATH FOUND
            if temp_g_score < g_score[neighbor]:
                came_from[neighbor] = current
                g_score[neighbor] = temp_g_score
                #OPENS PATH OPTION IN NODE IF IT WAS UNEXPLORED
                if neighbor not in open_set_hash:
                    count += 1
                    open_set.put((g_score[neighbor], count, neighbor))
                    open_set_hash.add(neighbor)
                    neighbor.make_open()
        draw()
        if current != start:
            current.make_closed()

#DEPTH FIRST SEARCH ALGORITHM
def dfs(draw, grid, start, end):
    count = 0
    open_set = LifoQueue() #LAST IN FIRST OUT QUEUE TO AGGRESIVELY SEARCH
    open_set.put((count, start))
    came_from = {} #KEEPS TRACK OF WHERE EACH EXPLORED NODE CAME FROM
    open_set_hash = {start} #LIST OF EXPLORED NODES

    while not open_set.empty():
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()

        current = open_set.get()[1]

        #SHOWS FOUND PATH
        if current == end:
            reconstruct_path(came_from, end, draw)
            end.make_end()
            start.make_start()
            return True

        #PUTS EACH AVAILABLE NEIGHBOR OF CURRENT NODE INTO LIFO QUEUE
        for neighbor in current.neighbors:
            if not neighbor in open_set_hash:
                came_from[neighbor] = current
            if neighbor not in open_set_hash:
                    count += 1
                    open_set.put((count, neighbor))
                    open_set_hash.add(neighbor)
                    if neighbor != end:
                        neighbor.make_open()
        draw()
        if current != start:
            current.make_closed()

###############################################################################################
# FUNCTIONS TO USE
###############################################################################################

#MAKES AREA FOR GRID SPACES
def make_grid(rows, width):
    grid = []
    gap = width // rows
    for i in range(rows):
        grid.append([])
        for j in range(rows):
            node = Node(i, j, gap, rows)
            grid[i].append(node)
    return grid

#DRAWS LINES TO VISUALLY DEFINE GRIDS
def draw_grid(win, rows, width):
    gap = width // rows
    for i in range(rows):
        pygame.draw.line(win, GREY, (0, i * gap), (width, i * gap))
    for j in range(rows):
        pygame.draw.line(win, GREY, (j * gap, 0), (j * gap, width))

#FILLS GRID WITH COLOR WHEN CLICKED AND UPDATES ENTIRE SCREEN
def draw(win, grid, rows, width):
    win.fill(WHITE)
    for row in grid:
        for node in row:
            node.draw(win)
    draw_grid(win, rows, width)
    pygame.display.update()

#TAKES PIXEL LOCATION AND DETERMINES GRID LOCATION
def get_clicked_pos(pos, rows, width):
    gap = width // rows
    y, x = pos
    row = y // gap
    col = x // gap
    return row, col

#COUNT PATH LENGTH
def path_length(grid):
    path_count = 0
    for row in grid:
        for node in row:
            if node.is_path():
                path_count += 1
    return path_count

#DETERMINES IF A PATH WAS FOUND
def path_stats(grid, start_time):
    for row in grid:
        for node in row:
            if node.is_path():
                print("--- %.2f seconds and %i path blocks ---" % ((time.time() - start_time), path_length(grid)))
                return
    print("No Path Was Found...")

#SAVES THE MAZE FILE
def save_mazefile(row, grid, node):
    try:
        #PROCESS TO STORE BARRIER LAYOUT IN THE FILE
        mazeLayout = []
        for row in grid:
            for node in row:
                if node.is_barrier():
                    mazeLayout.append("1")
                else:
                    mazeLayout.append("4")

        #Export data to pickle file with user choice of save location
        files = [('Pickle File', '*.pickle')]
        filename = tkf.asksaveasfilename(filetypes = files, defaultextension = files)
        with open(filename, 'wb') as f:
            pickle.dump(mazeLayout, f)
        print("Successfully saved the file.")

    except FileNotFoundError:
        print("Saving file failed. Reason: Aborted save process.")
    except:
        print("Saving file failed.")

#LOADS MAZE ONTO BOARD
def loading_maze(grid, mazeLayout):
    i = 0
    for row in grid:
        for node in row:
            if mazeLayout[i] == '1':
                node.make_barrier()
            else:
                node.reset()
            i+=1

#LOADS A MAZE FROM SELECTED FILE
def load_mazefile(grid):
    try:
        #USING TKINTER TO OPEN INTERFACE TO SELECT FILE
        files = [('Pickle File', '*.pickle')]
        filename = tkf.askopenfilename(filetypes = files, defaultextension = files)
        with open(filename, 'rb') as f:
            mazeLayout = pickle.load(f)

        #PROCESS TO BUILD MAZE BY READING FILE
        loading_maze(grid, mazeLayout)
        print("Successfully opened the maze layout file.")
    except FileNotFoundError:
        print("Opening file failed. Reason: Aborted opening process.")
    except:
        print("Opening file failed.")

# Instructions on use
def say_usermanual():
    print("")
    print("Welcome to the pathfinding algorithm visualizer by Jeffrey Collier!")
    print("---------------------------------------------------------------------")
    print("Current version: %s" % version)
    print("Last updated: %s" % lastUpdate)
    print("---------------------------------------------------------------------")
    print("• Your first left mouse click will place the start node, second click will place the end node, and then any remaining clicks will place walls.")
    print("• Your right mouse click will remove selected node.")
    print("• Keyboard press \"1\" runs the A* search algorithm.")
    print("• Keyboard press \"2\" runs the Dijkstra's algorithm.")
    print("• Keyboard press \"3\" runs the Depth-first search algorithm.")
    print("• Keyboard press \"R\" loads a random preloaded maze.")
    print("• Keyboard press \"C\" clears the board.")
    print("• Keyboard press \"S\" saves your current board layout.")
    print("• Keyboard press \"L\" loads a saved board layout.")
    print("• Keyboard press \"H\" prints this menu again for help.")

#LOADING A RANDOM PRELOADED MAZE
numberMemmory = [0]
def load_randomMaze(grid):
    try:
        lastNumber = numberMemmory.pop()

        currentNumber = random.randint(0,2)
        while currentNumber == lastNumber:
            currentNumber = random.randint(0,2)
        numberMemmory.append(currentNumber)

        listofmazes = [mazes.maze1, mazes.maze2, mazes.maze3]
        mazeLayout = listofmazes[currentNumber]
        loading_maze(grid, mazeLayout)
        print("Successfully loaded a random preloaded maze.")
    except Exception as e:
        print("Error: failed loading a random preloaded maze.")
        print(e)

###############################################################################################
# MAIN GAME LOOP
###############################################################################################

def main(win, width):
    say_usermanual()
    ROWS = 50
    grid = make_grid(ROWS, width)
    start = None #start node
    end = None #end node
    locked = False #board lock

    run = True
    while run:
        draw(win, grid, ROWS, width)
        for event in pygame.event.get():
            #QUITS PROGRAM IF REQUESTED
            if event.type == pygame.QUIT:
                run = False

            #CLEARS BOARD IF NEW ALGORITHM IS STARTED OR MOUSE BUTTON IS CLICKED
            if pygame.mouse.get_pressed()[0] or pygame.mouse.get_pressed()[2] or \
                (event.type == pygame.KEYDOWN and event.key == pygame.K_1) or \
                (event.type == pygame.KEYDOWN and event.key == pygame.K_2) or \
                (event.type == pygame.KEYDOWN and event.key == pygame.K_3):
                    if locked:
                        locked = False
                        for row in grid:
                            for node in row:
                                if node.is_open() or node.is_closed() or node.is_path():
                                    node.reset()

            #CREATES START, END, AND BARRIER NODES
            if pygame.mouse.get_pressed()[0]: #LEFT MOUSE CLICK
                pos = pygame.mouse.get_pos()
                row, col = get_clicked_pos(pos, ROWS, width)
                node = grid[row][col]
                if not start and node != end:
                    start = node
                    start.make_start()
                elif not end and node != start:
                    end = node
                    end.make_end()
                elif node != end and node != start:
                    node.make_barrier()

            #RESETS NODES WITH MOUSE CLICK
            elif pygame.mouse.get_pressed()[2]: #RIGHT MOUSE CLICK
                pos = pygame.mouse.get_pos()
                row, col = get_clicked_pos(pos, ROWS, width)
                node = grid[row][col]
                node.reset()
                if node == start:
                    start = None
                elif node == end:
                    end = None

            #KEYBOARD COMMANDS
            if event.type == pygame.KEYDOWN:
                #A* STAR ALGORITHM
                if event.key == pygame.K_1 and start and end: #RUNS PROGRAM
                    start_time = time.time()
                    for row in grid:
                        for node in row:
                            node.update_neighbors(grid)
                    astar(lambda: draw(win, grid, ROWS, width), grid, start, end)
                    locked = True
                    path_stats(grid, start_time)

                #DIJKSTRA'S ALGORITHM
                if event.key == pygame.K_2 and start and end: #RUNS PROGRAM
                    start_time = time.time()
                    for row in grid:
                        for node in row:
                            node.update_neighbors(grid)
                    dijkstra(lambda: draw(win, grid, ROWS, width), grid, start, end)
                    locked = True
                    path_stats(grid, start_time)

                #DEPTH FIRST SEARCH ALGORITHM
                if event.key == pygame.K_3 and start and end: #RUNS PROGRAM
                    start_time = time.time()
                    for row in grid:
                        for node in row:
                            node.update_neighbors(grid)
                    dfs(lambda: draw(win, grid, ROWS, width), grid, start, end)
                    locked = True
                    path_stats(grid, start_time)

                #LOADING A SAVED MAZE FILE ONTO BOARD
                if event.key == pygame.K_l:
                    grid = make_grid(ROWS, width)
                    start = None
                    end = None
                    load_mazefile(grid)

                #SAVING A MAZE BOARD INTO A FILE
                if event.key == pygame.K_s:
                    save_mazefile(row, grid, node)

                #CLEARS THE GRID BOARD
                if event.key == pygame.K_c:
                    start = None
                    end = None
                    grid = make_grid(ROWS, width)

                #PRINTS THE HELP MENU
                if event.key == pygame.K_h:
                    say_usermanual()

                #LOADS A RANDOM MAZE
                if event.key == pygame.K_r:
                    grid = make_grid(ROWS, width)
                    start = None
                    end = None
                    load_randomMaze(grid)

    pygame.quit()
main(WIN, WIDTH)