# -*- coding: utf-8 -*-
import re
import csv
import matplotlib.pyplot as plt

"""
This script is for plotting each node# (force/RMSE & TC(300K) diff from 112.1)
with classified color of each data#
"""

if __name__ == '__main__':
    datagrp=["2","10","20","40","60"]
    dlabels=['70','350','700','1400','2100']
    nodegrp=["50","100","500"]
    nodemrk=["x","+","."]
    dcolors=["red","orange","lime","b","cyan"]

    root="/home/okugawa/HDNNP/Si-190808"
    rstfile=root+"/result/RMSETCdata.csv"
    plotfile=root+"/result/grpplot/opt-node/node-all.png"

    #Read RMSETC.csv and make list of plotting data for each node#
    with open(rstfile, 'r') as rslt:
        readcsv = csv.reader(rslt)
        RMSETCd = [[] for i in range(len(nodegrp))]

        for row in readcsv:
            gnameall = re.split('[dn-]',row[0])
            dname= gnameall[1]
            nname= gnameall[2]
            if dname in datagrp and nname in nodegrp:
                dclr=dcolors[datagrp.index(dname)]
                nodenum=nodegrp.index(nname)
                RMSETCd[nodenum].append([float(row[1]),abs(float(row[2])),dclr])

    #Plotting all sample of force/RMSE and TC error for all node# group
    plottitle="Force/RMSE & TC err of each node# and data#"
    fig = plt.figure()
    ax = fig.add_subplot(111)
    plt.title(plottitle, fontsize=14)
    plt.xticks(fontsize=12)  #Font size of x-axis scale
    plt.yticks(fontsize=12)  #Font size of y-axis scale
    ax.set_xlabel("TC Err (from 112.1W/m-K:300K)",fontsize=12)
    ax.set_ylabel("force/RMSE (meV/A)",fontsize=12)
    ax.grid(True)
    plt.rcParams["legend.borderpad"] ='0.3'    #Space between handle & edge
    plt.rcParams["legend.handlelength"] ='1'   #Length of handle
    plt.rcParams["legend.handletextpad"] ='0.2' #Space between handle & label
    plt.rcParams["legend.columnspacing"] ='1'  #Space between columns
    plt.rcParams["legend.labelspacing"] ='0.2' #Space between row of label

    for i in range(len(nodegrp)):
        lbl="node="+nodegrp[i]
        for j,dtn in enumerate(RMSETCd[i]):
            ax.scatter(dtn[1],dtn[0],c=dtn[2],marker=nodemrk[i],s=20,alpha=0.7)

    left, right = ax.get_xlim()
    bottom, top = ax.get_ylim()
    ax.set_xlim(-1, right*1.1)
    ax.set_ylim(bottom, top*1.1)

    #ax1 is only for plotting legend of node#
    plt.rcParams["legend.edgecolor"] ='black'
    ax1 = ax.twinx()
    for i in range(len(nodegrp)):
        lbl="node="+nodegrp[i]
        ax1.scatter(dtn[1],dtn[0],c="black",marker=nodemrk[i],label=lbl)
    handler1, label1 = ax1.get_legend_handles_labels()
    leg1=ax.legend(handler1,label1,loc='upper left',ncol=3,fontsize=12,
              facecolor="lightgray") #Legend of node#
    fig.delaxes(ax1)

    #ax2 is only for plotting legend of all kind of data
    plt.rcParams["legend.edgecolor"] ='green'
    ax2 = ax.twinx()
    for j in range(len(dlabels)):
        lbl="data="+dlabels[j]
        clr=dcolors[j]
        ax2.scatter(dtn[1],dtn[0],c=clr,marker=',',label=lbl)
    handler2, label2 = ax2.get_legend_handles_labels()
    ax.legend(handler2, label2,loc='lower right',fontsize=12,
              title='Data#', title_fontsize=12)
    fig.delaxes(ax2)
    ax.add_artist(leg1)  #Write back legend of node#
    plt.savefig(plotfile)
    plt.close()


import bpy, mathutils

class RenderManager():

	files = None
	camera = None
	dummyObject = None
	zAxis = (0,0,1)

	blenderFilesDir = "."

	radius = 6378137

	def __init__(self, **kwargs):
		for k in kwargs:
			setattr(self, k, kwargs[k])

	def getBoundingBox(self):
		# perform context.scene.update(), otherwise o.matrix_world or o.bound_box are incorrect
		bpy.context.scene.update()
		xmin = float("inf")
		ymin = float("inf")
		xmax = float("-inf")
		ymax = float("-inf")
		for o in bpy.context.scene.objects:
			if o.type == "MESH" or o.type == "CURVE":
				for v in o.bound_box:
					(x,y,z) = o.matrix_world * mathutils.Vector(v)
					if x<xmin: xmin = x
					elif x>xmax: xmax = x
					if y<ymin: ymin = y
					elif y>ymax: ymax = y
		return {"xmin": xmin, "ymin": ymin, "xmax": xmax, "ymax": ymax}


#!/usr/bin/env python
# WS server that sends messages at random intervals
import asyncio
import datetime
import random
import websockets

async def time(websocket, path):
    print('Start... Loop előtt')
    while True:
        now = datetime.datetime.utcnow().isoformat() + 'Z'
        print(now)
        await websocket.send(now)
        await asyncio.sleep(random.random() * 3)

if __name__ == '__main__':
    print('Server is starting')
    start_server = websockets.serve(time, 'localhost', 5678)
    asyncio.get_event_loop().run_until_complete(start_server)
    print('Server started')
    asyncio.get_event_loop().run_forever()
    print('Forever vége')


""" 此文件用于解析rsl.out.0000文件 """

from datetime import datetime
# import matplotlib.pyplot as plt

class rslOutParser:
    """ 此类用于解析rsl.out.0000文件 """

    def __init__(self, rslFilePath):
        self.rslFilePath = rslFilePath
        self.dataLines = []

    def tryParse(self):
        """ a """
        print('aaaaa')
        with open(self.rslFilePath) as rslFile:
            self.dataLines = rslFile.readlines()
        # print(dataLines[:10])

        host = self.getHostRunningIn()
        cores = self.getCoresRunning()
        dateTimes, elapseSecs = self.getAllTime()

        # '%F %X 格式示例：2016-01-01 00:00:18'
        outData = {
            'host': host,
            'cores': cores,
            'dateTimes': [eachDateTime.strftime('%F %X') for eachDateTime in dateTimes],
            'elapseSecs': elapseSecs
        }

        # print('hahahaha', aa)
        # return str(aa)
        # return self.getDataLineNowTimeAndElapsedSec(firstIndex)
        return outData

    def getHostRunningIn(self):
        """ 此方法获取WRF在哪个主机上运行 """
        keyStr = 'hostname: '
        firstLine = self.dataLines[0]
        if keyStr in firstLine:
            return firstLine[firstLine.index(keyStr) + len(keyStr):].strip()
        else:
            return 'None'

    def getCoresRunning(self):
        """ 此方法获取用了几个核跑WRF """
        thirdLineWords = self.dataLines[2].split()
        x, y = thirdLineWords[3], thirdLineWords[-1]
        try:
            cores = int(x) * int(y)
        except ValueError:
            return 'unknown: parse error'
        return cores

    def findStartRunningIndex(self):
        """ 此方法找出第一行Timing for main:... """
        startCalIndex = 0
        headingStr = 'Timing for main:'
        while startCalIndex < len(self.dataLines):
            nowDataLine = self.dataLines[startCalIndex]
            if headingStr in nowDataLine and nowDataLine.index(headingStr) == 0:
                return startCalIndex
            startCalIndex += 1
        return 'None'

    def getAllTime(self):
        """_"""
        nowLineIndex = self.findStartRunningIndex()
        firstDateTime, _, _ = self.getDataLineNowTimeDomainAndElapsedSec(nowLineIndex)
        if firstDateTime is None:
            return [], []

        dateTimes, elapseSecs = [], []
        while nowLineIndex < len(self.dataLines):
            nowDateTime, nowDomain, nowElapseSec = self.getDataLineNowTimeDomainAndElapsedSec(nowLineIndex)
            if nowDateTime is None:
                nowLineIndex += 1
                continue
            dateTimes.append(nowDateTime)
            elapseSecs.append(nowElapseSec)
            nowLineIndex += 1

        # print(len(dateTimes), len(elapseSecs))

        return dateTimes, elapseSecs

    def getDataLineNowTimeDomainAndElapsedSec(self, lineIndex):
        """ 此方法获取一行log内的当前时间和运算时长 """
        dataLineWords = self.dataLines[lineIndex].split()
        try:
            nowDateTime = datetime.strptime(dataLineWords[4], '%Y-%m-%d_%H:%M:%S')
            domain = int(dataLineWords[-4][:-1])
            elapseSec = float(dataLineWords[-3])
        except ValueError:
            return None, None, None
        return nowDateTime, domain, elapseSec


from base import Screen
import curses
from math import sin, cos, pi
from time import sleep, strftime
from datetime import datetime
from threading import Lock, Thread

class TimeClock(Screen):
    name = '时钟'
    used_pairs = (
        (curses.COLOR_BLACK, curses.COLOR_WHITE),
        (curses.COLOR_YELLOW, -1),
        (curses.COLOR_RED, -1),
        (curses.COLOR_GREEN, -1)
    )
    styles = None
    clock_str = '''
          ***********
      ****    XII    ****
    **  XI            I  **
   *                       *
  * X                    II *
 *                           *
 *                           *
** IX          *         III **
 *                           *
 *                           *
  * VIII                 IV *
   *                       *
    **  VII           V  **
      ****     VI    ****
          ***********
'''
    # version 0
    # r'''
          # ***********
      # ****    XII    ****
    # **  XI           I   **
   # *               /       *
  # *  X         |  /     II  *
 # *             | /           *
 # *             |/            *
