import pygame
import random

def update_direction(keys):
    global direccioX, direccioY

    if keys[pygame.K_w]:
        direccioY = 1
        direccioX = -1

    if keys[pygame.K_a]:
        direccioX = 0
        direccioY = -1

    if keys[pygame.K_d]:
        direccioX = 1
        direccioY = -1

    if keys[pygame.K_s]:
        direccioY = 0
        direccioX = -1

def drawGrid(w, rows, surface):
	for i in range(1,rows+1):
		pygame.draw.line(surface, (0,255,0), (0,i*(w/rows)), (w,i*(w/rows)))
		pygame.draw.line(surface, (0,255,0), (i*(w/rows),0), (i*(w/rows), w))

def le_nek(snake_list, sf, snok):
    for m in snake_list:
        if m != snok:
            pygame.draw.rect(sf, (255,0,0), pygame.Rect(m[0]*(width/rows), m[1]*(width/rows), (width/rows)+1, (width/rows)+1))
        else:
            pygame.draw.rect(sf, (0,0,255), pygame.Rect(m[0]*(width/rows), m[1]*(width/rows), (width/rows)+1, (width/rows)+1))

def redrawWindow(surface, frux, fruy):
    global rows, width, square_xpos, square_ypos, x, y
    surface.fill((0,0,0))
    drawGrid(width,rows, surface)
    a = (255,0,0)

    if fruy == 25:
        fruy = 24
    if frux == 25:
        frux = 24
    pygame.draw.rect(surface, (255,255,255), pygame.Rect(frux*(width/rows), fruy*(width/rows), (width/rows)+1, (width/rows)+1))

def ptsdraw(surface):
    font = pygame.font.SysFont("Arial", 20)
    text = font.render('Points '+str(pts), True, (255,255,255))
    surface.blit(text,(0,0))

def main():
    global width, rows, square_xpos, square_ypos, direccioX, direccioY, x, y, frux, fruy, pts, win, dead
    dead = False
    width = 500
    rows = 25

    hs = []

    x = 0
    y = 0
    direccioX = -1
    direccioY = -1
    pts = 10000

    frux = random.randint(0, rows)
    fruy = random.randint(0, rows)

    snake_List = []
    Length_of_snake = 1

    pygame.init()

    win = pygame.display.set_mode((width,width))

    pygame.display.set_caption("SNAKE")

    square_xpos=20
    square_ypos=20
    dell = 100

    run = True
    startgem = True
    while run:

        if startgem == True:
            my_font = pygame.font.SysFont('times new roman', 90)
            my_fontt = pygame.font.SysFont('times new roman', 15)
            game_start_surface = my_font.render('Controls', True, (255,0,0))
            game_start_surfacee = my_fontt.render('"F" to start  "W""A""S""D" to move  "V" and "B" to turn speed up and down', True, (255,255,255))
            game_start_rect = game_start_surface.get_rect()
            game_start_rectt = game_start_surfacee.get_rect()
            game_start_rect.midtop = (width/2, width/4)
            game_start_rectt.midtop = (width/2, width/2)
            win.fill((0,0,0))
            win.blit(game_start_surface, game_start_rect)
            win.blit(game_start_surfacee, game_start_rectt)
            pygame.display.flip()
            keys = pygame.key.get_pressed()
            if keys[pygame.K_f]:
                startgem = False
                pygame.display.set_caption("SNAKE Points 0")
                x = 0
                y = 0
                direccioX = -1
                direccioY = -1
                pts = 10000
                Length_of_snake = 1

                frux = random.randint(1, rows)
                fruy = random.randint(1, rows)
                snake_List = []
                Length_of_snake = 1

        if startgem != True:
            drawGrid(width, rows, win)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False

        pygame.time.delay(dell)
        pygame.display.set_caption("SNAKE Points "+str(pts))

        # Speed
        keys = pygame.key.get_pressed()
        if keys[pygame.K_v]:
            dell -= 5
        if keys[pygame.K_b]:
            dell += 10

        snake_Head = []
        snake_Head.append(x)
        snake_Head.append(y)
        snake_List.append(snake_Head)

        if len(snake_List) > Length_of_snake:
            del snake_List[0]

        for k in snake_List[:-1]:
            if snake_Head == k:
                dead = True

        update_direction(pygame.key.get_pressed())

        # Movement
        if dead != True:

            if direccioX == 1:
                x += 1
                pts -= 10
            if direccioX == 0:
                x -= 1
                pts -= 10
            if direccioY == 1:
                y -= 1
                pts -= 10
            if direccioY == 0:
                y += 1
                pts -= 10

        # Constraints
        if snake_Head[0] > rows-1:
            dead = True
            x = rows
        elif snake_Head[1] > rows-1:
            dead = True
            y = rows
        if snake_Head[0] < 0:
            dead = True
            x = -1
        elif snake_Head[1] < 0:
            dead = True
            y = -1

        if x == frux and y == fruy:
            frux = random.randint(0, rows-1)
            fruy = random.randint(0, rows-1)
            pts += 1500
            Length_of_snake += 1

        if dead != True:
            redrawWindow(win, frux, fruy)
            le_nek( snake_List, win, snake_Head)
            ptsdraw(win)
            pygame.display.update()

        if dead == True:
            my_font = pygame.font.SysFont('times new roman', 90)
            my_fontt = pygame.font.SysFont('times new roman', 15)
            game_over_surface = my_font.render('YOU DIED', True, (255,0,0))
            game_over_surfacee = my_fontt.render('"F" to restart', True, (255,255,255))
            s_surfacee = my_fontt.render('Score: '+str(pts+10), True, (255,255,255))
            hs.append(pts+10)
            hs_surfacee = my_fontt.render('Highscore: '+str(max(hs)), True, (255,255,255))

            game_over_rect = game_over_surface.get_rect()
            game_over_rectt = game_over_surfacee.get_rect()
            s_rectt = s_surfacee.get_rect()
            hs_rectt = hs_surfacee.get_rect()
            game_over_rect.midtop = (width/2, width/4)
            game_over_rectt.midtop = (width/2, width/1.75)
            hs_rectt.midtop = (width/2, width/2)
            s_rectt.midtop = (width/2, width/1.87)
            win.fill((0,0,0))
            win.blit(game_over_surface, game_over_rect)
            win.blit(game_over_surfacee, game_over_rectt)
            win.blit(s_surfacee, s_rectt)
            win.blit(hs_surfacee, hs_rectt)
            pygame.display.flip()
            keys = pygame.key.get_pressed()
            if keys[pygame.K_f]:
                dead = False
                pygame.display.set_caption("SNAKE Points 0")
                x = 0
                y = 0
                direccioX = -1
                direccioY = -1
                pts = 10000
                Length_of_snake = 1

                frux = random.randint(1, rows)
                fruy = random.randint(1, rows)
                snake_List = []
                Length_of_snake = 1

    pygame.quit()

main()


import logging
import os

#--- Configuration parameters ---

loc_folder = ''  # I/O folder can be set here (Ex:loc_folder = 'C:/Users/USER/PycharmProjects/ProjectA'//). An empty string will look for resources in the root folder.
crypto = ['BTC/TUSD']  # Market pair chosen for trading
target_name = "Return"  # name of the target column
interval = '1s'  # Sample or Candle size of the imported dataset
future_points = 20  # Once the dataset is resampled this value defines how far in the future are we trying to predict (Ex: If resample size = 10 seconds and future_points = 18, the program is predicting at 180 seconds into the future)
resample_co = 15  # Resample size expressed as a number of intervals
time_steps = 30 * future_points  # "future_points" will be used as a base to calculate the sequence length of the LSTM RNN
sequences = 20  # This value represents the number of sequences that are being used in order to train the LSTM.
dense_units = 1  # Dense layer units, this should stay at 1.
dataset_name = 'candles.csv'  # Name of the data set that will be shared and updated constantly by the modeler and trader tools
hyper_loops = 100  # Number of random hyperparameter combinations to try for each combination of indicators that the tuner uses.
comb_loops = 50  # Number of random indicator combinations to try for each combination of hyperparameters that the tuner uses.
desc = 'A'  # Add a description to the output files of the modeler in case we need to save different models
validations = 5  # When modeler_mode is set to 'parallel' the model will calculate the accuracy N times with the same hyperparameters and set the score based on this average in order to get a more consistent value by train/test the model at different times.
#arr_list = ['so_data', 'ema_data', 'rsi_data']  # Once the best indicators have been found by the tuner and in order to use that combination in the modeler is necessary to update this list accordingly with that combination first.
model_reset = 14400  # Time expressed in seconds after which the modeler is reset to find a new model in order to keep the model updated.
modeler_mode = 'solo'  # Set if the modeler will work by itself in "solo" mode or in "parallel" mode with the trader to keep the model updated.
test_size = 0.3  # Fraction of the dataset that will be reserved for testing during modeling
asset_precision = 5
threshold_time = 21600  # A dataset will be built and constantly updated from the websocket seconds-candle stream of prices. This value sets a limit expressed in seconds on the length of this dataset
best_thres = 0  # The trader will use this value to filter weak predictions below a certain absolute value defined by this variable so that orders are only executed when signals are strong enough.
verbose = 0  # Level of description while executing model methods.

# Hyperparameter Tuner Configuration
'''This tool can be set in three different ways depending on the "hyper_mode" value:
hyper_mode = 0 --> Search for the best hyperparameters randomly in a number of combinations of indicators defined by comb_loops.
hyper_mode = 1 --> Search for the best hyperparameters randomly only in the combination of indicators defined in arr_list.
hyper_mode = 2 --> Search for the best hyperparameters randomly on all combinations of indicators.
hyper_mode = 3 --> Uses the default hyperparameter values on all random combinations of indicators.
'''
hyper_mode = 1  # The Hyperparameter finder can try different hyperparameters on all the indicators at the same time or try different Hyperparamenter on all combinations of

# Indicator list for the tuner, modeler, and trader:
arr_list = ['so', 'pvo', 'ema', 'rsi', 'srsi', 'cci', 'psar', 'vwap']  # Once the best indicators have been found by the tuner and in order to use that combination in the modeler is necessary to update this list accordingly with that combination first.
#arr_list = ['srsi', 'cci']
# No indicatorsLayers: 10 Epochs: 90 Batch Size: 15 Dropout Rate: 0.2 Learning Rate: 1e-06 Input Act. Func: tanh Output Act. Func: tanh

# Default hyperparameters for the Tuner and Modeler (If hyper_mode = 2 the tuner will use these):
layers = 64
dropout = 0.2
learning_rate = 0.000001
act_input = 'tanh'
act_output = 'tanh'
epochs = 120
batch_size = 10

# Logger config

file = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'Logger.log')
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s',
                    datefmt='%m-%d %H:%M', filename=file, filemode='a')
console = logging.StreamHandler()
logging.getLogger('').addHandler(console)
logger = logging.getLogger('Log')

# Files config
model_name = f'model-{desc}.h5'
scalers_name = f'scalers-{desc}.pkl'


import pygame, random
pygame.init()

#--------Данные---------
screen_width = 640
screen_height = 512
screen = pygame.display.set_mode([screen_width, screen_height])
pygame.display.set_caption("Сражайся")
pygame.key.set_repeat(2,10)

bg = pygame.image.load('back.jpg')

font = pygame.font.Font(None, 50)
tekst = font.render("ВЫ ПРОИГРАЛИ", 1, [0, 0, 0])