# **  IX         *        III  **
 # *              \            *
 # *               \           *
  # * VIII          \     IV  *
   # *               \       *
    # **  VII         \ V  **
      # ****     VI    ****
          # ***********
    # '''
    clock_height = 15
    clock_width = 31
    clock_radio = 15
    clock_hour_rate = 0.4
    clock_min_rate = 0.55
    clock_sec_rate = 0.65
    '''convert_table = [
        (pi/2, '|'), (pi/3, '/'), (pi/6, '-'), (0, '-'),
        (pi/6*11, '-'), (pi/3*5, '\\'), (pi/2*3, '|'),
        (pi/3*4, '/'), (pi/6*7, '-'), (pi, '-'),
        (pi/6*5, '-'), (pi/3*2, '\\')
    ]'''

    end = False
    pad = None
    thread = None

    def __init__(self, scr):
        # if TimeClock.styles is None:
            # curses.init_pair(3, curses.COLOR_BLACK, curses.COLOR_WHITE)
            # curses.init_pair(4, curses.COLOR_YELLOW, -1)
            # curses.init_pair(5, curses.COLOR_RED, -1)
            # curses.init_pair(6, curses.COLOR_GREEN, -1)
            # TimeClock.styles = [
                # curses.color_pair(3), curses.color_pair(4),
                # curses.color_pair(5), curses.color_pair(6)
            # ]
        Screen.__init__(self, scr)
        self.lock = Lock()
        self.alloc_space()
        self.draw_clock()

    def alloc_space(self):
        if self.maxx < 33 or self.maxy < 23:
            self.clock_scr = None
            self.time_scr = self.scr.subwin(1,
                self.maxx, (self.maxy - 1) // 2, 0)
            self.mode = 1
            return
        self.mode = 0
        width = TimeClock.clock_width
        height = TimeClock.clock_height
        clock_x = (self.maxx - width) // 2
        if self.maxy > 30:
            clock_y = (self.maxy - height - 6) // 2
            time_y = clock_y + height + 3
        else:
            clock_y = 3
            time_y = clock_y + height + 2
        self.clock_scr = self.scr.derwin(height, width, clock_y, clock_x)
        self.time_scr = self.scr.derwin(1, self.maxx, time_y, 0)
        # mode:
        # 0: clock painted
        # 1: no clock painting

    def get_clock_pad(self):
        pad = curses.newpad(
            TimeClock.clock_height,
            TimeClock.clock_width)
        border_style, char_style = TimeClock.styles[0:2]
        x = y = 0
        for ch in TimeClock.clock_str[1:-1]:
            if ch == '\n' or ch == '\r':
                x = 0
                y += 1
                continue
            if ch == '*':
                pad.addch(y, x, '■', border_style)
            elif ch != ' ':
                pad.addch(y, x, ch, char_style)
            x += 1
        return pad

    def draw_clock(self):
        if self.end or self.mode == 1:
            return

        # load time
        now = datetime.now()
        hour, minute, second = now.hour, now.minute, now.second
        del now

        # self.clock_scr.erase()
        min_style, hour_style, sec_style = TimeClock.styles[1:]
        '''x = y = 0
        for ch in TimeClock.clock_str[1:-1]:
            if ch == '\n' or ch == '\r':
                x = 0
                y += 1
                continue
            if ch == '*':
                self.clock_scr.addch(y, x, '■', border_style)
            elif ch != ' ':
                self.clock_scr.addch(y, x, ch, char_style)
            x += 1'''
        if self.pad is None:
            self.pad = self.get_clock_pad()
        endy = TimeClock.clock_height - 1
        endx = TimeClock.clock_width - 1
        self.pad.overwrite(self.clock_scr, 0, 0, 0, 0, endy, endx)

        # draw clock hands
        self.draw_hand(second / 5,
            TimeClock.clock_sec_rate, sec_style)
        self.draw_hand(minute / 5 + second / 300,
            TimeClock.clock_min_rate, min_style)
        self.draw_hand(hour + minute / 60,
            TimeClock.clock_hour_rate, hour_style)

        # end drawing
        self.clock_scr.refresh()

    def draw_hand(self, time, rate, style):
        # for style, see 'version 0' comment
        rd = TimeClock.clock_radio
        start_y = rd // 2
        start_x = rd
        # rad, ch = self.__convert_table[time % 12]
        time %= 12
        block = (15 - time) % 12
        rad = block * pi / 6
        if 1.8 <= time <= 4.2 or 7.8 <= time <= 10.2:
            ch = '-'
        elif 2.5 <= block <= 3.5 or 8.5 <= block <= 9.5:
            ch = '|'
        elif 0 <= block <= 3 or 6 <= block <= 9:
            ch = '/'
        else:
            ch = '\\'
        sin_rad = sin(rad)
        cos_rad = cos(rad)
        end_y = int(-sin_rad * rd / 2 * rate) + start_y
        end_x = int(cos_rad * rd * rate) + start_x
        if end_x == start_x:
            if end_y > start_y:
                for y in range(start_y + 1, end_y + 1):
                    self.clock_scr.addch(y, end_x, ch, style)
            else:
                for y in range(start_y - 1, end_y - 1, -1):
                    self.clock_scr.addch(y, end_x, ch, style)
            return
        if end_y == start_y:
            if end_x > start_x:
                for x in range(start_x + 1, end_x + 1):
                    self.clock_scr.addch(end_y, x, ch, style)
            else:
                for x in range(start_x - 1, end_x - 1, -1):
                    self.clock_scr.addch(end_y, x, ch, style)
            return
        if 1.5 <= time <= 4.5 or 7.5 <= time <= 10.5:
            tan = (end_y - start_y) / (end_x - start_x)
            if end_x > start_x:
                for x in range(start_x + 1, end_x + 1):
                    y = tan * (x - start_x) + start_y
                    self.clock_scr.addch(round(y), x, ch, style)
            else:
                for x in range(start_x - 1, end_x - 1, -1):
                    y = tan * (x - start_x) + start_y
                    self.clock_scr.addch(round(y), x, ch, style)
        else:
            tan = (end_x - start_x) / (end_y - start_y)
            if end_y > start_y:
                for y in range(start_y + 1, end_y + 1):
                    x = tan * (y - start_y) + start_x
                    self.clock_scr.addch(y, round(x), ch, style)
            else:
                for y in range(start_y - 1, end_y - 1, -1):
                    x = tan * (y - start_y) + start_x
                    self.clock_scr.addch(y, round(x), ch, style)

    def draw_time(self):
        if self.end:
            return
        msg = strftime('%a %b %d %Y %H:%M:%S %Z')
        x = (self.maxx - len(msg)) // 2
        self.time_scr.erase()
        self.time_scr.addstr(0, x, msg)
        self.time_scr.refresh()

    def on_load(self):
        self.draw_time()
        self.draw_clock()
        self.thread = Thread(target=self.target)
        self.thread.start()

    def target(self):
        while not self.end:
            for i in range(8):
                sleep(0.1)
                if self.end:
                    return
            self.lock.acquire()
            self.draw_time()
            self.draw_clock()
            self.lock.release()

    def on_unload(self):
        self.lock.acquire()
        self.end = True
        self.scr.erase()
        self.scr.refresh()
        self.lock.release()
        if self.thread is not None:
            self.thread.join()
            self.thread = None

    def on_resize(self):
        self.lock.acquire()
        Screen.on_resize(self)
        self.alloc_space()
        self.scr.erase()
        self.scr.refresh()
        self.draw_clock()
        self.draw_time()
        self.lock.release()

export = TimeClock


from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from optimum.pipelines import pipeline
from pathlib import Path

class OptimizedM100Model:
    def __init__(self, model_path, src_lang, tgt_lang):
        model_path = Path(model_path)
        assert model_path.exists(), "Model path does not exist"
        print("start loading the model........")
        self._model = ORTModelForSeq2SeqLM.from_pretrained(model_path)
        print("Model loaded successfully!")
        self._tokenizer = AutoTokenizer.from_pretrained(model_path)
        print("Tokenizer loaded successfully")
        self.pipeline = pipeline(f"translation_{src_lang}_to_{tgt_lang}", model=self._model, tokenizer=self._tokenizer)
        print("Pipeline created successfully")

    def predict_raw(self, X):
        data_to_translate = X.get("data")
        output = self.pipeline(data_to_translate)
        return output

    def health_status(self):
        text_to_translate = {"data": "Hello, my name is Espoir Murhabazi,  I am a Software Engineer from Congo DRC but living in UK"}
        translation = self.predict_raw(text_to_translate)
        assert len(translation) == 1, "health check returning bad translation"
        assert translation[0].get("translation_text") is not None, "health check returning bad translation"
        return translation[0].get("translation_text")


import torch
import torch.nn as nn
from torch import optim
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
import matplotlib.pyplot as plt

train_data = np.load("E:\\quant_research\\train the rank of ten points\\RNN_point\\data\\train_data_10num.npy")
train_aim = np.load("E:\\quant_research\\train the rank of ten points\\RNN_point\\data\\train_label_10num.npy")
train_data = train_data.reshape(train_data.shape[0],10,1)
train_data = train_data.swapaxes(0, 1)
train_data = torch.from_numpy(train_data).type(torch.FloatTensor)
train_aim = torch.from_numpy(train_aim).type(torch.FloatTensor)

test_data = np.load("E:\\quant_research\\train the rank of ten points\\RNN_point\\data\\test_data_10num.npy")
test_aim = np.load("E:\\quant_research\\train the rank of ten points\\RNN_point\\data\\test_label_10num.npy")
test_data = test_data.reshape(test_data.shape[0],10,1)
test_data = test_data.swapaxes(0, 1)
test_data = torch.from_numpy(test_data).type(torch.FloatTensor)
test_aim = torch.from_numpy(test_aim).type(torch.FloatTensor)

class Encoder(nn.Module):

    def __init__(self, input_size, hidden_size, batch_size, bidirectional=True):
        super(Encoder, self).__init__()
        self.hidden_size = hidden_size
        self.input_size = input_size
        self.batch_size = batch_size
        self.bidirectional = bidirectional

        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=False, bidirectional=bidirectional)

    def forward(self, inputs, hidden):
        output, hidden = self.lstm(inputs, hidden)
        return output, hidden

    def init_hidden(self):
        return (torch.zeros(1 + int(self.bidirectional), self.batch_size, self.hidden_size),
                torch.zeros(1 + int(self.bidirectional), self.batch_size, self.hidden_size))   #(num_layers * num_directions, batch, hidden_size)

class AttentionDecoder(nn.Module):

    def __init__(self, hidden_size, output_size, batch_size, vocab_size,seq_len):
        super(AttentionDecoder, self).__init__()
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.batch_size = batch_size
        self.seq_len = seq_len
        self.vocab_size = vocab_size

        self.attn = nn.Linear(hidden_size + output_size + vocab_size, 1)
        self.lstm = nn.LSTM(hidden_size + vocab_size, output_size)
        self.final = nn.Linear(output_size, vocab_size)

    def init_hidden(self):
        return (torch.zeros(1, self.batch_size, self.output_size),
                torch.zeros(1, self.batch_size, self.output_size))

    def forward(self, decoder_hidden, encoder_outputs, input):
        seq = 0
        weights= []
        i = 0
        output = torch.zeros(self.batch_size, self.vocab_size)
        for i in range(len(encoder_outputs)):
            weights.append(self.attn(torch.cat((decoder_hidden[0][:].squeeze(0),encoder_outputs[i],output), dim=1)))

        normalized_weight = F.softmax(torch.cat(weights, 1), 1)
        normalized_weights = normalized_weight

        attn_applied = torch.bmm(normalized_weight.unsqueeze(1),
                                 encoder_outputs.transpose(0,1))
        input_lstm = torch.cat((attn_applied.transpose(0,1)[0], output),
                               dim=1)  # if we are using embedding, use embedding of input here instead

        output_, hidden = self.lstm(input_lstm.unsqueeze(0), decoder_hidden)

        output = self.final(output_[0]) #output 为（vocab_size, output_size）
        #output = self.final2(output)

        # hidden0 = hidden[0].transpose(0, 1).reshape(batch_size, 1, -1).transpose(0, 1)
        # hidden1 = hidden[1].transpose(0, 1).reshape(batch_size, 1, -1).transpose(0, 1)
        # decoder_hidden = (hidden0, hidden1)
        # decoder_hiddens = decoder_hidden

        out = F.softmax(output,1)

        return out

seq_len = 10
input_size = 1
hidden_size = 2
batch_size = train_data.shape[1]
bidirectional = True
output_size = hidden_size * (1 + bidirectional)
vocal_size = 10

input = []
for i in range(10):
    m = np.ones((10000,10))*i
    input.append(m)
input = np.array(input)
input = torch.from_numpy(input).type(torch.FloatTensor)

class pointer_atten(nn.Module):
    def __init__(self):
        super(pointer_atten, self).__init__()
        self.layer1 = Encoder(input_size = input_size,
                              hidden_size = hidden_size,
                              batch_size = batch_size,
                              bidirectional=True)
        self.layer2 = AttentionDecoder(
            hidden_size = hidden_size * (1 + bidirectional),
            output_size = output_size,
            batch_size = batch_size,
            vocab_size = vocal_size,
            seq_len = 1
        )

    def forward(self,x):
        output, hidden = self.layer1.forward(x, self.layer1.init_hidden())
        hidden0 = hidden[0].transpose(0, 1).reshape(batch_size, 1, -1).transpose(0, 1)
        hidden1 = hidden[1].transpose(0, 1).reshape(batch_size, 1, -1).transpose(0, 1)
        decoder_hidden = (hidden0, hidden1)
        encoder_outputs = output
        last_output = self.layer2.forward(decoder_hidden, output, input)

        return last_output

Net = pointer_atten()

learning_rate = 0.05
Loss = nn.MSELoss(reduction='mean')
optimizer = torch.optim.Adam(Net.parameters(), lr=learning_rate)

###########################################
# train
###########################################
loss_list = []
True_list = []
num_epochs = 10000
epoch = 10000
batch = train_aim.detach().numpy().size

Net.load_state_dict(torch.load('E:\\quant_research\\train the rank of ten points\\RNN_point\\net_10num\\net720.pkl'))

for epoch in range(1000):
    train_data = Variable(train_data,requires_grad=True)
    train_aim = Variable(train_aim,requires_grad=True)

    # Forward pass
    outputs = Net(train_data)
    loss = Loss(outputs, train_aim)
    loss_list.append(loss)

    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch) % 10 == 0:
        print ('Epoch [{}/{}], Loss: {:.4f}'
               .format(epoch+1,num_epochs,loss.item()))

        is_not = outputs.detach().numpy() - train_aim.detach().numpy()
        is_not = np.where(is_not < -0.1, 10, is_not)
        is_not = np.where(is_not < 0.1, 1, 0)
        T_pre = np.nansum(is_not)
        True_rate = T_pre / batch
        True_list.append(True_rate)
        print('accuracy of prediction in training data:', True_rate)

    if epoch % 10 ==0:
        torch.save(Net.state_dict(), 'E:\\quant_research\\train the rank of ten points\\\RNN_point\\net_10num\\net{}.pkl'.format(epoch))

loss_array = np.array(loss_list)
true_array = np.array(True_list)
np.save('E:\\quant_research\\train the rank of ten points\\\RNN_point\\loss',loss_array)
np.save('E:\\quant_research\\train the rank of ten points\\\RNN_point\\true',true_array)

loss_array = np.load('E:\\quant_research\\train the rank of ten points\\\RNN_point\\loss.npy',allow_pickle=True)
true_array = np.load('E:\\quant_research\\train the rank of ten points\\\RNN_point\\true.npy')

outputs = Net(train_data)
loss = Loss(outputs, train_aim)
label = np.argmax(outputs.detach().numpy(),axis = 1)
label_aim = np.argmax(train_aim.detach().numpy(),axis = 1)
True_rate = np.sum(label == label_aim) / 10000
print('loss in testing data:%.5f，accuracy of prediction in testing data:%.5f'%(loss,True_rate))

outputs = Net(test_data)
loss = Loss(outputs, test_aim)
label = np.argmax(outputs.detach().numpy(),axis = 1)
label_aim = np.argmax(test_aim.detach().numpy(),axis = 1)
True_rate = np.sum(label == label_aim) / 10000
print('loss in training data:%.5f，accuracy of prediction in training data:%.5f'%(loss,True_rate))


# 10 그래프 이론 - 커리큘럼
# Solved Date: 22.07.17.
import sys
from collections import deque
import heapq

read = sys.stdin.readline

# 위상정렬하고, 더 오래 걸리는 값을 항상 저장해줌
def book_topological_sort(indegrees, graph, costs):
    answer = [cost for cost in costs]
    queue = deque()
    for node, indegree in enumerate(indegrees):
        if indegree == 0:
            queue.append(node)

    while queue:
        node = queue.popleft()
        for next_node in graph[node]:
            answer[next_node] = max(answer[next_node], answer[node] + costs[next_node])
            indegrees[next_node] -= 1
            if indegrees[next_node] ==  0:
                queue.append(next_node)

    return "\n".join([str(x) for x in answer])

# heapq를 통해 정렬하면 항상 최소로 강의 종료시간이 나온다는 점을 이용
def topological_sort(indegrees, graph, costs):
    hq = []
    for node, indegree in enumerate(indegrees):
        if indegree == 0:
            heapq.heappush(hq, (costs[node], node))

    answer = ["" for _ in range(len(indegrees))]
    while hq:
        current_cost, node = heapq.heappop(hq)
        answer[node] = str(current_cost)
        for next_node in graph[node]:
            indegrees[next_node] -= 1
            if indegrees[next_node] == 0:
                heapq.heappush(hq, (current_cost + costs[next_node], next_node))
    return "\n".join(answer)

def main():
    n = int(read().rstrip())
    indegrees = [0 for _ in range(n)]
    graph = [[] for _ in range(n)]
    costs = [0 for _ in range(n)]
    for node in range(n):
        for value_index, value in enumerate((int(x) for x in read().split())):
            if value_index == 0:
                costs[node] = value
                continue

            if value == -1:
                break

            indegrees[node] += 1
            graph[value-1].append(node)
    return print(book_topological_sort(indegrees, graph, costs))

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-

import json
import logging
import traceback
from datetime import datetime, timedelta

import odoo
from odoo import _, models, fields, api
from odoo.api import Environment
from odoo.tools import DEFAULT_SERVER_DATETIME_FORMAT as DATETIME_FORMAT

from ..api import AsyncDB

_logger = logging.getLogger(__name__)

class TaskAbstract(models.AbstractModel):

    _name = "oe.task.abstract"

    task_id = fields.Char(_('task id'))
    task_name = fields.Char(_('task name'))
    task_args = fields.Char(_('task arguments'))
    task_kwargs = fields.Char(_('task kwargs'))

    status = fields.Selection([
        ('PENDING', '待执行'),
        ('STARTED', '执行中'),
        ('SUCCESS', '完成'),
        ('FAILURE','执行失败'),
        ('REVOKED','已撤销'),
        ('RETRY','已重执行'), # RECEIVED
    ], string=_('state'), default='PENDING')

    priority = fields.Integer('priority')
    queue = fields.Char('queue')
    countdown = fields.Integer('countdown', default=0)

class TaskTask(models.Model):

    _name = 'oe.task'
    _description = u'Task'
    _inherit = ['oe.task.abstract']

    def run(self, tasks):
        for task in tasks:
            task_args = json.loads(task['task_args'])
            task_kwargs = json.loads(task['task_kwargs'])
            dbname = task_args.pop(0)
            uid = task_args.pop(0)
            model_name = task_args.pop(0)
            method = task_args.pop(0)
            ids = task_args.pop(0)

            if task['countdown']>0:
                _now = datetime.now()
                if _now < task['create_date'] + timedelta(seconds=task['countdown']):
                    continue

            _context = 'context' in task_kwargs and task_kwargs.pop('context') or {}
            env = Environment(self.env.cr, uid, _context)
            #env = self.with_context().env
            Model = env[model_name]

            try:
                objs = Model.sudo().search([('id', 'in', ids)])
                getattr(env.registry[model_name], method)(objs, *task_args, **task_kwargs)
                env.cr.commit()
            except Exception as exc:
                env.cr.rollback()
                self.env['oe.task.result'].sudo().create({
                    'task_id': task['id'],
                    'task_name': task['task_name'],
                    'task_args': task['task_args'],
                    'task_kwargs': task['task_kwargs'],
                    'status': 'FAILURE',
                    'traceback': '{}'.format(traceback.format_exc()),
                })
            finally:
                #task.unlink()
                self.delete(task)
                self.env.cr.commit()

    def delete(self,task):
        self.env.cr.execute('delete from oe_task where id=%s'%task['id'])

    def _process(self):
        while True:
            import time;time.sleep(2)
            db = odoo.sql_db.db_connect(self.env.cr.dbname)
            tasks = []
            with db.cursor() as cr:
                cr.execute("select * from oe_task where status='PENDING'")
                tasks = cr.dictfetchall()
            self.run(tasks)
            break

    @AsyncDB()
    @api.model
    def get_count(self):
        _logger.info('>>> get_task_count result %s', self.search_count([]))

    @AsyncDB(countdown=10)
    @api.model
    def test_countdown(self):
        _logger.info('>>> It is time to do')

    @AsyncDB()
    @api.model
    def test_trace_log(self):
        d = {}
        v = d['k']


import logging
import typing

from flask import render_template

import api
import api_types
from soda_api import client, LICENSE_DATASET

log = logging.getLogger(__name__)

def license_lookup(license: str) -> str:
    if license:
        try:
            results = client.get(
                LICENSE_DATASET, limit=1, where=f"license like '{license.upper()}'"
            )
            if not results:
                html = (
                    # Disable reason: No need to fix line length for strings
                    "
No vehicle found for this license in public dataset"
                    "

(not all undercover vehicles have available information)
"
                )
            else:
                r = results[0]
                html = render_template("license.html", **r)

        except Exception as err:
            log.exception(err)
            html = f"
Error: {err}"

        log.info("Final HTML:\n{}".format(html))
        return html

    else:
        return ""

def filter_records(
    records: typing.List[api_types.Record],
) -> typing.List[api_types.Record]:
    filtered: typing.List[api_types.Record] = []
    for idx, record in enumerate(records):
        next_idx = idx + 1
        if next_idx == len(records):
            # Reached the end of the list
            break
        next_record = records[next_idx]
        # Override date keys so they're the same for all comparisons
        if {**record, "date": True} != {**next_record, "date": True}:  # type: ignore
            filtered.append(record)
    return filtered

def name_lookup(
    metadata: api_types.DatasetMetadata,
    entities: typing.List[api_types.Entity],
    strict_search: bool = None,
    historical: bool = False,
    show_full_history: bool = False,
    **kwargs,
) -> typing.Tuple[str, bool]:
    if strict_search is None:
        strict_search = True
        for fuzzy_entity in [entity for entity in entities if entity["is_fuzzy"]]:
            if (
                fuzzy_entity["query_param"] in kwargs
                and kwargs[fuzzy_entity["query_param"]] != ""
            ):
                strict_search = False
                break

    try:
        records = api.get_results(metadata, strict_search, historical, **kwargs)
        if not show_full_history and historical:
            records = filter_records(records or [])
        html = (
            api.render_historical_officers(records, metadata)
            if historical
            else api.render_officers(records, metadata)
        )
    except Exception as err:
        log.exception(err)
        html = f"
Error: {err}"

    log.info("Final HTML:\n{}".format(html))
    return html, strict_search


from pymongo import MongoClient

# Create a pymongo client
client = MongoClient("localhost", 27017)

# Get the database instance
db = client["mydb"]

# db collection
pytech = db["PyTech"]

# insert 3 students
records = [
    {
        "student_id": "1007",
        "first_name": "Fred",
        "last_name": "Jones"
    },
    {
        "student_id": "1008",
        "first_name": "Jane",
        "last_name": "Doe"
    },
    {
        "student_id": "1009",
        "first_name": "Tim",
        "last_name": "Barnes"
    }
]

for record in records:
    new_student_Id = pytech.insert_one(record).inserted_id
    print(new_student_Id)

# display all documents in the collection
docs = pytech.find()

for doc in docs:
    print(doc)

# display a single document by student_id
print(pytech.find_one({"student_id": "1008"}))


import os
import pickle
import argparse
import numpy as np
import torch

import utils
import attacks

class TargetedIndexedDataset():
    def __init__(self, dataset, classes):
        self.dataset = dataset
        self.classes = classes

    def __getitem__(self, idx):
        x, y, ii = self.dataset[idx]
        y += 1
        if y >= self.classes: y -= self.classes

        return x, y, ii

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