#--------Спрайты игрока---------
player_down = [pygame.image.load('player_down1.png'), pygame.image.load('player_down2.png'), pygame.image.load('player_down3.png')]
player_right = [pygame.image.load('player_right1.png'), pygame.image.load('player_right2.png'), pygame.image.load('player_right3.png')]
player_left = [pygame.image.load('player_left1.png'), pygame.image.load('player_left2.png'), pygame.image.load('player_left3.png')]
player_up = [pygame.image.load('player_up1.png'), pygame.image.load('player_up2.png'), pygame.image.load('player_up3.png')]

player = [player_down, player_right, player_left, player_up]

#--------Первый враг---------
enemy1_down = [pygame.image.load('enemy1_down1.png'), pygame.image.load('enemy1_down2.png'), pygame.image.load('enemy1_down3.png')]
enemy1_right = [pygame.image.load('enemy1_right1.png'), pygame.image.load('enemy1_right2.png'), pygame.image.load('enemy1_right3.png')]
enemy1_left = [pygame.image.load('enemy1_left1.png'), pygame.image.load('enemy1_left2.png'), pygame.image.load('enemy1_left3.png')]
enemy1_up = [pygame.image.load('enemy1_up1.png'), pygame.image.load('enemy1_up2.png'), pygame.image.load('enemy1_up3.png')]

enemy_one = [enemy1_down, enemy1_right, enemy1_left, enemy1_up]

#--------Второй враг---------
enemy2_down = [pygame.image.load('enemy2_down1.png'), pygame.image.load('enemy2_down2.png'), pygame.image.load('enemy2_down3.png')]
enemy2_right = [pygame.image.load('enemy2_right1.png'), pygame.image.load('enemy2_right2.png'), pygame.image.load('enemy2_right3.png')]
enemy2_left = [pygame.image.load('enemy2_left1.png'), pygame.image.load('enemy2_left2.png'), pygame.image.load('enemy2_left3.png')]
enemy2_up = [pygame.image.load('enemy2_up1.png'), pygame.image.load('enemy2_up2.png'), pygame.image.load('enemy2_up3.png')]

enemy_two = [enemy2_down, enemy2_right, enemy2_left, enemy2_up]

#--------Третий враг---------
enemy3_down = [pygame.image.load('enemy3_down1.png'), pygame.image.load('enemy3_down2.png'), pygame.image.load('enemy3_down3.png')]
enemy3_right = [pygame.image.load('enemy3_right1.png'), pygame.image.load('enemy3_right2.png'), pygame.image.load('enemy3_right3.png')]
enemy3_left = [pygame.image.load('enemy3_left1.png'), pygame.image.load('enemy3_left2.png'), pygame.image.load('enemy3_left3.png')]
enemy3_up = [pygame.image.load('enemy3_up1.png'), pygame.image.load('enemy3_up2.png'), pygame.image.load('enemy3_up3.png')]

enemy_three = [enemy3_down, enemy3_right, enemy3_left, enemy3_up]

#--------Спрайты лазера---------
laser_en = [pygame.image.load('laser_en_down.png'), pygame.image.load('laser_en_right.png'), pygame.image.load('laser_en_left.png'), pygame.image.load('laser_en_up.png')]
laser_pl = [pygame.image.load('laser_pl_down.png'), pygame.image.load('laser_pl_right.png'), pygame.image.load('laser_pl_left.png'), pygame.image.load('laser_pl_up.png')]

#--------------Стены------------------
class Wall:
    def __init__(self, x1, y1, x2, y2):  #x1, y1 - верхняя левая точка, x1, y2 - нижняя правая точка
        self.x1 = x1
        self.y1 = y1
        self.x2 = x2
        self.y2 = y2

wall1 = Wall(0, 0, 640, 64)
wall2 = Wall(0, 0, 32, 512)
wall3 = Wall(0, 448, 640, 512)
wall4 = Wall(608, 0, 640, 512)

wall5 = Wall(64, 96, 96, 160)
wall6 = Wall(128, 64, 160, 160)
wall7 = Wall(192, 96, 320, 160)
wall8 = Wall(448, 64, 480, 160)
wall9 = Wall(544, 96, 576, 224)

wall10 = Wall(32, 192, 256, 256)
wall11 = Wall(288, 192, 320, 288)
wall12 = Wall(352, 128, 384, 352)
wall13 = Wall(384, 128, 416, 192)
wall14 = Wall(384, 224, 448, 288)
wall15 = Wall(480, 192, 512, 288)

wall16 = Wall(64, 320, 288, 384)
wall17 = Wall(320, 384, 352, 448)
wall18 = Wall(416, 352, 480, 416)
wall19 = Wall(448, 320, 480, 352)
wall20 = Wall(512, 320, 544, 384)
wall21 = Wall(544, 256, 576, 416)

#---------------Массив со стенами-----------------
walls = [wall1, wall2, wall3, wall4, wall5, wall6, wall7, wall8, wall9, wall10, wall11,
         wall12, wall13, wall14, wall15, wall16, wall17, wall18, wall19, wall20, wall21]

#----------Класс для врагов и игрока------------
class Eye:
    def __init__(self, sprite, direction, x, y, shootSpeed):
        self.sprite = sprite                                #массив изображений
        self.direction = direction                          #текущее направление
        self.shootSpeed = shootSpeed                        #скорость стрельбы (игроком не используется)
        self.shoot=0                                        #счетчик для регулирования частоты выстрелов (игроком не используется)
        self.num = 0                                        #номер спрайта для анимации
        self.image = self.sprite[self.direction][self.num]  #кадр в текущий момент
        self.x = x
        self.y = y
        self.exist = True                                   #если false - исчезает из виду и затем удаляется из массива. Так же отвечает за смерть игрока

    #----------Метод, отвечает за передвижение врага----------
    def move(self):

        #------------Cтолкновения со стенами-----------
        for wall in walls:
            if self.direction == 0 and self.y+32 == wall.y1 and self.x >= wall.x1 and self.x < wall.x2:
                self.turn()
                break
            elif self.direction == 1 and self.x+32 == wall.x1 and self.y >= wall.y1 and self.y < wall.y2:
                self.turn()
                break
            elif self.direction == 2 and self.x == wall.x2 and self.y >= wall.y1 and self.y < wall.y2:
                self.turn()
                break
            elif self.direction == 3 and self.y == wall.y2 and self.x >= wall.x1 and self.x < wall.x2:
                self.turn()
                break
        #----------Если не было - враг двигается в зависимости от направления-----------
        else:
            if self.direction == 0:
                self.y += 16
            elif self.direction == 1:
                self.x += 16
            elif self.direction == 2:
                self.x -= 16
            elif self.direction == 3:
                self.y -= 16
            self.anim()

    #--------Метод, поворачивает врага при столкновении со стеной---------
    def turn(self):
        if self.direction == 0 or self.direction == 3:
            self.direction = random.choice([1,2])
            self.move()
        elif self.direction == 1 or self.direction == 2:
            self.direction = random.choice([0,3])
            self.move()
        self.anim()

    #---------Метод, анимация движения---------
    def anim(self):
        if self.num == 0:
            self.num = 1
        elif self.num == 1:
            self.num = 2
        elif self.num == 2:
            self.num = 1
        self.image = self.sprite[self.direction][self.num]

    #-----------Метод, выстреливает лазером и следит за частотой стрельбы-----------
    def shooting(self):
        global lasers

        self.shoot += 1
        if self.shoot == self.shootSpeed:
            lasers.append(Laser(laser_en, self.x, self.y, self.direction))
            self.shoot = 1

class Laser:
    def __init__(self, sprite, x, y, direction):
        self.sprite = sprite
        self.direction = direction

        if self.direction == 0:  #Рассчитывает координаты, в зависимости от направления
            self.x = x
            self.y = y+32
        if self.direction == 1:
            self.x = x+32
            self.y = y
        if self.direction == 2:
            self.x = x-32
            self.y = y
        if self.direction == 3:
            self.x = x
            self.y = y-32

        self.exist = True  #если false - убирает лазер из виду, и затем удаляет из массива

    #----------Рассчет столкновений и движения-----------
    def move(self):

        check = True

        for wall in walls:

            #------------Проверка столкновений. Лазер уничтожается при столкновении со стеной----------------
            if self.direction == 0:
                if self.y+24 >= wall.y1 and self.y+16 <=wall.y2 and self.x >= wall.x1 and self.x < wall.x2:
                    self.exist = False
                    check = False
                    break

            elif self.direction == 1:
                if self.x+24 >= wall.x1 and self.x+16 <=wall.x2 and self.y >= wall.y1 and self.y < wall.y2:
                    self.exist = False
                    check = False
                    break

            elif self.direction == 2:
                if self.x+8 >= wall.x1 and self.x+16 <=wall.x2 and self.y >= wall.y1 and self.y < wall.y2:
                    self.exist = False
                    check = False
                    break

            elif self.direction == 3:
                if self.y+8 >= wall.y1 and self.y+16 <=wall.y2 and self.x >= wall.x1 and self.x < wall.x2:
                    self.exist = False
                    check = False
                    break

        #-------------Проверка на столкновение лазера врага и игрока-------------
        if abs(laser.x - player.x) <32 and abs(laser.y - player.y) <32 and self.sprite == laser_en:
            self.exist = False
            player.exist = False
            check = False

        #------------Если координаты лазера и врага пересекаются. Так же если лазер принадлежит игроку. Враги не будут убивать друг друга----------------
        for enemy in enemies:
            if abs(self.x - enemy.x) <32 and abs(self.y - enemy.y) <32 and self.sprite == laser_pl:
                self.exist = False
                enemy.exist = False
                check = False
                break

        #---------Если ничего из вышеперечисленного не случилось, лазер двигается----------
        if check:
            if self.direction == 0:
                self.y += 24
            elif self.direction == 1:
                self.x += 24
            elif self.direction == 2:
                self.x -= 24
            elif self.direction == 3:
                self.y -= 24

#----------Создаем трех врагов и игрока. Враги идут в массив--------
enemy1 = Eye(enemy_one, 0, 32, 64, 20)
enemy2 = Eye(enemy_two, 0, 576, 64, 25)
enemy3 = Eye(enemy_three, 2, 576, 416, 15)
player = Eye(player, 1, 32, 416, 0)
enemies = [enemy1, enemy2, enemy3]

lasers = [] #Лазеров пока нет
score = 0

running = True