def get_args():
    parser = argparse.ArgumentParser()
    utils.add_shared_args(parser)

    parser.add_argument('--adv-type', type=str, default='robust-pgd',
                        choices=['robust-pgd', 'diff-aug-pgd'])
    parser.add_argument('--targeted', action='store_true')

    parser.add_argument('--samp-num', type=int, default=1,
                        help='set the number of samples for calculating expectations')

    parser.add_argument('--resume', action='store_true',
                        help='set resume')
    parser.add_argument('--resume-path', type=str, default=None,
                        help='set where to resume the model')

    return parser.parse_args()

def regenerate_def_noise(def_noise, model, criterion, loader, defender, cpu, logger):
    cnt = 0
    for x, y, ii in loader:
        cnt += 1
        logger.info('progress [{}/{}]'.format(cnt, len(loader)) )

        if not cpu: x, y = x.cuda(), y.cuda()
        delta = defender.perturb(model, criterion, x, y)
        def_noise[ii] = delta.cpu().numpy()

def save_checkpoint(save_dir, save_name, model, optim, log, def_noise=None):
    torch.save({
        'model_state_dict': utils.get_model_state(model),
        'optim_state_dict': optim.state_dict(),
        }, os.path.join(save_dir, '{}-model.pkl'.format(save_name)))
    with open(os.path.join(save_dir, '{}-log.pkl'.format(save_name)), 'wb') as f:
        pickle.dump(log, f)
    if def_noise is not None:
        def_noise = (def_noise * 255).round()
        assert (def_noise.max()<=127 and def_noise.min()>=-128)
        def_noise = def_noise.astype(np.int8)
        with open(os.path.join(save_dir, '{}-def-noise.pkl'.format(save_name)), 'wb') as f:
            pickle.dump(def_noise, f)

def main(args, logger):
    ''' init model / optim / loss func '''
    model = utils.get_arch(args.arch, args.dataset)
    optim = utils.get_optim(
        args.optim, model.parameters(),
        lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum)
    criterion = torch.nn.CrossEntropyLoss()

    ''' get Tensor train loader '''
    train_loader = utils.get_indexed_tensor_loader(
        args.dataset, batch_size=args.batch_size, root=args.data_dir, train=True)

    dataset = train_loader.loader.dataset
    ascending = True
    if args.targeted:
        if args.dataset == 'cifar10': classes = 10
        elif args.dataset == 'cifar100': classes = 100
        elif args.dataset == 'imagenet-mini': classes = 100
        else: raise ValueError
        dataset = TargetedIndexedDataset(dataset, classes)
        ascending = False

    train_loader = utils.Loader(dataset, batch_size=args.batch_size, shuffle=False, drop_last=False)

    ''' get train transforms '''
    train_trans = utils.get_transforms(
        args.dataset, train=True, is_tensor=True)

    ''' get (original) test loader '''
    test_loader = utils.get_indexed_loader(
        args.dataset, batch_size=args.batch_size, root=args.data_dir, train=False)

    if args.adv_type == 'robust-pgd':
        defender = attacks.RobustPGDAttacker(
            samp_num     = args.samp_num,
            trans        = train_trans,
            radius       = args.pgd_radius,
            steps        = args.pgd_steps,
            step_size    = args.pgd_step_size,
            random_start = args.pgd_random_start,
            ascending    = ascending,
        )
    elif args.adv_type == 'diff-aug-pgd':
        defender = attacks.DiffAugPGDAttacker(
            samp_num     = args.samp_num,
            trans        = train_trans,
            radius       = args.pgd_radius,
            steps        = args.pgd_steps,
            step_size    = args.pgd_step_size,
            random_start = args.pgd_random_start,
            ascending    = ascending,
        )
    else: raise ValueError

    ''' initialize the defensive noise (for unlearnable examples) '''
    data_nums = len( train_loader.loader.dataset )
    if args.dataset == 'cifar10' or args.dataset == 'cifar100':
        def_noise = np.zeros([data_nums, 3, 32, 32], dtype=np.float16)
    elif args.dataset == 'tiny-imagenet':
        def_noise = np.zeros([data_nums, 3, 64, 64], dtype=np.float16)
    elif args.dataset == 'imagenet-mini':
        def_noise = np.zeros([data_nums, 3, 256, 256], dtype=np.float16)
    else:
        raise NotImplementedError

    start_step = 0
    log = dict()

    if not args.cpu:
        model.cuda()
        criterion = criterion.cuda()

    if args.resume:
        state_dict = torch.load( os.path.join(args.resume_path) )
        model.load_state_dict( state_dict['model_state_dict'] )
        optim.load_state_dict( state_dict['optim_state_dict'] )
        del state_dict

    if args.parallel:
        model = torch.nn.DataParallel(model)

    logger.info('Noise generation started')

    regenerate_def_noise(
        def_noise, model, criterion, train_loader, defender, args.cpu, logger)

    logger.info('Noise generation finished')

    save_checkpoint(args.save_dir, '{}-fin'.format(args.save_name), model, optim, log, def_noise)

    return

if __name__ == '__main__':
    args = get_args()
    logger = utils.generic_init(args)

    logger.info('EXP: robust minimax pgd perturbation')
    try:
        main(args, logger)
    except Exception as e:
        logger.exception('Unexpected exception! %s', e)


from tqdm import tqdm
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

from pyflann import FLANN
from scipy.stats import gaussian_kde
from sklearn.neighbors import KernelDensity

import tool

class Coverage:
    def __init__(self, model, layer_size_dict, hyper=None, **kwargs):
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.model = model
        self.model.to(self.device)
        self.layer_size_dict = layer_size_dict
        self.init_variable(hyper, **kwargs)

    def init_variable(self):
        raise NotImplementedError

    def calculate(self):
        raise NotImplementedError

    def coverage(self):
        raise NotImplementedError

    def save(self):
        raise NotImplementedError

    def load(self):
        raise NotImplementedError

    def build(self, data_loader):
        print('Building is not needed.')

    def assess(self, data_loader):
        for data, *_ in tqdm(data_loader):
            if isinstance(data, tuple):
                data = (data[0].to(self.device), data[1].to(self.device))
            else:
                data = data.to(self.device)
            self.step(data)

    def step(self, data):
        cove_dict = self.calculate(data)
        gain = self.gain(cove_dict)
        if gain is not None:
            self.update(cove_dict, gain)

    def update(self, all_cove_dict, delta=None):
        self.coverage_dict = all_cove_dict
        if delta:
            self.current += delta
        else:
            self.current = self.coverage(all_cove_dict)

    def gain(self, cove_dict_new):
        new_rate = self.coverage(cove_dict_new)
        return new_rate - self.current

class SurpriseCoverage(Coverage):
    def init_variable(self, hyper, min_var, num_class):
        self.name = self.get_name()
        assert self.name in ['LSC', 'DSC', 'MDSC']
        assert hyper is not None
        self.threshold = hyper
        self.min_var = min_var
        self.num_class = num_class
        self.data_count = 0
        self.current = 0
        self.coverage_set = set()
        self.mask_index_dict = {}
        self.mean_dict = {}
        self.var_dict = {}
        self.kde_cache = {}
        self.SA_cache = {}
        for (layer_name, layer_size) in self.layer_size_dict.items():
            self.mask_index_dict[layer_name] = torch.ones(layer_size[0]).type(torch.LongTensor).to(self.device)
            self.mean_dict[layer_name] = torch.zeros(layer_size[0]).to(self.device)
            self.var_dict[layer_name] = torch.zeros(layer_size[0]).to(self.device)

    def get_name(self):
        raise NotImplementedError

    def build(self, data_loader):
        print('Building Mean & Var...')
        for i, (data, label) in enumerate(tqdm(data_loader)):
            # print(data.size())
            if isinstance(data, tuple):
                data = (data[0].to(self.device), data[1].to(self.device))
            else:
                data = data.to(self.device)
            self.set_meam_var(data, label)
        self.set_mask()
        print('Building SA...')
        for i, (data, label) in enumerate(tqdm(data_loader)):
            if isinstance(data, tuple):
                data = (data[0].to(self.device), data[1].to(self.device))
            else:
                data = data.to(self.device)
            label = label.to(self.device)
            self.build_SA(data, label)
        self.to_numpy()
        if self.name == 'LSC':
            self.set_kde()
        if self.name == 'MDSC':
            self.compute_covinv()

    def assess(self, data_loader):
        for i, (data, label) in enumerate(tqdm(data_loader)):
            if isinstance(data, tuple):
                data = (data[0].to(self.device), data[1].to(self.device))
            else:
                data = data.to(self.device)
            label = label.to(self.device)
            self.step(data, label)

    def step(self, data, label):
        cove_set = self.calculate(data, label)
        gain = self.gain(cove_set)
        if gain is not None:
            self.update(cove_set, gain)

    def set_meam_var(self, data, label):
        batch_size = label.size(0)
        layer_output_dict = tool.get_layer_output(self.model, data)
        for (layer_name, layer_output) in layer_output_dict.items():
            self.data_count += batch_size
            self.mean_dict[layer_name] = ((self.data_count - batch_size) * self.mean_dict[layer_name] + layer_output.sum(0)) / self.data_count
            self.var_dict[layer_name] = (self.data_count - batch_size) * self.var_dict[layer_name] / self.data_count \
            + (self.data_count - batch_size) * ((layer_output - self.mean_dict[layer_name]) ** 2).sum(0) / self.data_count ** 2

    def set_mask(self):
        feature_num = 0
        for layer_name in self.mean_dict.keys():
            self.mask_index_dict[layer_name] = (self.var_dict[layer_name] >= self.min_var).nonzero()
            feature_num += self.mask_index_dict[layer_name].size(0)
        print('feature_num: ', feature_num)

    def build_SA(self, data_batch, label_batch):
        SA_batch = []
        batch_size = label_batch.size(0)
        layer_output_dict = tool.get_layer_output(self.model, data_batch)
        for (layer_name, layer_output) in layer_output_dict.items():
            SA_batch.append(layer_output[:, self.mask_index_dict[layer_name]].view(batch_size, -1))
        SA_batch = torch.cat(SA_batch, 1) # [batch_size, num_neuron]
        # print('SA_batch: ', SA_batch.size())
        SA_batch = SA_batch[~torch.any(SA_batch.isnan(), dim=1)]
        SA_batch = SA_batch[~torch.any(SA_batch.isinf(), dim=1)]
        for i, label in enumerate(label_batch):
            if int(label.cpu()) in self.SA_cache.keys():
                self.SA_cache[int(label.cpu())] += [SA_batch[i].detach().cpu().numpy()]
            else:
                self.SA_cache[int(label.cpu())] = [SA_batch[i].detach().cpu().numpy()]

    def to_numpy(self):
        for k in self.SA_cache.keys():
            self.SA_cache[k] = np.stack(self.SA_cache[k], 0)

    def set_kde(self):
        raise NotImplementedError

    def compute_covinv(self):
        raise NotImplementedError

    def calculate(self):
        raise NotImplementedError

    def update(self, cove_set, delta=None):
        self.coverage_set = cove_set
        if delta:
            self.current += delta
        else:
            self.current = self.coverage(self.coverage_set)

    def coverage(self, cove_set):
        return len(cove_set)

    def gain(self, cove_set_new):
        new_rate = self.coverage(cove_set_new)
        return new_rate - self.current

    def save(self, path):
        print('Saving recorded %s in %s...' % (self.name, path))
        state = {
            'coverage_set': list(self.coverage_set),
            'mask_index_dict': self.mask_index_dict,
            'mean_dict': self.mean_dict,
            'var_dict': self.var_dict,
            'SA_cache': self.SA_cache
        }
        torch.save(state, path)

    def load(self, path):
        print('Loading saved %s in %s...' % (self.name, path))
        state = torch.load(path)
        self.coverage_set = set(state['coverage_set'])
        self.mask_index_dict = state['mask_index_dict']
        self.mean_dict = state['mean_dict']
        self.var_dict = state['var_dict']
        self.SA_cache = state['SA_cache']
        loaded_cov = self.coverage(self.coverage_set)
        print('Loaded coverage: %f' % loaded_cov)