while running:
    screen.blit(bg, (0, 0))

    #------------Враги, двигаются, стреляют--------
    if enemies:
        for enemy in enemies:
            if enemy.exist:
                screen.blit(enemy.image, (enemy.x, enemy.y))
                enemy.move()
                enemy.shooting()

        #--------Если враг был уничтожен, он удаляется из массива--------
        for enemy in enemies:
            if not enemy.exist:
                enemies.remove(enemy)

    #--------Если массив пуст, игрок победил----------
    else:
        tekst = font.render("ВЫ ПОБЕДИЛИ", 1, [0, 0, 0])
        player.end = True

    #--------Игрок двигается-----------
    screen.blit(player.image, (player.x, player.y))
    player.anim()

    #------------Если игрок был уничтожен, игра проиграна---------
    if not player.exist:
        screen.blit(tekst, (200, 200))
        pygame.display.flip()
        running = False
        pygame.time.wait(2000)

    #----------Если лазеры существуют, они двигаются----------
    if lasers and player.exist:
        for laser in lasers:
            screen.blit(laser.sprite[laser.direction], (laser.x, laser.y))
            laser.move()

        #------------Если лазер столкнулся с чем то, он удаляется из массива----------
        for laser in lasers:
            if laser.exist == False:
                lasers.remove(laser)

    for i in pygame.event.get():
        if i.type == pygame.QUIT:
            running = False

        elif i.type == pygame.KEYDOWN:

            #----------Клавиша ПРОБЕЛ создает лазер игрока----------
            if i.key == pygame.K_SPACE:

                #-----------Проверка. Пока существует лазер игрока, невозможно выстрелить еще раз--------
                for laser in lasers:
                    if laser.sprite == laser_pl and laser.exist:
                        break
                else:
                    lasers.append(Laser(laser_pl, player.x, player.y, player.direction))

            #----------Если игрок сталкивается со стеной, движения не происходит--------------
            if i.key == pygame.K_DOWN:
                player.direction = 0
                for wall in walls:
                    if player.y+32 == wall.y1 and player.x+32 > wall.x1 and player.x < wall.x2:
                        break
                else:
                    player.y +=16

            elif i.key == pygame.K_RIGHT:
                player.direction = 1
                for wall in walls:
                    if player.x+32 == wall.x1 and player.y+32 > wall.y1 and player.y < wall.y2:
                        break
                else:
                    player.x += 16

            elif i.key == pygame.K_LEFT:
                player.direction = 2
                for wall in walls:
                    if player.x == wall.x2 and player.y+32 > wall.y1 and player.y < wall.y2:
                        break
                else:
                    player.x -= 16

            elif i.key == pygame.K_UP:
                player.direction = 3
                for wall in walls:
                    if player.y == wall.y2 and player.x+32 > wall.x1 and player.x < wall.x2:
                        break
                else:
                    player.y -= 16

    #---------Обновить изображение персонажа, чтобы не подлагивал после смены направления-----------
    player.image = player.sprite[player.direction][player.num]

    pygame.display.flip()
    pygame.time.delay(40)

pygame.quit()


import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
from os import path
import scipy.integrate as integrate

class QPendulumEnv(gym.Env):
    metadata = {
        'render.modes' : ['human', 'rgb_array'],
        'video.frames_per_second' : 30
    }

    def __init__(self):
        self.max_speed=100
        self.max_torque=0.5

        self.viewer = None
        self.g = 10.
        self.m = 0.15
        self.l = 0.5
        self.mu = 0.05
        self.dt = 0.02
        high = np.array([1,1, np.inf])
        self.action_space = spaces.Box(low=-self.max_torque, high=self.max_torque, shape=(1,), dtype=np.float32)
        self.observation_space = spaces.Box(low=-high, high=high, dtype=np.float32)

        self.seed()

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def x_next(self, x,u):
        x_theta = x[0] + x[1]*self.dt
        x_theta_dot = x[1] + (self.g/self.l*np.sin(x[0])-self.mu/(self.m*self.l**2)*x[1]+1/(self.m*self.l**2)*(u))*self.dt
        return np.array([x_theta, x_theta_dot])

    def step(self,u):
        th, thdot = self.state # th := theta

        u = np.clip(u, -self.max_torque, self.max_torque)[0]
        #u = u[0]
        d_u = u - np.copy(self.last_u)
        self.last_u = u # for rendering
        th_normalized = angle_normalize(th)
        #print(th_normalized)
        costs = th_normalized**2 + 0.1*thdot**2 + 0.001*u**2 #+ 1*d_u**2

        #newthdot = np.clip(newthdot, -self.max_speed, self.max_speed) #pylint: disable=E1111
        self.state = self.x_next(self.state, u)
        return self._get_obs(), -costs, False, {}

    def reset(self):
        #sign = [-1, 1]
        #theta = np.pi - self.np_random.uniform(low=0.1, high=1./2.*np.pi)
        #theta = np.random.uniform(low=np.pi/3, high=2.*np.pi-np.pi/3)
        #theta = self.np_random.choice(sign)*theta
        #theta_dot = self.np_random.uniform(-0.5, 0.5)
        #self.state = np.array([theta, theta_dot])
        high = np.array([np.pi, 3])
        self.state = self.np_random.uniform(low=-high, high=high)
        self.last_u = 0
        return self._get_obs()

    def _get_obs(self):
        theta, thetadot = self.state
        return np.array([np.cos(theta), np.sin(theta), thetadot])

    def render(self, mode='human'):

        if self.viewer is None:
            from gym.envs.classic_control import rendering
            self.viewer = rendering.Viewer(500,500)
            self.viewer.set_bounds(-2.2,2.2,-2.2,2.2)
            rod = rendering.make_capsule(1, .2)
            rod.set_color(.8, .3, .3)
            self.pole_transform = rendering.Transform()
            rod.add_attr(self.pole_transform)
            self.viewer.add_geom(rod)
            axle = rendering.make_circle(.05)
            axle.set_color(0,0,0)
            self.viewer.add_geom(axle)
            #fname = path.join(path.dirname(__file__), "assets/clockwise.png")
            #self.img = rendering.Image(fname, 1., 1.)
            self.imgtrans = rendering.Transform()
            #self.img.add_attr(self.imgtrans)

        #self.viewer.add_onetime(self.img)
        self.pole_transform.set_rotation(self.state[0] + np.pi/2)
        if self.last_u:
            self.imgtrans.scale = (-self.last_u/2, np.abs(self.last_u)/2)

        return self.viewer.render(return_rgb_array = mode=='rgb_array')

    def close(self):
        if self.viewer:
            self.viewer.close()
            self.viewer = None

def angle_normalize(x):
    return (((x+np.pi) % (2*np.pi)) - np.pi)


import numpy as np
import matplotlib.pyplot as plt

def logtanh(x):
    return -np.log(np.tanh(x / 2))

def plot_large_x():
    import matplotlib.pyplot as plt
    x_series = np.arange(0, 35, 0.0001)
    exact_vals = logtanh(x_series)
    t = np.exp(-x_series)

    approx_vals = 2 * t
    rel_error = np.abs(exact_vals - approx_vals) / exact_vals
    plt.semilogy(x_series, rel_error)

    approx_vals = 2 * (t + t ** 3 / 3)
    rel_error = np.abs(exact_vals - approx_vals) / exact_vals
    plt.semilogy(x_series, rel_error)

    approx_vals = 2 * (t + t ** 3 / 3 + t ** 5 / 5)
    rel_error = np.abs(exact_vals - approx_vals) / exact_vals
    plt.semilogy(x_series, rel_error)
    plt.legend(['First-order', 'Third-order', 'Fifth-order'])
    plt.grid()
    plt.show()

def plot_small_x():
    logx = np.arange(-13, 0, 0.001)
    x_series = np.exp(logx)
    exact_vals = logtanh(x_series)

    approx_vals = -np.log(x_series / 2)
    rel_error = np.abs(exact_vals - approx_vals) / exact_vals
    plt.semilogy(x_series, rel_error)

    approx_vals = -np.log(x_series / 2) + x_series ** 2 / 12
    rel_error = np.abs(exact_vals - approx_vals) / exact_vals
    plt.semilogy(x_series, rel_error)
    plt.xscale('log')
    plt.legend(['First-order', 'Second-order'])
    plt.grid()
    plt.show()

if __name__ == '__main__':
    plot_small_x()
    plot_large_x()


import ply.lex as lex
import ply.yacc as yacc
# This is heavily based on Michael Fogelman's solution
# (https://www.michaelfogleman.com/aoc18/#20)
# and only refactored so I understood it better

directions = {
    'N': (0, -1),
    'S': (0, 1),
    'E': (1, 0),
    'W': (-1, 0)
    }

# lexer rules
tokens = ['LPAREN', 'RPAREN', 'OR', 'DIR']

t_LPAREN = r'\('
t_RPAREN = r'\)'
t_OR = r'\|'
t_DIR = r'[NESW]+'

# parser rules
def p_pattern(t):
    'pattern : concat'
    t[0] = t[1]

def p_pattern_or(t):
    'pattern : pattern OR concat'
    t[0] = ('or', t[1], t[3])

def p_concat(t):
    'concat : term'
    t[0] = t[1]

def p_concat_term(t):
    'concat : concat term'
    t[0] = ('concat', t[1], t[2])

def p_term(t):
    'term : DIR'
    t[0] = ('dir', t[1])

def p_term_empty(t):
    'term :'
    t[0] = ('dir', '')

def p_term_paren(t):
    'term : LPAREN pattern RPAREN'
    t[0] = t[2]

def parse(text):
    lexer = lex.lex(errorlog=lex.NullLogger())
    parser = yacc.yacc()
    text = text.strip().replace('^', '').replace('$', '')
    return parser.parse(text, lexer=lexer)

def visit(node, doors, visited):
    if node[0] == 'concat':
        visited = visit(node[1], doors, visited)
        visited = visit(node[2], doors, visited)
        return visited
    elif node[0] == 'or':
        path_a = visit(node[1], doors, visited)
        path_b = visit(node[2], doors, visited)
        return list(set(path_a) | set(path_b))
    else:
        for d in node[1]:
            dx, dy = directions[d]
            doors |= set((x, y, x + dx, y + dy) for x, y in visited)
            doors |= set((x + dx, y + dy, x, y) for x, y in visited)
            visited = [(x + dx, y + dy) for x, y in visited]
        return visited

def locate_doors(text):
    parse_tree = parse(text)
    doors = set()
    visit(parse_tree, doors, [(0, 0)])
    return doors

def search(doors):
    distances = {}
    queue = [(0, 0, 0)]
    while queue:
        x, y, d = queue.pop()
        if (x, y) in distances and distances.get((x, y)) <= d:
            continue
        distances[(x, y)] = d
        for dx, dy in directions.values():
            if (x, y, x + dx, y + dy) in doors:
                queue.append((x + dx, y + dy, d + 1))
    return distances

with open("input.txt", 'r') as myfile:
    input = myfile.readline()
doors = locate_doors(input)
distances = search(doors)
print("Part 1: ", max(distances.values()))
print("Part 2: ", sum(x >= 1000 for x in distances.values()))


#encoding
import xml.etree.ElementTree as ET
employees = ET.Element('Employees')
emp1 = ET.SubElement(employees,'Employee')
id1 = ET.SubElement(emp1,'id')
id1.text = '1'
name1 = ET.SubElement(emp1,'name')
name1.text = 'Mary'
# id2 = ET.SubElement(emp2,'id')
# id2.text = '2'
# name2 = ET.SubElement(emp2,'name')
# name2.text = 'John'
print(ET.tostring(employees))
et = ET.ElementTree(employees)
with open(r'c:\Users\user\PycharmProjects\pythonProject\class7\aa.xml',mode = 'wb') as f:
    et.write(f,encoding = 'UTF-8')


import PIL.Image
import tensorflow as tf
import numpy as np
import pandas as pd
from tensorflow import keras
from keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
from keras.datasets import cifar10
from keras.layers import Dense, Dropout, Conv2D, MaxPool2D, Flatten, Input
from keras.models import Model
from PIL import Image, ImageDraw

def draw_rectangle(draw, p1, p2):
    draw.line([(p1[0], p1[1]),
               (p1[0], p2[1]),
               (p2[0], p2[1]),
               (p2[0], p1[1]),
               (p1[0], p1[1])], "yellow", 5)
    return img

img = Image.open(r'..\data\bild1.jpg')
img = img.resize((round(img.size[0]/10), round(img.size[1]/10)), resample=Image.Resampling.LANCZOS)
# img.rotate(90, expand=True)

#draw = ImageDraw.Draw(img)
draw = ImageDraw.Draw(img)
#img = draw_rectangle(draw, [50, 50], [150, 150])

# plt.imshow(img)
# plt.show()

# load fitted model from "car_detector"
model = tf.keras.models.load_model(r"..\models\cnn_car_detector")

# sliding window over image to find car
# cropping
window_size = 100
step_size = 20
nsteps_x = int(img.size[0] / step_size) + 1
nsteps_y = int(img.size[1] / step_size) + 1