class NLC(Coverage):
    def init_variable(self, hyper=None):
        assert hyper is None, 'NLC has no hyper-parameter'
        self.estimator_dict = {}
        self.current = 1
        for (layer_name, layer_size) in self.layer_size_dict.items():
            self.estimator_dict[layer_name] = tool.Estimator(feature_num=layer_size[0])

    def calculate(self, data):
        stat_dict = {}
        layer_output_dict = tool.get_layer_output(self.model, data)
        for (layer_name, layer_output) in layer_output_dict.items():
            info_dict = self.estimator_dict[layer_name].calculate(layer_output.to(self.device))
            stat_dict[layer_name] = (info_dict['Ave'], info_dict['CoVariance'], info_dict['Amount'])
        return stat_dict

    def update(self, stat_dict, gain=None):
        if gain is None:
            for i, layer_name in enumerate(stat_dict.keys()):
                (new_Ave, new_CoVariance, new_Amount) = stat_dict[layer_name]
                self.estimator_dict[layer_name].Ave = new_Ave
                self.estimator_dict[layer_name].CoVariance = new_CoVariance
                self.estimator_dict[layer_name].Amount = new_Amount
            self.current = self.coverage(self.estimator_dict)
        else:
            (delta, layer_to_update) = gain
            for layer_name in layer_to_update:
                (new_Ave, new_CoVariance, new_Amount) = stat_dict[layer_name]
                self.estimator_dict[layer_name].Ave = new_Ave
                self.estimator_dict[layer_name].CoVariance = new_CoVariance
                self.estimator_dict[layer_name].Amount = new_Amount
            self.current += delta

    def coverage(self, stat_dict):
        val = 0
        for i, layer_name in enumerate(stat_dict.keys()):
            # Ave = stat_dict[layer_name].Ave
            CoVariance = stat_dict[layer_name].CoVariance
            # Amount = stat_dict[layer_name].Amount
            val += self.norm(CoVariance)
        return val

    def gain(self, stat_new):
        total = 0
        layer_to_update = []
        for i, layer_name in enumerate(stat_new.keys()):
            (new_Ave, new_CoVar, new_Amt) = stat_new[layer_name]
            value = self.norm(new_CoVar) - self.norm(self.estimator_dict[layer_name].CoVariance)
            if value > 0:
                layer_to_update.append(layer_name)
                total += value
        if total > 0:
            return (total, layer_to_update)
        else:
            return None

    def norm(self, vec, mode='L1', reduction='mean'):
        m = vec.size(0)
        assert mode in ['L1', 'L2']
        assert reduction in ['mean', 'sum']
        if mode == 'L1':
            total = vec.abs().sum()
        elif mode == 'L2':
            total = vec.pow(2).sum().sqrt()
        if reduction == 'mean':
            return total / m
        elif reduction == 'sum':
            return total

    def save(self, path):
        print('Saving recorded NLC in %s...' % path)
        stat_dict = {}
        for layer_name in self.estimator_dict.keys():
            stat_dict[layer_name] = {
                'Ave': self.estimator_dict[layer_name].Ave,
                'CoVariance': self.estimator_dict[layer_name].CoVariance,
                'Amount': self.estimator_dict[layer_name].Amount
            }
        torch.save({'stat': stat_dict}, path)

    def load(self, path):
        print('Loading saved NLC from %s...' % path)
        ckpt = torch.load(path)
        stat_dict = ckpt['stat']
        for layer_name in stat_dict.keys():
            self.estimator_dict[layer_name].Ave = stat_dict[layer_name]['Ave']
            self.estimator_dict[layer_name].CoVariance = stat_dict[layer_name]['CoVariance']
            self.estimator_dict[layer_name].Amount = stat_dict[layer_name]['Amount']

class LSC(SurpriseCoverage):
    def get_name(self):
        return 'LSC'

    def set_kde(self):
        for k in self.SA_cache.keys():
            if self.num_class <= 1:
                self.kde_cache[k] = gaussian_kde(self.SA_cache[k].T)
            else:
                self.kde_cache[k] = KernelDensity(kernel='gaussian', bandwidth=0.2).fit(self.SA_cache[k])
            # The original LSC uses the `gaussian_kde` function, however, we note that this function
            # frequently crashes due to numerical issues, especially for large `num_class`.

    def calculate(self, data_batch, label_batch):
        cove_set = set()
        SA_batch = []
        batch_size = label_batch.size(0)
        layer_output_dict = tool.get_layer_output(self.model, data_batch)
        for (layer_name, layer_output) in layer_output_dict.items():
            SA_batch.append(layer_output[:, self.mask_index_dict[layer_name]].view(batch_size, -1))
        SA_batch = torch.cat(SA_batch, 1).detach().cpu().numpy() # [batch_size, num_neuron]
        for i, label in enumerate(label_batch):
            SA = SA_batch[i]
            # if (np.isnan(SA).any()) or (not np.isinf(SA).any()):
            #     continue
            if self.num_class <= 1:
                lsa = np.asscalar(-self.kde_cache[int(label.cpu())].logpdf(np.expand_dims(SA, 1)))
            else:
                lsa = np.asscalar(-self.kde_cache[int(label.cpu())].score_samples(np.expand_dims(SA, 0)))
            if (not np.isnan(lsa)) and (not np.isinf(lsa)):
                cove_set.add(int(lsa / self.threshold))
        cove_set = self.coverage_set.union(cove_set)
        return cove_set

class DSC(SurpriseCoverage):
    def get_name(self):
        return 'DSC'

    def calculate(self, data_batch, label_batch):
        cove_set = set()
        SA_batch = []
        batch_size = label_batch.size(0)
        layer_output_dict = tool.get_layer_output(self.model, data_batch)
        for (layer_name, layer_output) in layer_output_dict.items():
            SA_batch.append(layer_output[:, self.mask_index_dict[layer_name]].view(batch_size, -1))
        SA_batch = torch.cat(SA_batch, 1).detach().cpu().numpy() # [batch_size, num_neuron]
        for i, label in enumerate(label_batch):
            SA = SA_batch[i]

            # # using numpy
            # dist_a_list = np.linalg.norm(SA - self.SA_cache[int(label.cpu())], axis=1)
            # idx_a = np.argmin(dist_a_list, 0)

            dist_a_list = torch.linalg.norm(
                torch.from_numpy(SA).to(self.device) - torch.from_numpy(self.SA_cache[int(label.cpu())]).to(self.device), dim=1)
            idx_a = torch.argmin(dist_a_list, 0).item()

            (SA_a, dist_a) = (self.SA_cache[int(label.cpu())][idx_a], dist_a_list[idx_a])
            dist_a = dist_a.cpu().numpy()

            dist_b_list = []
            for j in range(self.num_class):
                if ( j != int(label.cpu()) ) and ( j in self.SA_cache.keys() ):
                    # # using numpy
                    # dist_b_list += np.linalg.norm(SA - self.SA_cache[j], axis=1).tolist()
                    dist_b_list += torch.linalg.norm(
                        torch.from_numpy(SA).to(self.device) - torch.from_numpy(self.SA_cache[j]).to(self.device),
                        dim=1).cpu().numpy().tolist()

            dist_b = np.min(dist_b_list)
            dsa = dist_a / dist_b if dist_b > 0 else 1e-6
            if (not np.isnan(dsa)) and (not np.isinf(dsa)):
                cove_set.add(int(dsa / self.threshold))
        cove_set = self.coverage_set.union(cove_set)
        return cove_set

class MDSC(SurpriseCoverage):
    def get_name(self):
        return 'MDSC'

    def set_mask(self):
        feature_num = 0
        for layer_name in self.mean_dict.keys():
            self.mask_index_dict[layer_name] = (self.var_dict[layer_name] >= self.min_var).nonzero()
            feature_num += self.mask_index_dict[layer_name].size(0)
        print('feature_num: ', feature_num)
        self.estimator = tool.Estimator(feature_num=feature_num, num_class=self.num_class)

    def build_SA(self, data_batch, label_batch):
        SA_batch = []
        batch_size = label_batch.size(0)
        layer_output_dict = tool.get_layer_output(self.model, data_batch)
        for (layer_name, layer_output) in layer_output_dict.items():
            SA_batch.append(layer_output[:, self.mask_index_dict[layer_name]].view(batch_size, -1))
        SA_batch = torch.cat(SA_batch, 1) # [batch_size, num_neuron]
        # print('SA_batch: ', SA_batch.size())
        SA_batch = SA_batch[~torch.any(SA_batch.isnan(), dim=1)]
        SA_batch = SA_batch[~torch.any(SA_batch.isinf(), dim=1)]
        stat_dict = self.estimator.calculate(SA_batch, label_batch)
        self.estimator.update(stat_dict)

    def calculate(self, data_batch, label_batch):
        cove_set = set()
        SA_batch = []
        batch_size = label_batch.size(0)
        layer_output_dict = tool.get_layer_output(self.model, data_batch)
        for (layer_name, layer_output) in layer_output_dict.items():
            SA_batch.append(layer_output[:, self.mask_index_dict[layer_name]].view(batch_size, -1))
        SA_batch = torch.cat(SA_batch, 1) # [batch_size, num_neuron]
        mu = self.estimator.Ave[label_batch]
        covar_inv = self.estimator.CoVarianceInv[label_batch]
        mdsa = (
            torch.bmm(torch.bmm((SA_batch - mu).unsqueeze(1), covar_inv),
            (SA_batch - mu).unsqueeze(2))
        ).sqrt()
        # [bs, 1, n] x [bs, n, n] x [bs, n, 1]
        # [bs, 1]
        mdsa = mdsa.view(batch_size, -1)
        mdsa = mdsa[~torch.any(mdsa.isnan(), dim=1)]
        mdsa = mdsa[~torch.any(mdsa.isinf(), dim=1)]
        mdsa = mdsa.view(-1)
        if len(mdsa) > 0:
            mdsa_list = (mdsa / self.threshold).cpu().numpy().tolist()
            mdsa_list = [int(_mdsa) for _mdsa in mdsa_list]
            cove_set = set(mdsa_list)
            cove_set = self.coverage_set.union(cove_set)
        return cove_set

    def compute_covinv(self):
        self.estimator.invert()

    def save(self, path):
        print('Saving recorded %s in %s...' % (self.name, path))
        state = {
            'coverage_set': list(self.coverage_set),
            'mask_index_dict': self.mask_index_dict,
            'mean_dict': self.mean_dict,
            'var_dict': self.var_dict,
            'mu': self.estimator.Ave,
            'covar': self.estimator.CoVariance,
            'amount': self.estimator.Amount
        }
        torch.save(state, path)

    def load(self, path):
        print('Loading saved %s in %s...' % (self.name, path))
        state = torch.load(path)
        self.coverage_set = set(state['coverage_set'])
        self.mask_index_dict = state['mask_index_dict']
        self.mean_dict = state['mean_dict']
        self.var_dict = state['var_dict']
        self.estimator.Ave = state['mu']
        self.estimator.CoVariance = state['covar']
        self.estimator.Amount = state['amount']
        loaded_cov = self.coverage(self.coverage_set)

class NC(Coverage):
    def init_variable(self, hyper):
        assert hyper is not None
        self.threshold = hyper
        self.coverage_dict = {}
        for (layer_name, layer_size) in self.layer_size_dict.items():
            self.coverage_dict[layer_name] = torch.zeros(layer_size[0]).type(torch.BoolTensor).to(self.device)
        self.current = 0

    def calculate(self, data):
        cove_dict = {}
        layer_output_dict = tool.get_layer_output(self.model, data)
        for (layer_name, layer_output) in layer_output_dict.items():
            scaled_output = tool.scale(layer_output)
            mask_index = scaled_output > self.threshold
            is_covered = mask_index.sum(0) > 0
            cove_dict[layer_name] = is_covered | self.coverage_dict[layer_name]
        return cove_dict

    def coverage(self, cove_dict):
        (cove, total) = (0, 0)
        for layer_name in cove_dict.keys():
            is_covered = cove_dict[layer_name]
            cove += is_covered.sum()
            total += len(is_covered)
        return (cove / total).item()

    def save(self, path):
        print('Saving recorded NC in %s...' % path)
        torch.save(self.coverage_dict, path)

    def load(self, path):
        print('Loading saved NC from %s...' % path)
        self.coverage_dict = torch.load(path)

class KMNC(Coverage):
    def init_variable(self, hyper):
        assert hyper is not None
        self.k = int(hyper)
        self.name = 'KMNC'
        self.range_dict = {}
        coverage_multisec_dict = {}
        for (layer_name, layer_size) in self.layer_size_dict.items():
            num_neuron = layer_size[0]
            coverage_multisec_dict[layer_name] = torch.zeros((num_neuron, self.k + 1)).type(torch.BoolTensor).to(self.device)
            self.range_dict[layer_name] = [torch.ones(num_neuron).to(self.device) * 10000, torch.ones(num_neuron).to(self.device) * -10000]
        self.coverage_dict = {
            'multisec': coverage_multisec_dict
        }
        self.current = 0

    def build(self, data_loader):
        print('Building range...')
        for data, *_ in tqdm(data_loader):
            if isinstance(data, tuple):
                data = (data[0].to(self.device), data[1].to(self.device))
            else:
                data = data.to(self.device)
            self.set_range(data)

    def set_range(self, data):
        layer_output_dict = tool.get_layer_output(self.model, data)
        for (layer_name, layer_output) in layer_output_dict.items():
            cur_max, _ = layer_output.max(0)
            cur_min, _ = layer_output.min(0)
            is_less = cur_min < self.range_dict[layer_name][0]
            is_greater = cur_max > self.range_dict[layer_name][1]
            self.range_dict[layer_name][0] = is_less * cur_min + ~is_less * self.range_dict[layer_name][0]
            self.range_dict[layer_name][1] = is_greater * cur_max + ~is_greater * self.range_dict[layer_name][1]

    def calculate(self, data):
        multisec_cove_dict = {}
        lower_cove_dict = {}
        upper_cove_dict = {}
        layer_output_dict = tool.get_layer_output(self.model, data)
        for (layer_name, layer_output) in layer_output_dict.items():
            [l_bound, u_bound] = self.range_dict[layer_name]
            num_neuron = layer_output.size(1)
            multisec_index = (u_bound > l_bound) & (layer_output >= l_bound) & (layer_output <= u_bound)
            multisec_covered = torch.zeros(num_neuron, self.k + 1).type(torch.BoolTensor).to(self.device)
            div_index = u_bound > l_bound
            div = (~div_index) * 1e-6 + div_index * (u_bound - l_bound)
            multisec_output = torch.ceil((layer_output - l_bound) / div * self.k).type(torch.LongTensor).to(self.device) * multisec_index
            # (1, k), index 0 indicates out-of-range output

            index = tuple([torch.LongTensor(list(range(num_neuron))), multisec_output])
            multisec_covered[index] = True
            multisec_cove_dict[layer_name] = multisec_covered | self.coverage_dict['multisec'][layer_name]

        return {
            'multisec': multisec_cove_dict
        }

    def coverage(self, cove_dict):
        multisec_cove_dict = cove_dict['multisec']
        (multisec_cove, multisec_total) = (0, 0)
        for layer_name in multisec_cove_dict.keys():
            multisec_covered = multisec_cove_dict[layer_name]
            num_neuron = multisec_covered.size(0)
            multisec_cove += torch.sum(multisec_covered[:, 1:])
            multisec_total += (num_neuron * self.k)
        multisec_rate = multisec_cove / multisec_total
        return multisec_rate.item()

    def save(self, path):
        print('Saving recorded %s in \%s' % (self.name, path))
        state = {
            'range': self.range_dict,
            'coverage': self.coverage_dict
        }
        torch.save(state, path)

    def load(self, path):
        print('Loading saved %s from %s' % (self.name, path))
        state = torch.load(path)
        self.range_dict = state['range']
        self.coverage_dict = state['coverage']

class SNAC(KMNC):
    def init_variable(self, hyper=None):
        assert hyper is None
        self.name = 'SNAC'
        self.range_dict = {}
        coverage_upper_dict = {}
        for (layer_name, layer_size) in self.layer_size_dict.items():
            num_neuron = layer_size[0]
            coverage_upper_dict[layer_name] = torch.zeros(num_neuron).type(torch.BoolTensor).to(self.device)
            self.range_dict[layer_name] = [torch.ones(num_neuron).to(self.device) * 10000, torch.ones(num_neuron).to(self.device) * -10000]
        self.coverage_dict = {
            'upper': coverage_upper_dict
        }
        self.current = 0

    def calculate(self, data):
        upper_cove_dict = {}
        layer_output_dict = tool.get_layer_output(self.model, data)
        for (layer_name, layer_output) in layer_output_dict.items():
            [l_bound, u_bound] = self.range_dict[layer_name]
            num_neuron = layer_output.size(1)
            upper_covered = (layer_output > u_bound).sum(0) > 0
            upper_cove_dict[layer_name] = upper_covered | self.coverage_dict['upper'][layer_name]

        return {
            'upper': upper_cove_dict
        }

    def coverage(self, cove_dict):
        upper_cove_dict = cove_dict['upper']
        (upper_cove, upper_total) = (0, 0)
        for layer_name in upper_cove_dict.keys():
            upper_covered = upper_cove_dict[layer_name]
            upper_cove += upper_covered.sum()
            upper_total += len(upper_covered)
        upper_rate = upper_cove / upper_total
        return upper_rate.item()

class NBC(KMNC):
    def init_variable(self, hyper=None):
        assert hyper is None
        self.name = 'NBC'
        self.range_dict = {}
        coverage_lower_dict = {}
        coverage_upper_dict = {}
        for (layer_name, layer_size) in self.layer_size_dict.items():
            num_neuron = layer_size[0]
            coverage_lower_dict[layer_name] = torch.zeros(num_neuron).type(torch.BoolTensor).to(self.device)
            coverage_upper_dict[layer_name] = torch.zeros(num_neuron).type(torch.BoolTensor).to(self.device)
            self.range_dict[layer_name] = [torch.ones(num_neuron).to(self.device) * 10000, torch.ones(num_neuron).to(self.device) * -10000]
        self.coverage_dict = {
            'lower': coverage_lower_dict,
            'upper': coverage_upper_dict
        }
        self.current = 0

    def calculate(self, data):
        lower_cove_dict = {}
        upper_cove_dict = {}
        layer_output_dict = tool.get_layer_output(self.model, data)
        for (layer_name, layer_output) in layer_output_dict.items():
            [l_bound, u_bound] = self.range_dict[layer_name]

            lower_covered = (layer_output < l_bound).sum(0) > 0
            upper_covered = (layer_output > u_bound).sum(0) > 0

            lower_cove_dict[layer_name] = lower_covered | self.coverage_dict['lower'][layer_name]
            upper_cove_dict[layer_name] = upper_covered | self.coverage_dict['upper'][layer_name]

        return {
            'lower': lower_cove_dict,
            'upper': upper_cove_dict
        }

    def coverage(self, cove_dict):
        lower_cove_dict = cove_dict['lower']
        upper_cove_dict = cove_dict['upper']

        (lower_cove, lower_total) = (0, 0)
        (upper_cove, upper_total) = (0, 0)
        for layer_name in lower_cove_dict.keys():
            lower_covered = lower_cove_dict[layer_name]
            upper_covered = upper_cove_dict[layer_name]

            lower_cove += lower_covered.sum()
            upper_cove += upper_covered.sum()

            lower_total += len(lower_covered)
            upper_total += len(upper_covered)
        lower_rate = lower_cove / lower_total
        upper_rate = upper_cove / upper_total
        return (lower_rate + upper_rate).item() / 2

class TKNC(Coverage):
    def init_variable(self, hyper):
        assert hyper is not None
        self.k = int(hyper)
        self.coverage_dict = {}
        for (layer_name, layer_size) in self.layer_size_dict.items():
            num_neuron = layer_size[0]
            self.coverage_dict[layer_name] = torch.zeros(num_neuron).type(torch.BoolTensor).to(self.device)
        self.current = 0

    def calculate(self, data):
        cove_dict = {}
        layer_output_dict = tool.get_layer_output(self.model, data)
        for (layer_name, layer_output) in layer_output_dict.items():
            batch_size = layer_output.size(0)
            num_neuron = layer_output.size(1)
            # layer_output: (batch_size, num_neuron)
            _, idx = layer_output.topk(min(self.k, num_neuron), dim=1, largest=True, sorted=False)
            # idx: (batch_size, k)
            covered = torch.zeros(layer_output.size()).to(self.device)
            index = tuple([torch.LongTensor(list(range(batch_size))), idx.transpose(0, 1)])
            covered[index] = 1
            is_covered = covered.sum(0) > 0
            cove_dict[layer_name] = is_covered | self.coverage_dict[layer_name]
        return cove_dict

    def coverage(self, cove_dict):
        (cove, total) = (0, 0)
        for layer_name in cove_dict.keys():
            is_covered = cove_dict[layer_name]
            cove += is_covered.sum()
            total += len(is_covered)
        return (cove / total).item()

    def save(self, path):
        print('Saving recorded TKNC in %s' % path)
        torch.save(self.coverage_dict, path)

    def load(self, path):
        print('Loading saved TKNC from %s' % path)
        self.coverage_dict = torch.load(path)

class TKNP(Coverage):
    def init_variable(self, hyper):
        assert hyper is not None
        self.k = hyper
        layer_pattern = {}
        network_pattern = set()
        self.current = 0
        for (layer_name, layer_size) in self.layer_size_dict.items():
            self.layer_pattern[layer_name] = set()
        self.coverage_dict = {
            'layer_pattern': layer_pattern,
            'network_pattern': network_pattern
        }

    def calculate(self, data):
        layer_pat = {}
        layer_output_dict = tool.get_layer_output(self.model, data)
        topk_idx_list = []
        for (layer_name, layer_output) in layer_output_dict.items():
            num_neuron = layer_output.size(1)
            _, idx = layer_output.topk(min(self.k, num_neuron), dim=1, largest=True, sorted=True)
            # idx: (batch_size, k)
            pat = set([str(s) for s in list(idx[:, ])])
            topk_idx_list.append(idx)
            layer_pat[layer_name] = set.union(pat, self.layer_pattern[layer_name])
        network_topk_idx = torch.cat(topk_idx_list, 1)
        network_pat = set([str(s) for s in list(network_topk_idx[:, ])])
        network_pat = set.union(network_pat, self.network_pattern)
        return {
            'layer_pattern': layer_pat,
            'network_pattern': network_pat
        }

    def coverage(self, cove_dict, mode='network'):
        assert mode in ['network', 'layer']
        if mode == 'network':
            return len(cove_dict['network_pattern'])
        cnt = 0
        if mode == 'layer':
            for layer_name in cove_dict['layer_pattern'].keys():
                cnt += len(cove_dict['layer_pattern'][layer_name])
        return cnt

    def save(self, path):
        print('Saving recorded TKNP in %s' % path)
        torch.save(self.coverage_dict, path)

    def load(self, path):
        print('Loading saved TKNP from %s' % path)
        self.coverage_dict = torch.load(path)