y_data = []
for i in range(nsteps_x):
    for j in range(nsteps_y):
        img_cropped = img.crop((i * step_size, j * step_size,
                                i * step_size + window_size, j * step_size + window_size))\
                                .resize((32, 32), resample=Image.Resampling.BICUBIC)

        data = np.asarray(img_cropped).astype(np.float32) / 255.
        data = data.reshape((1, 32, 32, 3))
        y_data.append(model.predict(data)[0][0])

assert len(y_data) == nsteps_x * nsteps_y, "length of prediction array is not as expected"
y_data = np.reshape(y_data, (nsteps_x, nsteps_y))

# (i_opt, j_opt) = np.unravel_index(y_data.argmax(), y_data.shape)
# point_corner_opt = [i_opt * step_size, j_opt * step_size]
# point_corner_opt2 = [i_opt * step_size + window_size, j_opt * step_size + window_size]
# draw = ImageDraw.Draw(img)
# img_signed = draw_rectangle(draw, point_corner_opt, point_corner_opt2)

points = []
points2 = []
img_signed = img.copy()
for ix in range(y_data.shape[0]):
    for iy in range(y_data.shape[1]):
        if y_data[ix, iy] >= 0.4:
            points.append([ix * step_size, iy*step_size])
            points2.append([ix * step_size + window_size, iy*step_size + window_size])
            draw = ImageDraw.Draw(img_signed)
            img_signed = draw_rectangle(draw, points[-1], points2[-1])

plt.imshow(img_signed)
plt.show()

print('finished')


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jun 19 10:27:05 2019

@author: Kai Sandbrink

"""

from sklearn.metrics import r2_score, mean_squared_error
from sklearn.model_selection import train_test_split
from scipy.optimize import minimize
import numpy as np
import pandas as pd
import pickle, os, argparse
import multiprocessing as mp
import h5py
from sklearn.metrics import roc_auc_score
from sklearn.preprocessing import label_binarize
from sklearn.svm import LinearSVC
from sklearn.multiclass import OneVsRestClassifier
from sklearn.linear_model import LinearRegression, Ridge
import time
import resource
import os

# GLOBAL PARS
t_stride = 2
ntime=320
metrics = ['RMSE', 'r2', 'PCC']
nmetrics = len(metrics)
#datafolder = 'data/'

# %% UTILITY FUNCTIONS

def compute_metrics(y, pred):
    """ Computes RMSE, R2, and correlation performance metrics given true and predicted values

    Parameters
    ----------
    y : np.array of floats, unit activations
    pred : np.array of floats, corresponding predicted values

    Returns
    -------
    [rmse, r2, cor] : prediction metrics
    """
    try:
        rmse = np.sqrt(mean_squared_error(y.flatten(), pred.flatten()))
    except:
        print("MSE Error")
        print(pred)
        print(y)
    r2 = r2_score(y.flatten(), pred.flatten())
    cor = np.corrcoef(y.flatten(), pred.flatten())[0,1]

    #print(r2)
    assert r2 <=1 or np.isnan(r2), 'illegal r2 score!!!! %s %s %s' %(r2, y[:10], pred[:10])

    return [rmse, r2, cor]

def compute_dist_metric(y, pred):
    """ Computes RMSE, R2, and correlation performance metrics given true and predicted values

    Parameters
    ----------
    y : np.array of floats [nr_samples, n_features], unit activations
    pred : np.array of floats [nr_samples, n_features], corresponding predicted values

    Returns
    -------
    dist: int, mean distance between y and pred
    """

    dist = np.linalg.norm(y - pred, axis=1)
    mean_dist = dist.mean(axis=0)
    return mean_dist

def lstring(ilayer):
    ''' Get the formatted string depending on layer

    Arguments
    ---------
    ilayer : int, layer idx (-1 is spindles)

    Returns
    -------
    layer : str
    '''
    if ilayer==-1:
        layer='data'
    else:
        layer = 'l%d' %ilayer
    return layer

def get_binidx(theta, nbins=8):
    binidx = int(( theta + np.pi ) / (2*np.pi/nbins))
    return binidx

def get_centers(fmapntime, ilayer, model, ntime = 320):
    """ Computes the temporal centers of the convolutional centers

    Parameters
    ----------
    fmapntime : int, length of temporal axis in given feature map
    ilayer : int, current layer of network, spindles = -1
    model : dict -> Config, information of model
    ntime : int, starting width

    Returns
    -------
    centers : list of ints, indices of temporal centers
    """

    #s_stride = model['s_stride']
    t_stride = model['t_stride']
    mtype = model['type']

    centers = np.arange(ntime)

    #if spindles skip
    #if(ilayer == -1):
    #    ilayer = 0

    for i in np.arange(0, ilayer + 1): #excludes spindles
        if(mtype == 'S' and i >= 4):
            centers = np.array([centers[i*t_stride] for i in range(int(np.ceil(len(centers)/t_stride)))])
        elif(mtype == 'ST'):
            centers = np.array([centers[i*t_stride] for i in range(int(np.ceil(len(centers)/t_stride)))])

    assert len(centers) == fmapntime, "Time dimensions mismatch!!!! centers: %d fmap: %d" %(len(centers), fmapntime)
    return centers

def read_layer_reps(ilayer, runinfo, model):
    '''Retrieve generated representations for the queried layer, adapted from Pranav 2020-12-01.

    Parameters
    ----------
    ilayer : int, index of current layer, spindles = -1
    runinfo : RunInfo, contains run configuration
    model : dict, contains model configuration

    Returns
    -------
    layer : array of floats [n_samples, n_features, n_time] or [n_samples, n_features, n_time, n_channels],
        storing current layer activation
    '''

    layer = lstring(ilayer)

    if ilayer == -1:
        lo = pickle.load(open(os.path.join(runinfo.datafolder(model), layer + '.pkl'), 'rb'))
    else:
        try:
            lo = pickle.load(open(os.path.join(runinfo.datafolder(model), layer + '.pkl'), 'rb'))
        except:
            h5filepath = os.path.join(runinfo.datafolder(model), layer + '.hdf5')

            reps_layer = h5py.File(h5filepath, 'r')
            num_datasets = len(list(reps_layer))
            ds_shape = list(reps_layer.get('0').shape)

            batch_size = ds_shape[0]
            ds_shape[0] = batch_size*num_datasets
            lo = np.zeros((ds_shape))
            for i in range(num_datasets):
                lo[batch_size*i : batch_size*(i+1)] = reps_layer.get(str(i))[()]

    if(model['type'] == 'LSTM' and ilayer == model['nlayers'] - 1):
        lo = lo.swapaxes(-2,-1)
        print("lo axes swapped. layer representations shape %s " %str(lo.shape))

    if(runinfo.verbose):
        print("read layer represenations. shape: ", lo.shape)

    return lo

# %% DIRECTION TUNING UTILITY FCTS

def unit_vector(vector):
    """ Returns the unit vector of the vector.  """
    return vector / np.linalg.norm(vector)

def angle_between(v1, v2):
    """ Returns the angle in radians between vectors 'v1' and 'v2'::

            >>> angle_between((1, 0, 0), (0, 1, 0))
            1.5707963267948966
            >>> angle_between((1, 0, 0), (1, 0, 0))
            0.0
            >>> angle_between((1, 0, 0), (-1, 0, 0))
            3.141592653589793

    https://stackoverflow.com/questions/2827393/angles-between-two-n-dimensional-vectors-in-python

    Arguments
    ---------
    v1 : np.array [2,]
    v2 : np.array [2,]

    Returns
    -------
    float, angle between v1 and v2
    """

    v1_u = unit_vector(v1)
    v2_u = unit_vector(v2)
    return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))
        #clip solves rounding errors

def angle_xaxis(v):
    """  Returns the angle between a vector and the x-axis ([1,0] vector) according to the rule
        Positive angle => positive y, negative angle => negative y
        Returns nan if v is 0

    Arguments
    ---------
    v : np.array [2,]

    Returns
    -------
    float, angle between v and the [1,0] vector
    """

    if not np.any(v):
        return np.nan
    elif v[1] >= 0:
        return angle_between(v, [1, 0])
    else:
        return - angle_between(v, [1,0])

def get_polar(X):
    ''' Converts Cartesian coordinates to polar

    Parameters
    ----------
    X : np.array [number of samples or timepoints, 2], Cartesian coordinates

    Returns
    -------
    Xp : np.array [number of samples or timepoints, 2], Polar coordinates
    '''

    Xp = np.zeros(X.shape)
    Xp[:,0] = np.apply_along_axis(np.linalg.norm, 1, X)
    Xp[:,1] = np.apply_along_axis(angle_xaxis, 1, X)
    return Xp

# %% DIRECTION AND VELOCITY TUNING FUNCTIONS

def linreg(X_train, X_test, Y_train, Y_test):
    ''' Performs linear regression with train and test set

    Parameters
    ----------
    X_train : np.array [number of samples,]
    X_test : np.array [number of samples,]
    Y_train : np.array [number of samples,]
    Y_test : np.array [number of samples,]

    Returns
    -------
    trainmets : list of floats [3], metrics for training fit specified in compute_metrics function
    testmets : list of floats [3], test metrics specified in compute_metrics function
    '''

    if(np.any(np.isnan(X_train))):
        print("Warning: Contains nans")
    c = np.linalg.lstsq(X_train, Y_train, rcond=None)[0]
    trainmets = np.concatenate((compute_metrics(Y_train, X_train @ c), c, np.zeros(3 - len(c))))
    testmets = np.concatenate((compute_metrics(Y_test, X_test @ c), c, np.zeros(3 - len(c))))

    #print("LinReg Fit Score: " + str(trainmets[:3]))
    #print("LinReg Test Score: " + str(testmets[:3]))
    #print("Coefficients: " + str(c[:3]))
    return trainmets, testmets

def feature_set(Xcols_train, Xcols_test, Y_train, Y_test):
    ''' Combines different features together and excludes invalid values

    Arguments
    ---------
    Xcols_train : tuple of np.arrays of shape [nr of samples, k features]
    Xcols_test : tuple of np.arrays
    Y_train : np.array
    Y_test : np.array

    Returns
    -------
    X_train : np.array [nr of samples, nr of features]
    X_test : np.array [nr of samples, nr of features]
    Y_train : np.array [nr of samples,]
    Y_test : np.array [nr of samples,]

    '''

    X_train = np.column_stack(Xcols_train).squeeze()
    X_test = np.column_stack(Xcols_test).squeeze()
    assert len(X_train) == len(Y_train), "Dimensions mismatch!!! X: %s , Y: %s" %(str(X_train.shape), str(Y_train.shape))
    if len(X_train.shape) > 1:
        nna = ~np.any(np.isnan(X_train), axis=1)
    else:
        nna = ~np.isnan(X_train)
    X_train = X_train[nna]
    Y_train = Y_train[nna]

    if len(X_test.shape) > 1:
        nna = ~np.any(np.isnan(X_test), axis=1)
    else:
        nna = ~np.isnan(X_test)
    X_test = X_test[nna]
    Y_test = Y_test[nna]

    return X_train, X_test, Y_train, Y_test

# %% DIRECTION AND VELOCITY TUNING
def tune_row_vel(X_train, X_test, Y_train, Y_test, row, isPolar = True):
    ''' Fit tuning curves for a single neuron / unit ("row") for velocity inputs (and similar kinematic features) for four different kinds of models
        1. simple linear regression over input features of X
        2. Directional tuning curves
        3. Velocity / Speed Dependence
        4. Direction x Velocity tuning curve

    Arguments
    ---------
    X : np.array [nr of samples, 2]
    Y : np.array [nr of samples]
    row : int, row index
    isPolar : bool, is result already boolean?

    Returns
    -------
    tuple of
        row : int, row index
        rowtraineval : np.array [4, 6] for four different model types. Cols 0-2: Metrics from compute_metrics, Cols 3-5: Linear Regression coeffs
        rowtesteval : np.array [4, 6] for four different model types. Cols 0-2: Metrics from compute_metrics, Cols 3-5: Linear regression coeffs
    '''

    rowtraineval = np.zeros((4,6))
    rowtesteval = np.zeros((4,6))

    ##BASELINE LINREG MODEL
    Xtform_train, Xtform_test, Ytform_train, Ytform_test = feature_set(
            (X_train, np.ones_like(X_train[:,0])),
            (X_test, np.ones_like(X_test[:,0])),
            Y_train,
            Y_test
            )

    rowtraineval[0], rowtesteval[0] = linreg(Xtform_train, Xtform_test, Ytform_train, Ytform_test)

    #Change to polar coordinates
    if(not isPolar):
        #print("converting to polar...")
        X_train = get_polar(X_train)
        X_test = get_polar(X_test)

    ##DIR DEP
    Xtform_train, Xtform_test, Ytform_train, Ytform_test = feature_set(
        (np.cos(X_train[:,1]), np.sin(X_train[:,1]), np.ones_like(X_train[:,0])),
        (np.cos(X_test[:,1]), np.sin(X_test[:,1]), np.ones_like(X_test[:,0])),
        Y_train,
        Y_test
        )

    rowtraineval[1], rowtesteval[1] = linreg(Xtform_train, Xtform_test, Ytform_train, Ytform_test)

    ##VEL DEP
    Xtform_train, Xtform_test, Ytform_train, Ytform_test = feature_set(
        (X_train[:,0], np.ones_like(X_train[:,0])),
        (X_test[:,0], np.ones_like(X_test[:,0])),
        Y_train,
        Y_test
        )
    rowtraineval[2], rowtesteval[2] = linreg(Xtform_train, Xtform_test, Ytform_train, Ytform_test)

    ##DIR PLUS VEL DEP
    Xtform_train, Xtform_test, Ytform_train, Ytform_test = feature_set(
        (X_train[:,0] * np.cos(X_train[:,1]), X_train[:,0] * np.sin(X_train[:,1]), np.ones_like(X_train[:,0])),
        (X_test[:,0] * np.cos(X_test[:,1]), X_test[:,0] * np.sin(X_test[:,1]), np.ones_like(X_test[:,0])),
        Y_train,
        Y_test
        )

    rowtraineval[3], rowtesteval[3] = linreg(Xtform_train, Xtform_test, Ytform_train, Ytform_test)

    return (row, rowtraineval, rowtesteval)

# %% LABEL SPECIFICITY

def tune_row_label(X_train, X_test, Y_train, Y_test, node):
    """Perform ANOVA-like analysis to determine how much of the variance can be
    explained by tuning to labels for a single unit ("row").

    Arguments
    ----------
    X: labels
    Y: node activitiy for each sample
    row: row index

    Returns
    ----------
    node: node index
    nodeeval: np.array
    """

    ##RESHAPE OPERATIONS
    X_train = label_binarize(X_train, [ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13., 14., 15., 16., 17., 18., 19.,])
    X_test = label_binarize(X_test, [ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13., 14., 15., 16., 17., 18., 19.,])
    Y_train = Y_train.reshape(-1,1)
    Y_test = Y_test.reshape(-1,1)

    try:
        svm = OneVsRestClassifier(LinearSVC(max_iter=10, verbose=0))
        svm.fit(Y_train, X_train)

        nodetraineval = roc_auc_score(X_train, svm.decision_function(Y_train))
        try:
            nodetesteval = roc_auc_score(X_test, svm.decision_function(Y_test))
        except ValueError as err:
            print("test evaluation failed. %s" %err)
            nodetesteval = 0.5

    except ValueError as err:
         print("fitting SVM failed. %s" %err)
         nodetraineval = 0.5
         nodetesteval = 0.5

    #print("Train eval: %s" %str(nodetraineval))
    #print("Test eval: %s" %str(nodetesteval))

    return node, np.array([(nodetraineval-0.5)*2]), np.array([(nodetesteval-0.5)*2])

# %% TUNE

def tune(X, fset, Y, centers, nmods, nmets, ilayer, mmod='std', pool=None):
    ''' Makes calls to tuning curves for individual rows using multiprocessing

       Arguments
       ---------
       X : np.array
       fset : string, name of kinematic feature
       Y : hidden layer activations to be fitted
       centers : np.array of ints, mask for downsampling that contains the integer positions for temporal dimension that correspond to units in hidden layers
       nmods : number of different types of models that will be tested for this kinematic feature
       nmets : number of metrics to be used (usually 3 to correspond to output of compute_metrics)
       mmod : modifier for the kinematic feature type (default 'std')

       Returns
       -------
       trainevals : np.array [nr rows / units, nr of tuning curves, nr of metrics]
       testevals : np.array [nr rows / units, nr of tuning curves, nr of metrics]

    '''

    Yshape = Y.shape
    if len(Yshape) == 3:
        ravelshape = (Y.shape[1], nmods, nmets)
        unravelshape = (Y.shape[1], nmods, nmets)
    else:
        ravelshape = (Y.shape[1] * Y.shape[3], nmods, nmets )
        unravelshape = (Y.shape[1], Y.shape[3], nmods, nmets)

    trainevals =  np.empty(ravelshape)
    testevals = np.empty(ravelshape)
        #Axis 0-k: lo dims except time
        #Axis k+1: (0) LinReg (1) Dir Tuning (2) Vel Tuning (3) Dir + Vel Tuning
        #Axis k+2: (0) RMSE (1) r2 (2) PCC

    if(fset != 'labels'):
        n_cpus = 10
    else:
        n_cpus = 5
    print('Max CPU Count: %d , using %d ' %(mp.cpu_count(), n_cpus))
    pool = mp.Pool(n_cpus)
    print("Pool started")

    results = []

    # Resize Y so that feature maps are appended as new rows in first feature map
    Y = Y.swapaxes(1,2).reshape((Y.shape[0], Y.shape[2], -1)).swapaxes(1,2)
    print("Y reshaped")

    for irow in range(Y.shape[1]):
        print("Layer %d, Node %d / %d" %(ilayer, irow, Y.shape[1]))

        if(len(X.shape) > 1):
            x = X[..., centers]
        else:
            x = X

        y = Y[:,irow]

        ### DO TRAIN-TEST SPLIT
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.20, random_state=42)

        ##RESHAPE FOR LINEAR REGRESSION

        if len(x.shape) > 1: ##FOR TIME-BASED DATA ( 2 COMPS PER TIMESPACE IS USE CASE)

            tcoff = sum(np.where(centers <= 32, True, False))
            x_train = x_train[...,tcoff:ntime-tcoff]
            x_test = x_test[...,tcoff:ntime-tcoff]
            y_train = y_train[:,tcoff:ntime-tcoff]
            y_test = y_test[:,tcoff:ntime-tcoff]
            x_train = x_train.swapaxes(1,2).reshape((-1,2))
            x_test = x_test.swapaxes(1,2).reshape((-1,2))
        elif fset == 'labels':
            temp_train = np.ones_like(y_train)
            temp_test = np.ones_like(y_test)
            x_train = (temp_train.swapaxes(0,1) * x_train).swapaxes(0,1)
            x_train = x_train.reshape((-1,))
            x_test = (temp_test.swapaxes(0,1) * x_test).swapaxes(0,1)
            x_test = x_test.reshape((-1,))

        y_train = y_train.reshape((-1,))
        y_test = y_test.reshape((-1,))

        if fset == 'acc':
            results.append(pool.apply_async(tune_row_vel, args=(x_train, x_test, y_train, y_test,irow, True)))
        elif fset == 'vel' or fset == 'eepolar' or fset == 'ang' or fset=='angvel':
            results.append(pool.apply_async(tune_row_vel, args=(x_train, x_test, y_train, y_test ,irow, True)))
        elif fset == 'ee':
            results.append(pool.apply_async(tune_row_vel, args=(x_train, x_test,y_train, y_test,irow, False)))
            #sys.stdout.flush()
        elif fset == 'labels':
            results.append(pool.apply_async(tune_row_label, args=(x_train, x_test, y_train, y_test, irow)))

    results = [r.get() for r in results]

    pool.close()
    print("pool closed")
    pool.join()
    print("pool joined")

    for result in results:
        trainevals[result[0]] = result[1]
        testevals[result[0]] = result[2]

    #shape evals back to original shape
    trainevals = trainevals.reshape(unravelshape)
    testevals = testevals.reshape(unravelshape)

    return trainevals, testevals

def tune_decoding(X, fset, Y, centers, ilayer, mmod, alpha = None):
    """ Computes decoding accuracies for given kinetic variable

    Arguments
    ----------

    Returns
    -------
    trainevals : np.array of floats [nr_features, 3]
    testevals : np.array of floats [nr_features, 3]
    coefs : np.array of floats [nr_features, nr_coefs]
    """

    print("Running decoding with alpha parameter %f ..." %alpha)

    if(len(X.shape) > 1):
        X = X[..., centers]

    # Resize Y so that feature maps are appended as new rows in first feature map
    Y = Y.swapaxes(1,2).reshape((Y.shape[0], Y.shape[2], -1)).swapaxes(1,2)

    print("X has size %s and Y %s" %(str(X.shape), str(Y.shape)))

    X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.20, random_state=42)

    if fset != 'labels':
        # reshape so that both X and Y are in format [samples x timepoints, features] (except for labels)
        print(X.shape)
        assert len(X.shape) > 1, "X has shape 1 %s" %str(X.shape)

        X_train = X_train.swapaxes(1,2).reshape((-1, X_train.shape[1]))
        if len(Y_train.shape) > 2:
            Y_train = Y_train.swapaxes(1,2).reshape((-1, Y_train.shape[1]))
        else:
            Y_train = Y_train.reshape((-1, Y_train.shape[1]))

        X_test = X_test.swapaxes(1,2).reshape((-1, X_test.shape[1]))
        if len(Y_test.shape) > 2:
            Y_test = Y_test.swapaxes(1,2).reshape((-1, Y_test.shape[1]))
        else:
            Y_test = Y_test.reshape((-1, Y_test.shape[1]))
    # for labels: reshape so that both are in format [samples, timepoints x features]
    else:
        print("About to binarize shape %s" %str(X_train.shape))
        X_train = label_binarize(X_train, [ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13., 14., 15., 16., 17., 18., 19.,])
        print("New shape of X %s"%str(X.shape))
        Y_test = Y_test.reshape((Y_test.shape[0], -1))

        print("About to binarize shape %s" %str(X_test.shape))
        X_train = label_binarize(X_train, [ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13., 14., 15., 16., 17., 18., 19.,])
        print("New shape of X %s"%str(X_train.shape))
        Y_test = Y_test.reshape((Y_test.shape[0], -1))

    #switch kin vars / labels to Y and neuron firing ratest to X
    X_temp_train = X_train
    X_train = Y_train
    Y_train = X_temp_train

    X_temp_test = X_test
    X_test = Y_test
    Y_test = X_temp_test

    if len(Y_train.shape) > 1:
        nna = ~np.any(np.isnan(Y_train), axis=1)
    else:
        nna = ~np.isnan(Y_train)
    X_train = X_train[nna]
    Y_train = Y_train[nna]

    if len(Y_test.shape) > 1:
        nna = ~np.any(np.isnan(Y_test), axis=1)
    else:
        nna = ~np.isnan(Y_test)
    X_test = X_test[nna]
    Y_test = Y_test[nna]

    trainevals = []
    testevals = []
    coefs = []

    assert Y_train.shape[1] > 1, 'Y is supposed to have multiple targets/columns'

    assert Y_train.size < X_train.size, "somehow size of target is bigger than size of training, are you sure that you are decoding the kinematic variable and not the inverse?"

    if(fset == 'ee'):
        ee_lms = []

    if (len(Y_train) >= 1 and len(Y_test)  >= 1):

        for target in range(Y_train.shape[1]):

            if fset != 'labels':

                if alpha is None:
                    lm = LinearRegression().fit(X_train, Y_train[:,target])

                else:
                    lm = Ridge(alpha = alpha).fit(X_train, Y_train[:,target])
                    print("Running regression with alpha %f ..." %alpha)

                trainevals.append(compute_metrics(Y_train[:, target], lm.predict(X_train)))
                testevals.append(compute_metrics(Y_test[:, target], lm.predict(X_test)))

                coefs.append(lm.coef_.copy())

                if fset == 'ee':
                    ee_lms.append(lm)

            else:
                svm = OneVsRestClassifier(LinearSVC(max_iter=10, verbose=0))
                try:
                    svm.fit(X_train, Y_train)
                except ValueError as err:
                    print("fitting SVM failed. %s" %err)
                    nodetraineval = 0.5

                try:
                    print("About to calculate ROC: X_train Shape: %s, Y_train Shape: %s" %(str(X_train.shape), str(Y_train.shape)))
                    nodetraineval = roc_auc_score(Y_train, svm.decision_function(X_train))
                except ValueError as err:
                    print("train evaluation failed. %s" %err)
                    nodetraineval = 0.5

                try:
                    nodetesteval = roc_auc_score(Y_test, svm.decision_function(X_test))
                except ValueError as err:
                    print("test evaluation failed. %s" %err)
                    nodetesteval = 0.5

                #print("Train eval: %s" %str(nodetraineval))
                #print("Test eval: %s" %str(nodetesteval))

                trainevals.append(np.array([(nodetraineval-0.5)*2]*3))
                testevals.append(np.array([(nodetesteval-0.5)*2]*3))
                coefs.append([np.nan]*3)

    else:
        trainevals.append([np.nan]*3)
        testevals.append([np.nan]*3)
        coefs.append([np.nan]*3)

    trainevals = np.vstack(trainevals)
    testevals = np.vstack(testevals)
    coefs = np.vstack(coefs)

    #for ees: compute "mean distance" metric and store in previous PPC columns
    if fset == 'ee':
        Y_train_pred = np.hstack([ee_lms[i].predict(X_train).reshape(-1,1) for i in range(2)])
        Y_test_pred = np.hstack([ee_lms[i].predict(X_test).reshape(-1,1) for i in range(2)])

        print("Shape of Y_train_pred: ", Y_train_pred.shape)
        assert Y_train_pred.shape[1] == 2 and len(Y_train_pred.shape) == 2, "unexpected shape of predictions, Y_train_pred shape " + str(Y_train_pred.shape)

        trainevals[:,2] = compute_dist_metric(Y_train, Y_train_pred)
        testevals[:,2] = compute_dist_metric(Y_test, Y_test_pred)

        print("Rewrote the PPC column :)")

    return (trainevals, testevals, coefs)

def tune_layer(X, fset, xyplmvt, runinfo, ilayer, mmod, model, t_stride=2, pool=None, alpha=None):
    """Performs several layer-wide calculations for fitting the tuning curves and saves output to file

    Arguments
    ---------
    X : np.array of floats, [nr_samples, nr_features, 320 timepoints]
    fset : string, kinematic tuning curve name
    xyplmvt : np.array of bools [nr_samples,], mask restricting to trials / movements that occur in the desired plane or orientation
    runinfo : RunInfo (dict extension)
    ilayer : int, index of current layer, -1 = spindles
    mmod : string, modifier to model or kinematic tuning curve type
    model : dict, information on model
    t_stride : int, temporal stride used when applying convolutional filters

    """

    savepath = runinfo.resultsfolder(model, fset)

    if mmod == 'decoding':
        savepath = runinfo.resultsfolder(model, '%s_%s' %(mmod, fset))

    try:
        os.makedirs(savepath)
    except:
        print('folder already exists')

    savename = 'l%d_%s_mets_%s_%s' %(ilayer+1, fset, mmod, runinfo.planestring())

    modelname = model['name']
    base = model['base']
    layer = lstring(ilayer)
    expid = runinfo['expid']

    #lo = pickle.load(open(os.path.join(runinfo.datafolder(model), layer + '.pkl'), 'rb'))
    lo = read_layer_reps(ilayer, runinfo, model)
    lo = lo[xyplmvt]

    centers = get_centers(lo.shape[2], ilayer, model)

    if mmod=='decoding':
        trainevals, testevals, coefs = tune_decoding(X, fset, lo, centers, ilayer, mmod, alpha)

        print('saving decoding reg coefs')
        if alpha is None:
            savefile = '%s/%s_coefs.npy' %(savepath, savename)
        else:
            savefile = '%s/%s_a%d_coefs.npy' %(savepath, savename, int(alpha*1000))

        np.save(savefile, coefs)

    else:
        if (fset == 'vel' or fset == 'acc' or fset == 'eepolar' or fset == 'ang' or fset=='angvel' or fset=='ee'):
            nmods = 4
            nmets = 6
        elif (fset == 'labels'):
            nmods = 1
            nmets = 1

        if len(X) > 1:
            trainevals, testevals = tune(X, fset, lo, centers, nmods, nmets, ilayer, mmod, pool=pool)
        else:
            trainevals = np.empty((0, nmods, nmets))
            testevals = np.empty((0, nmods, nmets))

    print("Layer %d Completed!" %ilayer)

    if alpha is None:
        savefile_train = '%s/%s_train.npy' %(savepath, savename)
        savefile_test = '%s/%s_test.npy' %(savepath, savename)
    else:
        savefile_train = '%s/%s_a%d_train.npy' %(savepath, savename,  int(alpha*1000))
        savefile_test = '%s/%s_a%d_test.npy' %(savepath, savename,  int(alpha*1000))

    np.save(savefile_train, trainevals)
    np.save(savefile_test, testevals)
    print("files saved")

# %% DATA READIN

def X_data(fset = 'vel', runinfo = dict({'orientation': 'hor', 'plane': 'all', 'datafraction': 0.5}),
           datafolder = '../data/spatial_temporal_4_8-16-16-32_64-64-64-64_5272/',
           threed= False,
           polar=True):
    """Returns input data for model. Uses global parameters.

    Arguments
    ---------
    fset : string, one of ['acc', 'ang', 'angvel', 'labels', 'ee', 'eepolar', 'muscle_coords', 'mf' (muscle firing rates)]
    runinfo : RunInfo (dict extension)
    datafolder : string, location where data is stored
    threed : bool, are all three coords needed or only the two in which movement occurs
    polar : bool, should data be converted to polar

    Returns
    ---------
    vels: np.array of floats [nsamples, 2, 320]
    xyplmvt: list of bools [nsamples_original], boolean mask to mask samples
        that are not in xy plane
    """

    ### INIT
    print(datafolder)
    try:
        kinvars = pd.read_hdf(os.path.join(datafolder, 'kinvars.hdf5'))
        mc = np.swapaxes(kinvars['muscle_coords'].values, 0, 1)
        ee = np.swapaxes(kinvars['endeffector_coords'].values, 0, 1)
    except:
        print('using alternate method for accessing kinvars files by directly accessing needed arrays (pandas causes error)')
        filename = os.path.join(datafolder, 'kinvars.hdf5')

        with h5py.File(filename, "r") as f:

            key = 'data'
            data_group = 'block0_values'
            data_values = f[key][data_group][:]

            ee = data_values[:,:,:3]
            ee = np.swapaxes(ee, 1, 2)
            ee = np.swapaxes(ee, 0, 2)

            mc = data_values[:,:,7:32]
            mc = np.swapaxes(mc, 1, 2)
            mc = np.swapaxes(mc, 0, 2)

    vels = np.gradient(ee)[2]

    # EXCLUDE SAMPLES THAT AREN'T IN XY PLANE
    if(runinfo['orientation']=='hor'):
        if runinfo['height']=='all':
            xyplmvt = (1 - np.any(vels[:,2,:], axis=1)).astype('bool')
        else:
            xyplmvt = np.all(ee[:,2,:] == runinfo['height'], axis=1)

        if not threed:
            ee = ee[:,:2]
            vels = vels[:,:2]
    else:
        if runinfo['height']=='all':
            xyplmvt = np.invert(np.any(vels[:,0,:], axis=1)) #movements in a yz plane
        else:
            xyplmvt = np.all(ee[:,0,:] == runinfo['height'], axis=1)

        if not threed:
            ee = ee[:, 1:]
            vels = vels[:, 1:]

    ee = ee[xyplmvt]
    vels = vels[xyplmvt]
    mc = mc[xyplmvt]
    X = vels

    if len(X) > 1:
        print(X.size)
        print(X.shape)

        if(fset == 'acc'):
            X = np.gradient(vels)[2]
            if polar:
                X = get_polar(X)
            #X = np.nan_to_num(X)
        elif(fset == 'ang'):
            Xa = np.apply_along_axis(angle_xaxis, 1, vels)
            Xb = np.gradient(Xa)[1]
            Xa = Xa[:,np.newaxis,:]
            Xb = Xb[:,np.newaxis,:]
            X = np.concatenate((Xb, Xa), axis=1)
        elif(fset == 'angvel'):
            X = np.apply_along_axis(angle_xaxis, 1, vels)
            X = np.gradient(X)[1]
            X = np.concatenate((X[:,np.newaxis,:], X[:,np.newaxis,:]), axis=1)
        elif(fset == 'labels'):
            X = pickle.load(open(os.path.join(datafolder, 'labels.pkl'), 'rb'))
            X = X[xyplmvt]
        elif(fset == 'ee'):
            X = ee
        elif(fset == 'eepolar'):
            x0 = ee[:,:,0]
            Xrel = [(ee[i].swapaxes(0,1) - x0[i]).swapaxes(0,1) for i in range(len(ee))]
            X[:,0,:] = np.apply_along_axis(np.linalg.norm, 1, Xrel)
            X[:,1,:] = np.apply_along_axis(angle_xaxis, 1, Xrel)

        elif(fset == 'vel'):

            if polar:
                X = get_polar(X)

        elif(fset == 'muscle_coords'):
            X = mc

        elif (fset == 'mf'):
            data = pickle.load(open(os.path.join(datafolder, 'data.pkl'), 'rb'))
            X = data[xyplmvt]

    return X, xyplmvt

# %% MAIN

def main(fset, runinfo, model, startlayer=-1, endlayer=8, mmod='std', pool=None, alpha = None):

    assert fset == 'vel' or fset == 'acc' or fset == 'eepolar' or fset == 'ang'\
        or fset == 'labels' or fset=='angvel' or fset=='ee', "Invalid fset!!!"

    assert mmod == 'std' or mmod=='decoding', 'Invalid mmod!!!'

    assert mmod == 'decoding' or alpha is None, "alpha values only accepted for decoding"

    modelname = model['name']
    nlayers = model['nlayers']
    base = model['base']

    print('evaluating %s model for %s %s, expid %d, plane %s ...' %(modelname, fset, mmod, runinfo['expid'], runinfo.planestring()))

    X, xyplmvt = X_data(fset, runinfo, datafolder=runinfo.datafolder(model))

    np.random.seed(42)

    for ilayer in np.arange(startlayer, min(endlayer, nlayers)):
        tune_layer(X, fset, xyplmvt, runinfo, ilayer, mmod, model, pool=pool, alpha=alpha)

if __name__=='__main__':

    parser = argparse.ArgumentParser(description='Tuning Properties Analysis for DeepDraw Models.')
    parser.add_argument('--fset', type=str, default = 'vel', help='Type of Features {"vel", "acc", "ang", "labels"}')
    parser.add_argument('--expid', type=int, help='Experiment ID Number')
    parser.add_argument('--startlayer', type=int, default=-1, help='Layer you want to start at')
    parser.add_argument('--endlayer', type=int, default=8, help='Layer you want to end before')
    parser.add_argument('--mmod', type=str, default='std', help='Model modifier')
    parser.add_argument('--mtype', type=str, default='S', \
                        help='Model type to be analyzed {(S)patial temporal, (T)emporal spatial, (S)patio(T)emporal, (F)ully (C)onnected}'
                        )

    args = parser.parse_args()

    main(args.fset, args.expid, args.startlayer, args.endlayer, args.mmod, args.mtype)


import os
import homology
import utils
import filters
import hpa
import go
import pandas as pd

def get_proteins(config_file):
    """
    Retrieve all proteins for all species
    :param str config_file: path to config file

    :return: dictionary with all proteins for all species. Key -> tax id, value -> dictionary: key -> protein id, value -> protein name"""
    proteins = {}
    hosts = utils.read_config(filepath=config_file, field='hosts')
    parasites = utils.read_config(filepath=config_file, field='parasites')
    urls = utils.read_config(filepath=config_file, field='urls')
    if "string_protein_url" in urls:
        string_file = urls['string_protein_url']
        if hosts is not None and parasites is not None:
            taxids = list(hosts.keys()) + list(parasites.keys())
            for taxid in taxids:
                proteins[taxid] = parse_proteins(string_file, taxid)

    return proteins

def parse_proteins(string_file, taxid):
    """
    Retrieve proteins for a given specie
    :param str string_file: url to string PPI file
    :param int taxid: taxonomic id of the species of interest

    :return: dictionary with all proteins. Key -> Ensembl protein id, value -> protein name
    """
    proteins = {}
    if string_file is not None:
        filename = utils.download_file(url=string_file.replace('TAXID', str(taxid)), data_dir='data')
        sp = utils.read_gzipped_file(filename)
        first = True
        for line in sp:
            if first:
                first = False
                continue

            data = line.decode("utf-8").rstrip().split('\t')
            identifier = data[0]
            name = data[1]
            proteins[identifier] = name

    return proteins

def get_tissue_cell_type_annotation(tissues, output_file):
    tissues_df = pd.concat({k: pd.Series(v) for k, v in tissues.items()}).reset_index()
    tissues_df = tissues_df.iloc[:, [0, 2]]
    tissues_df.columns = ['Gene', 'Tissue']
    tissues_df = tissues_df[tissues_df['Gene'].isin(proteins.keys())]
    hpa_data = hpa.parse_hpa(config_file, valid_proteins=proteins.keys())
    tissues_df = pd.merge(tissues_df, hpa_data, on=['Gene', 'Tissue'], how='left')

    utils.save_to_parquet(tissues_df, output_file)

def setup(config_file, output_file_path):
    """
    Downloads all necessary files according to the urls specified in the configuration file
    except the protein files and host PPIs that are downloaded for each species. The go terms
    will also be downloaded and formatted.

    :param str config_file: path to the configuration file
    """
    urls = utils.read_config(filepath=config_file, field='urls')
    for url_name in urls:
        url = urls[url_name]
        if url_name != "string_protein_url" and url_name != "string_ppi_url" and url_name != "string_go_url":
            filename = utils.download_file(url=url, data_dir=output_file_path)

    go.get_gene_ontology(config_file, output_dir=output_file_path)

if __name__ == "__main__":
    data_dir = 'data'
    config_file = 'config.yml'

    setup(config_file=config_file, output_file_path=data_dir)

    hosts = utils.read_config(filepath=config_file, field='hosts')
    parasites = utils.read_config(filepath=config_file, field='parasites')

    #Get host and parasite proteins
    proteins = get_proteins(config_file)

    #Apply filters -- secretome, tissue, cellular compartment context
    proteins = filters.get_secretome_predictions(config_file=config_file, secretome_dir='data/secretome_pred_input_data/input_data', valid_proteins=proteins)
    tissues = filters.apply_tissue_filter(config_file, proteins, cutoff=2.5)
    compartments = filters.apply_compartment_filter(config_file, proteins, cutoff=2.5)
    proteins = utils.merge_dict_of_dicts(dict_of_dicts=proteins)

    #Annotate tissue and cell type expression
    get_tissue_cell_type_annotation(tissues, output_file=os.path.join(data_dir, 'tissues_cell_types.parquet'))

    #Get eggnog groups and transfer PPIs
    valid_groups = homology.get_eggnog_groups(filepath=os.path.join(data_dir, '2759_members.tsv.gz'), proteins=proteins.keys())
    homology.get_links(filepath=os.path.join(data_dir, 'COG.links.detailed.v11.5.txt.gz'), valid_groups=valid_groups, proteins=proteins,
              ouput_filepath=os.path.join(data_dir, 'predictions.parquet'), config_file=config_file)

    predictions = pd.read_parquet(os.path.join(data_dir, 'predictions.parquet'))
    predictions = utils.annotate_alias_id(predictions_df=predictions,
                            taxids=list(parasites.keys()), config_file=config_file,
                            sources=['BLAST_UniProt_AC'], new_col="source_uniprot",
                            mapping_col="source")

    predictions = utils.annotate_alias_id(predictions_df=predictions,
                            taxids=list(hosts.keys()), config_file=config_file,
                            sources=['Ensembl_HGNC_UniProt_ID(supplied_by_UniProt)'],
                            new_col="target_uniprot", mapping_col="target")

    utils.save_to_parquet(df=predictions, output_file=os.path.join(data_dir, 'annotated_predictions.parquet'))


import argparse
import logging

import hail as hl

from gnomad.sample_qc.filtering import compute_stratified_metrics_filter
from gnomad.utils.annotations import bi_allelic_expr
from gnomad.utils.filtering import (
    filter_low_conf_regions,
    filter_to_autosomes,
)
from gnomad.utils.slack import slack_notifications
from ukbb_qc.resources.basics import (
    array_sample_map_ht_path,
    get_ukbb_data,
    logging_path,
)
from ukbb_qc.resources.resource_utils import CURRENT_FREEZE
from ukbb_qc.resources.sample_qc import (
    ancestry_hybrid_ht_path,
    platform_pca_assignments_ht_path,
    platform_pop_outlier_ht_path,
    qc_temp_data_prefix,
)
from ukbb_qc.slack_creds import slack_token
from ukbb_qc.utils.utils import annotate_interval_qc_filter, remove_hard_filter_samples

logging.basicConfig(format="%(levelname)s (%(name)s %(lineno)s): %(message)s")
logger = logging.getLogger("outlier_filter")
logger.setLevel(logging.INFO)

def run_sample_qc(mt: hl.MatrixTable) -> hl.Table:
    """
    Filter input MatrixTable to bi-allelic sites, removes problematic intervals, and computes sample QC metrics

    :param MatrixTable mt: Input MatrixTable on which sample QC metrics need to be computed
    :return: MatrixTable filtered to autosomes and high-confidence regions, with computed sample QC column annotations
    :rtype: hl.MatrixTable
    """
    # Filtering to bi-allelic sites outside of low confidence regions
    mt = filter_to_autosomes(mt)
    mt = filter_low_conf_regions(mt, filter_decoy=False)
    mt = mt.filter_rows(bi_allelic_expr(mt))

    logger.info("Starting sample QC...")
    mt = hl.sample_qc(mt)
    return mt.cols().select("sample_qc")

def main(args):

    hl.init(log="/outlier_detection.log", default_reference="GRCh38")

    data_source = "broad"
    freeze = args.freeze
    pop_assignment_method = args.pop_assignment_method

    try:
        if args.run_mini_qc:
            # NOTE: we run outlier detection without adj filtration to get better separation between high and low quality samples
            # this is per Julia's discussion with Konrad in #ukbb_qc
            # Need all workers for the mini qc
            mt = get_ukbb_data(data_source, freeze, adj=False)
            mt = annotate_interval_qc_filter(
                data_source,
                freeze,
                mt,
                cov_filter_field=args.cov_filter_field,
                pct_samples=args.pct_samples,
                autosomes_only=True,
            )
            mt = mt.filter_rows(mt.interval_qc_pass)

            logger.info("Filtering samples that fail hard filters...")
            mt = remove_hard_filter_samples(data_source, freeze, mt, gt_field="GT")
            logger.info(f"Sample count post filtering: {mt.count_cols()}")

            logger.info(
                "Running mini sample QC for platform- and population-specific filtering..."
            )
            sample_qc_ht = run_sample_qc(mt)
            sample_qc_ht = sample_qc_ht.repartition(args.n_partitions)
            sample_qc_ht.write(
                qc_temp_data_prefix(data_source, freeze)
                + "outlier_sample_qc_intervals.ht",
                args.overwrite,
            )

        sample_qc_ht = hl.read_table(
            qc_temp_data_prefix(data_source, freeze) + "outlier_sample_qc_intervals.ht"
        )
        strata = {}

        logger.info("Annotating population assignments...")
        pop_ht = hl.read_table(ancestry_hybrid_ht_path(data_source, freeze))
        sample_qc_ht = sample_qc_ht.annotate(
            qc_pop=pop_ht[sample_qc_ht.key][pop_assignment_method].pop
        )
        logger.info("Annotating inferred platform assignments...")
        platform_ht = hl.read_table(
            platform_pca_assignments_ht_path(data_source, freeze)
        )
        sample_qc_ht = sample_qc_ht.annotate(
            qc_platform=platform_ht[sample_qc_ht.key].qc_platform
        )
        logger.info("Annotating with batch (tranche) as a proxy for platform...")
        sample_map_ht = hl.read_table(array_sample_map_ht_path(freeze))
        sample_qc_ht = sample_qc_ht.annotate(
            batch=sample_map_ht[sample_qc_ht.key].batch
        )

        if args.population_filter:
            strata["qc_pop"] = sample_qc_ht.qc_pop
        else:
            pop_assignment_method = "no_pop"

        if args.platform_filter or args.batch_filter:
            if args.platform_filter:
                strata["qc_platform"] = sample_qc_ht.qc_platform
                platform_assignment_method = "qc_platform"

            if args.batch_filter:
                # NOTE: used tranche as a proxy for platform in tranche 2/freeze 5/200K
                strata["qc_platform"] = sample_qc_ht.batch
                platform_assignment_method = "batch"
        else:
            platform_assignment_method = "no_platform"

        # Make qc_metrics a dict (needs to be dict for compute_stratified_metrics_filter)
        metrics = args.filtering_qc_metrics.split(",")
        qc_metrics = {}
        for m in metrics:
            qc_metrics[m] = sample_qc_ht.sample_qc[f"{m}"]

        # For each platform and population, aggregate sample QC metrics and calculate the MAD/mean/stdev
        logger.info(
            "Flagging samples failing pop/platform-specific sample qc thresholds..."
        )
        pop_platform_filter_ht = compute_stratified_metrics_filter(
            sample_qc_ht, qc_metrics, strata
        )
        pop_platform_filter_ht = pop_platform_filter_ht.repartition(args.n_partitions)
        pop_platform_filter_ht.write(
            platform_pop_outlier_ht_path(
                data_source, freeze, pop_assignment_method, platform_assignment_method
            ),
            overwrite=args.overwrite,
        )
        num_pass = pop_platform_filter_ht.aggregate(
            hl.agg.count_where(hl.len(pop_platform_filter_ht.qc_metrics_filters) == 0)
        )
        logger.info(
            f"{num_pass} exome samples found passing pop/platform-specific filtering"
        )

    finally:
        logger.info("Copying hail log to logging bucket...")
        hl.copy_log(logging_path(data_source, freeze))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-o",
        "--overwrite",
        help="Overwrite all data from this subset (default: False)",
        action="store_true",
    )
    parser.add_argument(
        "--slack_channel", help="Slack channel to post results and notifications to."
    )
    parser.add_argument(
        "-f", "--freeze", help="Data freeze to use", default=CURRENT_FREEZE, type=int
    )
    parser.add_argument(
        "--n_partitions",
        help="Desired number of partitions. Used for ALL output HTs",
        default=5000,
        type=int,
    )

    parser.add_argument(
        "--run_mini_qc",
        help="Run mini sample qc needed for outlier filtering",
        action="store_true",
    )
    parser.add_argument(
        "--cov_filter_field",
        help="Coverage field to use to filter high coverage intervals",
        default="pct_samples_20x",
    )
    parser.add_argument(
        "--pct_samples",
        help="Percent samples at specified coverage to filter intervals",
        default=0.85,
        type=float,
    )
    platform_group = parser.add_mutually_exclusive_group()
    platform_group.add_argument(
        "--platform_filter",
        help="Include inferred platforms in outlier filtering",
        action="store_true",
    )
    platform_group.add_argument(
        "--batch_filter",
        help="Include batch (tranche) as proxy for platform in outlier filtering",
        action="store_true",
    )
    parser.add_argument(
        "--population_filter",
        help="Included inferred population in outlier filtering",
        action="store_true",
    )
    parser.add_argument(
        "--pop_assignment_method",
        help="Population assignment method to use for stratification",
        default="hybrid_pop_data",
        choices=["gnomad_pc_project_pop_data", "hybrid_pop_data"],
    )
    parser.add_argument(
        "--filtering_qc_metrics",
        help="List of QC metrics for filtering.",
        default=",".join(
            [
                "n_snp",
                "r_ti_tv",
                "r_insertion_deletion",
                "n_insertion",
                "n_deletion",
                "r_het_hom_var",
            ]
        ),
    )

    args = parser.parse_args()

    if args.slack_channel:
        with slack_notifications(slack_token, args.slack_channel):
            main(args)
    else:
        main(args)


"""
Design Hit Counter
Things to consider:
1. If a lot of hits within given time window / per sec
2. Multi-threading, if multiple threads report hit #Example: if a webpage, multiple clients click hit
"""

from collections import deque

"""
1. 简单版本，不考虑多线程
"""
class HitCounter:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self._counter = deque()
        self.th = 300

    def hit(self, timestamp: int) -> None:
        """
        Record a hit.
        @param timestamp - The current timestamp (in seconds granularity).
        """
        self._counter.append(timestamp)

    def getHits(self, timestamp: int) -> int:
        """
        Return the number of hits in the past 5 minutes.
        @param timestamp - The current timestamp (in seconds granularity).
        """
        cutoff = timestamp - self.th
        if self._counter and self._counter[-1] <= cutoff:
            return 0

        while self._counter:
            if self._counter[0] <= cutoff:
                self._counter.popleft()
            else:
                break

        return len(self._counter)

# python solution: O(1) Space, O(1) hit, O(S) gethits()
"""
using a tuple and referencing the variable at the same time so even if the context switch
happens in a multithreaded environment, both hits and times will be in sync