class CC(Coverage):
    '''
    Cluster-based Coverage, i.e., the coverage proposed by TensorFuzz
    '''
    def init_variable(self, hyper):
        assert hyper is not None
        self.threshold = hyper
        self.distant_dict = {}
        self.flann_dict = {}

        for (layer_name, layer_size) in self.layer_size_dict.items():
            self.flann_dict[layer_name] = FLANN()
            self.distant_dict[layer_name] = []

    def update(self, dist_dict, delta=None):
        for layer_name in self.distant_dict.keys():
            self.distant_dict[layer_name] += dist_dict[layer_name]
            self.flann_dict[layer_name].build_index(np.array(self.distant_dict[layer_name]))
        if delta:
            self.current += delta
        else:
            self.current += self.coverage(dist_dict)

    def calculate(self, data):
        layer_output_dict = tool.get_layer_output(self.model, data)
        dist_dict = {}
        for (layer_name, layer_output) in layer_output_dict.items():
            dist_dict[layer_name] = []
            for single_output in layer_output:
                single_output = single_output.cpu().numpy()
                if len(self.distant_dict[layer_name]) > 0:
                    _, approx_distances = self.flann_dict[layer_name].nn_index(np.expand_dims(single_output, 0), num_neighbors=1)
                    exact_distances = [
                        np.sum(np.square(single_output - distant_vec))
                        for distant_vec in self.distant_dict[layer_name]
                    ]
                    buffer_distances = [
                        np.sum(np.square(single_output - buffer_vec))
                        for buffer_vec in dist_dict[layer_name]
                    ]
                    nearest_distance = min(exact_distances + approx_distances.tolist() + buffer_distances)
                    if nearest_distance > self.threshold:
                        dist_dict[layer_name].append(single_output)
                else:
                    self.flann_dict[layer_name].build_index(single_output)
                    self.distant_dict[layer_name].append(single_output)
        return dist_dict

    def coverage(self, dist_dict):
        total = 0
        for layer_name in dist_dict.keys():
            total += len(dist_dict[layer_name])
        return total

    def gain(self, dist_dict):
        increased = self.coverage(dist_dict)
        return increased

    def save(self, path):
        print('Saving recorded CC in %s' % path)
        torch.save(self.distant_dict, path)

    def load(self, path):
        print('Loading saved CC from %s' % path)
        self.distant_dict = torch.load(path)

if __name__ == '__main__':
    pass


import sys
import signal
import socket
import gym
import json
import numpy as np
import cv2
import tqdm

from itertools import product
from time import sleep
from utils.virtual_controller import VirtualKeyboard
from utils.img import ImageCapture
from utils.utils import changeWindowName, kill_process, kill_steam, run_game, kill_game
from utils.wrappers import FrameStack, MaxAndSkipEnv, ScaledFloatFrame, WarpFrame
from speedhack.speedhack import run_speed_hack
from gym import spaces

gym.logger.set_level(40)

RECV_SIZE = 4096

class Isaac(gym.Env):
    def __init__(
        self,
        speed_hack=True,
        use_virtual_controller=True,
        max_timesteps=None,
        max_enemies_to_observe=1,
        max_enemy_projectiles_to_observe=15,
        max_tears_to_observe=5,
        _semaphore=None,
        _alter_window_name=True,
    ) -> None:
        super().__init__()
        if _semaphore:
            _semaphore.acquire(block=True)
        self.max_enemies = max_enemies_to_observe
        self.max_projectiles = max_enemy_projectiles_to_observe
        self.max_tears = max_tears_to_observe
        self.max_steps = max_timesteps
        self.use_virtual_controller = use_virtual_controller
        # generate action dict
        p = list(product(range(0, 5), repeat=2))
        self.action_dict = {
            k: self._tuple_to_vect(v) for k, v in zip(list(range(len(p))), p)
        }

        self.observation_shape = (
            2
            + 2
            + self.max_tears * 4
            + self.max_enemies * 4
            + self.max_projectiles * 4,
        )
        self.observation_space = spaces.Box(
            low=np.finfo(np.float32).min,
            high=np.finfo(np.float32).max,
            shape=self.observation_shape,
            dtype=np.float32,
        )

        # init action space
        self.action_space = spaces.Discrete(
            len(p),
        )

        self.gamestate = {}
        self.previous_gamestate = {}

        exe_path = "C:/program files (x86)/steam/Steam.exe"
        args = [
            "C:/program files (x86)/steam/Steam.exe",
            "-applaunch",
            "250900",
            "--set-stage=1",
            "--set-stage-type=0",
        ]

        kill_steam()
        self.game_process = run_game(exe_path, args)
        self.windown_name = (
            f"{self.game_process.pid}"
            if _alter_window_name
            else f"Binding of Isaac: Repentance"
        )
        if speed_hack:
            run_speed_hack("Binding of Isaac: Repentance", 500)

        # Open connection with game
        self.host, self.port = "127.0.0.1", 6942
        self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.s.setblocking(True)
        self.connected = False
        while not self.connected:
            try:
                self.s.connect((self.host, self.port))
                self.connected = True
                # print("Successful connection with game.", flush=True)
                changeWindowName("Binding of Isaac: Repentance", self.windown_name)
            except:
                # print("Failed to connect to game, retrying...", flush=True)
                sleep(1)
        self.controller = VirtualKeyboard(self.windown_name, wait=0.5)
        self.capturer = ImageCapture(self.windown_name)

        self._current_step = 0

        if _semaphore:
            _semaphore.release()

    @staticmethod
    def _tuple_to_vect(tuple, sizes=[4, 4]):
        vectors = [np.zeros(i) for i in sizes]
        for idx, v in enumerate(tuple):
            if v != 0:
                vectors[idx][v - 1] = 1.0
        return np.concatenate(vectors)

    def compute_reward(self):
        if self.previous_gamestate:
            dmg_taken = (
                self.gamestate["damagesTaken"] - self.previous_gamestate["damagesTaken"]
            )
            mob_dmg_taken = (
                self.gamestate["mobDamagesTaken"]
                - self.previous_gamestate["mobDamagesTaken"]
            )
            mob_dmg_taken = 1.0 if mob_dmg_taken != 0 else 0.0
            mob_killed = (
                self.gamestate["mobKilled"] - self.previous_gamestate["mobKilled"]
            )
            # ======
            is_too_close = 0
            for e in range(len(self.gamestate["enemy_positions"])):
                dist_from_enemy = np.sqrt(
                    (
                        self.gamestate["player_position"][0]
                        - self.gamestate["enemy_positions"][e][0]
                    )
                    ** 2
                    + (
                        self.gamestate["player_position"][1]
                        - self.gamestate["enemy_positions"][e][1]
                    )
                    ** 2
                )
                is_too_close += (
                    1.0 if dist_from_enemy > 270 or dist_from_enemy < 160 else -1.0
                )
            # ======
            has_died = self.gamestate["isDead"]
            cleared = not self.gamestate["isDead"] and self.gamestate["done"]
            return (
                is_too_close * -0.2
                + dmg_taken * -20.0
                + mob_killed * 20.0
                + has_died * -0.0
                + cleared * 20.0
                + mob_dmg_taken * 5.0
            )
        return 0.0

    def observe(self, gamestate):
        observation = np.zeros(self.observation_space.shape)
        ptr = 0
        # player infos
        observation[ptr : ptr + 4] = (
            gamestate["player_position"] + gamestate["player_velocity"]
        )
        observation[ptr : ptr + 4] /= [570, 410, 10, 10]
        ptr = ptr + 4
        # enemies infos
        for i in range(self.max_enemies):
            if i < len(gamestate["enemy_positions"]):
                observation[ptr : ptr + 5] = (
                    gamestate["enemy_positions"][i]
                    + gamestate["enemy_velocities"][i]
                    + [gamestate["enemy_healths"][i]]
                )
                observation[ptr : ptr + 5] /= [570, 410, 10, 10, 300]
            ptr += 4
        # enemies projectiles infos
        for i in range(self.max_projectiles):
            if i < len(gamestate["enemy_projectiles_positions"]):
                observation[ptr : ptr + 4] = (
                    gamestate["enemy_projectiles_positions"][i]
                    + gamestate["enemy_projectiles_velocities"][i]
                )
                observation[ptr : ptr + 4] /= [570, 410, 10, 10]
            ptr += 4
        # tears infos
        for i in range(self.max_tears):
            if i < len(gamestate["tears_positions"]):
                observation[ptr : ptr + 4] = (
                    gamestate["tears_positions"][i] + gamestate["tears_velocities"][i]
                )
                observation[ptr : ptr + 4] /= [570, 410, 10, 10]
            ptr += 4
        return observation

    def reset(self, return_info=False, options=None, seed=None):
        self._current_step = 0
        self.gamestate = {}
        self.previous_gamestate = {}
        self.controller.update(self.controller.get_null_action())
        request = {
            "action": "reset",
            "use_virtual_controller": self.use_virtual_controller,
        }

        self.s.sendall((json.dumps(request) + "\n").encode("utf-8"))
        msg = self.s.recv(RECV_SIZE)
        if msg == b"":
            raise Exception("Connection Lost")
        gamestate = json.loads(msg)

        # observation = self.capturer.capture_frame() # => not reliable
        observation = self.observe(gamestate)
        if not return_info:
            return observation
        else:
            return observation, {}

    def step(self, action):
        """obs, reward, done, infos"""
        act_to_send_low = np.where(self.action_dict[action][0:4] == 1.0)[0]
        act_to_send_high = np.where(self.action_dict[action][4::] == 1.0)[0]
        if not act_to_send_low:
            act_to_send_low = -1
        if not act_to_send_high:
            act_to_send_low = -1
        # Get subsidiary information like reward or done
        ## Send instruction to game
        request = {
            "action": "step",
        }

        if not self.use_virtual_controller:
            request["player_move"] = int(act_to_send_high)
            request["player_fire"] = int(act_to_send_low)

        self.s.sendall((json.dumps(request) + "\n").encode("utf-8"))

        if self.use_virtual_controller:
            self.controller.update(self.action_dict[action])

        ## Receive gamestate from game (if state is terminal for example)
        msg = self.s.recv(RECV_SIZE)
        if msg == b"":
            raise Exception("Connection Lost")
        gamestate = json.loads(msg)

        # Observe the environment
        # observation = self.capturer.capture_frame() # => do not work reliably
        observation = self.observe(gamestate)

        if gamestate["done"]:
            # Send input
            if self.use_virtual_controller:
                self.controller.update(self.controller.get_null_action())

        self.previous_gamestate = self.gamestate
        self.gamestate = gamestate

        reward = self.compute_reward()

        if self.max_steps and self._current_step >= self.max_steps:
            return observation, reward, True, gamestate

        self._current_step += 1
        return observation, reward, gamestate["done"], gamestate

    def close(self):
        print("Closing environement")
        # request = {"action": "close"}
        # self.s.sendall((json.dumps(request) + "\n").encode("utf-8"))
        self.s.close()
        kill_process(self.windown_name)

def signal_handler(sig, frame):
    kill_game()
    sys.exit(0)

# register environment
gym.envs.register(id="isaac-v0", entry_point="isaac_env:Isaac")

if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal_handler)

    # env = Isaac(speed_hack=True)
    env = gym.make("isaac-v0", speed_hack=False)
    obs = env.reset()
    done = False
    count = 0
    # for i in tqdm.trange(0, 1000):
    while True:
        if done:
            env.reset()
        action = env.action_space.sample()
        obs, reward, done, info = env.step(action)
        # print(f"Step {count} => {action}")
        count += 1


# -*- coding: utf-8 -*-

import select
import socket as sk
import logging

import common
import terminal

import commands.client
import commands.command

MAX_LOOP_TIME = 0.1  # s
server_ip = ("localhost", 23456)

class Client(common.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.initState(commands.client.state_setups,
                       commands.command.state_handlers, [0, 0])

    def commandInput(self, term, text):
        """ Called by the terminal when user input. """
        try:
            command = text.split(" ")
            assert len(command) > 0
            command[0] = command[0].strip().lower()

            logging.debug("Command: %s", " ".join(command))
            term.appendHistory(" ".join(command))

            if command[0] == "exit":  # Stops the client, always active.
                self.stop()
            else:
                self.commandHandle(self, command)
        except Exception:
            logging.exception("Exception in handleInput.")

    def connected(self):
        logging.info("Connected to %s.", self.address)

    def disconnected(self):
        logging.info("Disconnected from %s.", self.address)

def main(term):
    socket = sk.socket(sk.AF_INET, sk.SOCK_STREAM)

    logging.info("Connecting to server...")
    try:
        socket.connect(server_ip)
    except (ConnectionAbortedError,
            ConnectionRefusedError,
            ConnectionResetError):
        logging.error("Connection failed.")
    else:
        client = Client(socket, server_ip)
        try:
            client.start()
            term.callback = client.commandInput

            outputs = []

            # Loops when the socket is ready or every MAX_LOOP_TIME seconds.
            while client.is_alive():
                readable, writable, exceptional = select.select(
                    [socket], outputs, [socket], MAX_LOOP_TIME)

                # When receiving something.
                for s in readable:
                    try:
                        data = s.recv(1024)
                    except (ConnectionAbortedError,
                            ConnectionRefusedError,
                            ConnectionResetError):
                        logging.warning("Connection failed.")
                        data = None
                    if data:  # If data is None / b'': the server disconnected.
                        client.data_queue.put(data)  # TODO: public method
                    else:
                        client.stop()
                        if s in outputs:
                            outputs.remove(s)

                # When ready to send something.
                for s in writable:
                    logging.debug("Writable.")

                # When encountering an exception.
                for s in exceptional:
                    client.stop()
                    if s in outputs:
                        outputs.remove(s)
                    logging.warning("Socket exception.")
        except Exception:
            client.stop()
            raise

    socket.close()

if __name__ == "__main__":
    # Setting up "graphics".

    def commandInput(term, text):
        """ Handles terminal input before the client creation. """
        text = text.strip().lower()
        term.appendHistory(text)
        logging.warning("Unhandled Command: %s", text)

    term = terminal.Terminal(commandInput)

    logging.basicConfig(
        level=logging.DEBUG,
        format='[%(levelname)-5s] (%(threadName)-10s) %(message)s',
        handlers=[terminal.TerminalHandler(term)])  # logging.StreamHandler()

    # Main program.
    try:
        main(term)
    except Exception:
        logging.exception("Exception in \"main\".")

    term.stop()


import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score

import category_encoders as ce

import warnings
import sys
import os
import logging

import mlflow
import mlflow.sklearn
import dvc.api

# log warning messages
logging.basicConfig(level=logging.WARN)
logger = logging.getLogger(__name__)

# get data url from DVC
path = 'data/car_evaluation_processed.csv'
repo = '.'
version = 'v1'

data_url = dvc.api.get_url(
    path=path,
    repo=repo,
    rev = version
)

# define metrics for model perfromance evaluation
def model_eval(actual, pred):
    """
    Calculating the accuracy score between actual and predicted values

    Inp: actual, predicated values
    Out: accuracy score
    """
    accuracy = accuracy_score(actual, pred)

    return accuracy

# set an experiment name
mlflow.set_experiment("car-evaluation")

# log parameters
def log_data_params(data_url, data_version, data):
    """
    Logging data parameters to MLflow

    Inp: any
    Out: none
    """
    mlflow.log_param("data_url", data_url)
    mlflow.log_param("data_version", data_version)
    mlflow.log_param("num_rows", data.shape[0])
    mlflow.log_param("num_cols", data.shape[1])

# encode features
def encode_features(X_train, X_test, columns):
    """
    Encoding the features with ordinal encoder

    Inp: X_train, X_test, columns
    Out: X_train_encoded, X_test_encoded
    """

    oe = ce.OrdinalEncoder(cols=columns)
    X_train_encoded = oe.fit_transform(X_train)
    X_test_encoded = oe.transform(X_test)

    return X_train_encoded, X_test_encoded

# execute the training pipeline
if __name__ == "__main__":
    warnings.filterwarnings("ignore")

    # read data from the remote repository
    data = pd.read_csv(data_url, sep=",")

    # initialize mlflow
    with mlflow.start_run():

        # log data parameters
        log_data_params(data_url, version, data)

        # X and y split
        X = data.drop(['class'], axis=1)
        y = data['class']

        # split data into train and test
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

        # create an 'artifacts' directory
        directory = "./artifacts"
        if not os.path.exists(directory):
            os.makedirs(directory)

        # log artifacts: features
        X_cols = pd.DataFrame(list(X_train.columns))
        X_cols.to_csv('artifacts/features.csv', header=False, index=False)
        mlflow.log_artifact('artifacts/features.csv')

        # log artifacts: targets
        y_cols = pd.DataFrame(list(y_train.unique()))
        y_cols.to_csv('artifacts/targets.csv', header=False, index=False)
        mlflow.log_artifact('artifacts/targets.csv')

        # encode the features
        cols = ['buying', 'meant', 'doors', 'persons', 'lug_boot', 'safety']
        X_train, X_test = encode_features(X_train, X_test, cols)

        # set model parameters
        criterion = sys.argv[0] if len(sys.argv) > 1 else "entropy"
        max_depth = float(sys.argv[1]) if len(sys.argv) > 2 else 3

        # training the model
        dtc = DecisionTreeClassifier(criterion=criterion, max_depth=max_depth, random_state=42)
        dtc.fit(X_train, y_train)

        # predicted values
        y_pred = dtc.predict(X_test)

        # model performance using accuracy
        accuracy = model_eval(y_test, y_pred)

        # log the metric
        mlflow.log_metric("accuracy", accuracy)

        # print("Accuracy:, accuracy)

    mlflow.end_run()


import os
from re import L
import sys

import numpy as np
from numpy import asarray
import PIL
from PIL import Image, ImageDraw, ImageFont
from PIL.ExifTags import TAGS

np.set_printoptions(threshold=sys.maxsize)

def addition_decode_algo(stego_path):
    stego_image = Image.open(stego_path, 'r')
    stego_array = np.array(list(stego_image.getdata()))

    if stego_image.mode == 'RGB':
        n = 3
        # steg_image = PIL.Image.new(mode="RGB", size=(200, 200))
    elif stego_image.mode == 'RGBA':
        n = 4
        # steg_image = PIL.Image.new(mode="RGBA", size=(200, 200))

    total_pixels_steg = stego_array.size//n
    width_stego, height_stego = width_c, height_c = stego_image.size
    # print(stego_array.size)
    # print(width_c," ",height_c)
    total_pixels_c = width_c*height_c
    total_pixels_s = total_pixels_steg - total_pixels_c

    rows_s, cols_s = (total_pixels_s, n)
    rows_c, cols_c = (total_pixels_c, n)
    container_array = [[0]*cols_c]*rows_c
    sink_array = [[0]*cols_s]*rows_s

    for p in range(total_pixels_steg):
        for q in range(n):
            container_array[p][q] = stego_array[p][q]

    for p in range(total_pixels_steg, len(stego_array)):
        for q in range(0, n):
            sink_array[p-total_pixels_steg][q] = stego_array[p][q]

    # container_array=np.arange(0,total_pixels_c, 1,np.uint8)
    # container_array=np.reshape(container_array,(1080,1920))
    # container_image=Image.fromarray(container_array)
    # container_image.save(stego_path+"Container.jpg")

    # container_image = Image.fromarray(np.array(container_array).astype(np.uint8))
    # container_image= np.array(container_array).astype(np.uint8)

    # sink_image = Image.fromarray(np.array(sink_array).astype(np.uint8))
    # sink_image = np.array(sink_array).astype(np.uint8)

    # saving the final output
    # as a PNG file

    # container_image = container_image.save(stego_path+"ContainerImage.png")
    # sink_image = sink_image.save(stego_path+"SinkImage.png")

    container_image = Image.fromarray(
        container_array.astype('uint8'), stego_image.mode)
    sink_image = Image.fromarray(sink_array.astype('uint8'), stego_image.mode)

    return ""

def addition_encode_algo(container_path, sink_path):
    container_image = Image.open(container_path, 'r')
    sink_image = Image.open(sink_path, 'r')
    # steg_image = Image.open(container_path, 'r')

    destination_path = "Stego_.jpeg"

    # steg_image = steg_image.save(destination_pat

    # destination_path = os.path.splitext(container_path)[0]+"Encrypted.png"

    width_c, height_c = container_image.size
    print(width_c, " ", height_c)
    width_s, height_s = sink_image.size

    if(width_s > width_c or height_s > width_c):
        print('Error! Sink image too big!')
        return container_path
    # container_array = asarray(container_image)
    # sink_array=asarray(sink_image)
    # steg_array=asarray(steg_image)
    container_array = np.array(list(container_image.getdata()))
    sink_array = np.array(list(sink_image.getdata()))
    steg_array = container_array.copy()
    temp_array = container_array.copy()

    data_c = []  # r,g,b,i,j
    k = 0
    for i in range(height_c):
        temp = []
        for j in range(width_c):
            temp.append(container_array[k])
            k += 1
        data_c.append(temp)

    data_s = []  # r,g,b,i,j
    k = 0
    for i in range(height_s):
        temp = []
        for j in range(width_s):
            temp.append(sink_array[k])
            k += 1
        data_s.append(temp)

    steg_array = data_c.copy()

    n = 5

    # newsize = (width_c*n, height_c*n)
    # steg_array = steg_array.resize(newsize)

    image_resized = []
    for i in range(height_c):
        temp = []
        for j in range(width_c):
            for k in range(n):
                temp.append(data_c[i][j][0:3])
        for k in range(n):
            image_resized.append(temp)

    print(len(image_resized[0]))
    for i in range(len(data_s)):
        for j in range(len(data_s[0])):
            p = int((i+1)*(n-1))
            q = int((j+1)*(n-1))
            image_resized[p][q] = data_s[i][j][0:3]

    # img_res = []
    # for i in range(len(image_resized)):
    #     for j in range(len(image_resized[0])):
    #         img_res.append(image_resized[i][j])

    img_res = np.array(image_resized)

    steg_image = Image.fromarray(
        img_res, "RGB")

    # print(steg_image.size)
    steg_image = steg_image.save(destination_path)
    # print("Image Encoded Successfully")
    # print(width_c, " ", height_c)

    return destination_path