can achieve synchronized modification and reading
or using blocking queue

it can handle multiple hit at the same time, because this approach does not involve dynamic allocating the space
"""

class HitCounter2:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.th = 300
        self.q = [(0, 0)] * self.th

    def hit(self, timestamp: int) -> None:
        """
        Record a hit.
        @param timestamp - The current timestamp (in seconds granularity).
        """
        idx = timestamp % self.th
        time, hit = self.q[idx]
        if time != timestamp:
            # previous timestamp is outdated, safe to override
            self.q[idx] = timestamp, 1
        else:
            self.q[idx] = time, hit + 1

    def getHits(self, timestamp: int) -> int:
        """
        Return the number of hits in the past 5 minutes.
        @param timestamp - The current timestamp (in seconds granularity).
        """
        cnt = 0
        cutoff = timestamp - self.th
        for time, hit in self.q:
            if time > cutoff:
                cnt += hit
        return cnt


#!/usr/bin/env python
from pathlib import Path
from distutils.core import setup

README = Path(__file__).parent / "README.md"
with open(README) as fp:
    long_description = fp.read()

setup(
    name="PyAutoBlockchain",
    version="0.5",
    description="PAB is a framework that helps with development and automation of periodic tasks on blockchains.",
    long_description=long_description,
    long_description_content_type="text/markdown",
    author="Manuel Pepe",
    author_email="manuelpepe-dev@outlook.com.ar",
    url="https://github.com/manuelpepe/PyAutoBlockchain",
    include_package_data=True,
    package_data={"pab": ["resources/config.schema.json"]},
    entry_points={
        "console_scripts": [
            "pab = pab.cli:ep",
        ],
    },
    packages=[
        "pab",
        "pab.resources",
    ],
    install_requires=["web3", "hexbytes", "python-dotenv", "projectutils"],
    extras_require={"ui": ["pabui"]},
)


import logging
import re

logger = logging.getLogger(__name__)

import sublime

try:
    from DynamicCompletions import FileLoader
except ImportError as e:
    logger.error('DynamicCompletions package not installed')
    raise e

from .tools.classes import get_ring_file
from .tools.settings import get_completion_source_enabled_setting

from .misc.completion_types import (
    CT_ALIAS,
    CT_FOCUS_LOCAL,
    CT_SUBROUTINE,
    CT_LIST,

    CT_OBJECT,
    CT_RECORD,
    # CT_FILE,
    CT_KEY,
    CT_FIELD,
    # CT_INDEX,
    # CT_INDEXKEY,
    # CT_LONGLOCK,
)

class IncludeFileLoader(FileLoader):
    """
    Parent class for CompletionLoaders that load completions from an Include
    file.
    """

    LoadAsync = True

    @classmethod
    def view_scope(cls):
        """
        Return a scope to determine if the CompletionLoader will be enabled
        for a view.

        """
        return 'source.focus'

    @classmethod
    def view_check(cls, view):
        n = view.file_name()
        if (n is None) or (n == ''):
            return False
        return (get_ring_file(n) is not None)

    @classmethod
    def instances_for_view(cls, view):
        """
        Returns a list of instances of the given class to be used for the
        given view.

        """
        ring_file = get_ring_file(view.file_name())
        if ring_file is None:
            return []

        include_files = ring_file.get_include_files()
        if not include_files:
            return []

        instances = set()
        for f in include_files:
            logger.debug('%s;  f = %s', cls, f)
            try:
                instances.add(cls.Instances[f])
            except KeyError:
                instances.add(cls(file_path=f))
            except AttributeError:
                instances.add(cls(file_path=f))
        return instances

class AliasIncludeLoader(IncludeFileLoader):
    """
    Loads aliases defined in an Include File.
    """

    EmptyReturn = ([], sublime.INHIBIT_WORD_COMPLETIONS |
                   sublime.INHIBIT_EXPLICIT_COMPLETIONS)

    AliasMatcher = re.compile(r"^\s*:Alias\s+([^\n()\t]+?) *$", re.MULTILINE)

    @classmethod
    def completion_types(cls):
        """
        Return a set of completion types that this CompletionLoader can return.

        """
        return [CT_ALIAS]

    @classmethod
    def view_check(cls, view):
        if super().view_check(view):
            return get_completion_source_enabled_setting('Alias', 'Include')

    def load_completions(self, **kwargs):
        """
        Populate self.completions with Aliases from an Include file.

        """
        self.completions = set()
        ring_file = get_ring_file(self.file_path)
        for alias in ring_file.get_defined_aliases():
            alias = '@@' + alias + '()'
            self.completions.add((alias + '\tInclude', alias[2:]))

class LocalIncludeLoader(IncludeFileLoader):
    """
    Loads any locals used in an Include File.
    """

    EmptyReturn = ([], sublime.INHIBIT_WORD_COMPLETIONS |
                   sublime.INHIBIT_EXPLICIT_COMPLETIONS)

    @classmethod
    def completion_types(cls):
        """
        Return a list of completion types that this CompletionLoader can
        return.

        """
        return [CT_FOCUS_LOCAL]

    @classmethod
    def view_check(cls, view):
        if super().view_check(view):
            return get_completion_source_enabled_setting('Local', 'Include')

    def load_completions(self, **kwargs):
        """
        Populate self.completions with Locals from an Include file.

        """
        self.completions = set()
        ring_file = get_ring_file(self.file_path)
        for local in (ring_file.get_used_locals() |
                      ring_file.get_defined_locals()):
            self.completions.add((local + '\tInclude', local))

class ObjectIncludeLoader(IncludeFileLoader):
    """
    Loads DataDef completions from an Include File.
    """

    EmptyReturn = ([], sublime.INHIBIT_EXPLICIT_COMPLETIONS)

    @classmethod
    def completion_types(cls):
        """
        Return a list of completion types that this CompletionLoader can
        return.

        """
        return [CT_OBJECT, CT_RECORD, CT_KEY, CT_FIELD]

    @classmethod
    def view_check(cls, view):
        if super().view_check(view):
            return get_completion_source_enabled_setting('Object', 'Include')

    def load_completions(self, **kwargs):
        """
        Populate self.completions with objects from an Include file.

        """
        self.completions = dict()
        ring_file = get_ring_file(self.file_path)
        object_dict = ring_file.get_defined_objects()

        for key, object_set in object_dict.items():
            self.completions[key] = set(
                [(x + '\tInclude', x) for x in object_set])

class SubroutineIncludeLoader(IncludeFileLoader):
    """
    Loads subroutine completions from an Include File.
    """

    EmptyReturn = ([], sublime.INHIBIT_WORD_COMPLETIONS |
                   sublime.INHIBIT_EXPLICIT_COMPLETIONS)

    @classmethod
    def completion_types(cls):
        """
        Return a list of completion types that this CompletionLoader can
        return.

        """
        return [CT_SUBROUTINE]

    @classmethod
    def view_check(cls, view):
        if super().view_check(view):
            return get_completion_source_enabled_setting('Subroutine',
                                                         'Include')

    def load_completions(self, **kwargs):
        """
        Populate self.completions with subroutines from an Include file.

        """
        ring_file = get_ring_file(self.file_path)
        self.completions = set(
            [(x + '\tInclude', x) for x in
             ring_file.get_defined_subroutines()])

class ListIncludeLoader(IncludeFileLoader):
    """
    Loads subroutine completions from an Include File.
    """

    EmptyReturn = ([], sublime.INHIBIT_WORD_COMPLETIONS |
                   sublime.INHIBIT_EXPLICIT_COMPLETIONS)

    @classmethod
    def completion_types(cls):
        """
        Return a list of completion types that this CompletionLoader can
        return.

        """
        return [CT_LIST]

    @classmethod
    def view_check(cls, view):
        if super().view_check(view):
            return get_completion_source_enabled_setting('List',
                                                         'Include')

    def load_completions(self, **kwargs):
        """
        Populate self.completions with Lists from an Include file.

        """
        ring_file = get_ring_file(self.file_path)
        self.completions = set(
            [(x + '\tInclude', x) for x in ring_file.get_defined_lists()])