"""
Snake game - singleplayer, or multiplayer up to 3 players at total.
Author - David Z., I. year of Bc. MFF UK Charles university, informatics.
Programování I. winter semester 2020/2021

"""

import sys
import os
import random
import pygame
from pygame import Rect, Vector2, Color

def fix_path(p):
    return os.path.dirname(os.path.realpath(__file__)) + "/" + p        # help with links

def create_board(d = 857, i = 0, j = 0):            # draw squares -- board, d is exactly the number needed to draw them all
    for n in range(d):
        pygame.draw.polygon(screen, (255,255,255), [(i,j), (i+40,j), (i+40,j+40), (i,j+40), (i,j)], width=3)
        if i != width:
            i = i + 40
        else:
            j = j + 40
            i = -i

"""
                        ----
        APPLE
                        ----
"""

apple_position = []             # coordinates of all apples
class Apple:
    def __init__(self):
        self.alive = False
        self.apple = pygame.image.load(fix_path("image - APPLE.jpg"))
        self.applerect = self.apple.get_rect()

    def draw(self):
        if self.alive == False:                         # no apple is on the board, apple is yet to be drawn; random coordinates, which do not collide with other objects - such as snakes or poison jars
            while True:
                self.x = random.randint(20,940)
                self.y = random.randint(20,700)
                coord = [self.x, self.y]
                if self.x / 20 in [i for i in range(50) if i%2 !=0] and self.y / 20 in [i for i in range(50) if i%2 !=0]:       # math - we need our images to possibly spawn 40px to 40px from (0,0) each 40px
                    if (coord not in poison_position) and ([coord[0]-20, coord[1]-20] not in snake_position1) and ([coord[0]-20, coord[1]-20] not in snake_position2) and ([coord[0]-20, coord[1]-20] not in snake_position3):          # so we don't spawn it on poison or snake
                        self.applerect.center = self.x, self.y
                        self.alive = True
                        apple_position.append(coord)    # adding to a separate list, later on to determine if the head of snake eats the apple
                        break
        if self.alive == True:                          # apple is fresh and ready to eat, we just have to draw it on the screen
            screen.blit(self.apple,self.applerect)

    def check_if_apple_eaten(self):                     # if a snake ate an apple, we switch alive to False, so we will generate a new one
        global a
        if a == 0:
            self.alive = False
            a = 1

"""
                        ----
        POISON
                        ----
"""

i = 0                           # counting every frame; to add more poison over time
poison_position = []            # coords of all poison jars
class Poison:
    def __init__(self):
        self.alive = False
        self.poison = pygame.image.load(fix_path("image - POISON.jpg"))
        self.poisonrect = self.poison.get_rect()

    def draw(self):
        global i
        if i == 60:                                     # after a solid time we will add a new poison jar
            i = 0
            self.alive = False
        if self.alive == False:                         # snake ate a poison jar, what to do?
            while True:
                self.x = random.randint(20,940)
                self.y = random.randint(20,700)         # center of the poison jar picture - to fit a random square in my board
                coord = [self.x, self.y]
                if self.x / 20 in [i for i in range(50) if i%2 !=0] and self.y / 20 in [i for i in range(50) if i%2 !=0]:       # math - same as apple
                    if (coord not in apple_position) and ([coord[0]-20, coord[1]-20] not in snake_position1) and ([coord[0]-20, coord[1]-20] not in snake_position2) and ([coord[0]-20, coord[1]-20] not in snake_position3):                                                                             # check - we don't want to spawn poison onto an already existing apple
                        self.alive = True
                        poison_position.append(coord)
                        break
        if self.alive == True:                          # snake haven't eaten a poison jar yet - we keep drawing it
            for c in poison_position:
                self.x = c[0]
                self.y = c[1]
                self.poisonrect.center = self.x, self.y
                screen.blit(self.poison,self.poisonrect)

"""
                                            ----
            SNAKE FOR PLAYER 1
                                            ----
"""

length1 = 1                 # length of the snake
a = 1                       # apple reaction - we switch it to 0 if snake eats the apple, apple will redraw at the new position and a will return to 1
p = 1                       # poison reaction - snake is alive = 1, snake is dead = 0
queue1 = []                 # protection of snake comitting suicide
snake_road1 = []            # all coord of the snake's route
snake_position1 = []        # actual coords of the body - check if head collides with the rest of its body
class Snake1:
    def __init__(self, coord_x, coord_y):
        self.head = [coord_x, coord_y]
        self.head_checking = [coord_x + 20, coord_y + 20]
        self.alive = True

    def draw(self):                             # we take last couple of coords of the snake's road equivalent to snake's head and apples he has eaten so far, we draw it
        snake_road1.append(self.head)
        if length1 > 0:
            for i in range(1, length1 + 1):
                x = snake_road1[-i]
                pygame.draw.rect(screen, (255,0,0), (x[0], x[1], 40,40))
                snake_position1.append(x)
        while len(snake_position1) != length1:  # we get the momentarily coordination of the snake to a new list
            snake_position1.pop(0)

    def movement(self):                         # function to regulate direction
        global player1_movement_x, player1_movement_y, queue1
        if len(queue1) == 0:
            pass
        else:
            if queue1[0] == "UP":
                if player1_movement_y == 40:            # QUEUE - double clicking could make our snake kill itself - he would turn on the same line
                    queue1.pop(0)
                else:
                    player1_movement_y = -40
                    player1_movement_x = 0
                    queue1.pop(0)
            elif queue1[0] == "DOWN":
                if player1_movement_y == -40:
                    queue1.pop(0)
                else:
                    player1_movement_y = +40
                    player1_movement_x = 0
                    queue1.pop(0)
            elif queue1[0] == "RIGHT":
                if player1_movement_x == -40:
                    queue1.pop(0)
                else:
                    player1_movement_x = 40
                    player1_movement_y = 0
                    queue1.pop(0)
            elif queue1[0] == "LEFT":
                if player1_movement_x == 40:
                    queue1.pop(0)
                else:
                    player1_movement_x = -40
                    player1_movement_y = 0
                    queue1.pop(0)

    def eat_apple(self):                                    # if apple is eaten, then we add to the score, length of our snake and we reset apple
        global a, player1_score, length1
        if self.head_checking in apple_position:
            apple_position.remove(self.head_checking)
            player1_score += 1
            length1 += 1
            a = 0

    def eat_poison(self):                                   # if poison is eaten, then we remove poison and we kill our snake
        global p
        if self.head_checking in poison_position:
            poison_position.remove(self.head_checking)
            p = 0

    def collision_border(self):                             # collision with border or with snake itself kills our snake
        global p
        if self.head_checking[0] > 960 or self.head_checking[0] < 0:
            p = 0
        if self.head_checking[1] > 720 or self.head_checking[1] < 0:
            p = 0
        if snake_position1.count(self.head) > 1:
            p = 0

    def collision_with_players(self):                       # collision with players - snake kills itself
        global p
        if self.head in snake_position2:
            p = 0
        if self.head in snake_position3:
            p = 0

"""
                                            ----
            SNAKE FOR PLAYER 2
                                            ----
"""

length2 = 1                 # length of the snake
a = 1                       # apple reaction
q = 1                       # poison reaction - snake is alive
queue2 = []                 # protection of snake comitting suicide
snake_road2 = []            # all coord of the snake's route
snake_position2 = []        # actual coords of the body - check if head collides with the rest of its body
class Snake2:
    def __init__(self, coord_x, coord_y):
        self.head = [coord_x, coord_y]
        self.head_checking = [coord_x + 20, coord_y + 20]
        self.alive = True

    def draw(self):                               # we take last coords of the snake's road and we draw exactly how many apples we have eaten
        snake_road2.append(self.head)
        if length2 > 0:
            for i in range(1, length2 + 1):
                x = snake_road2[-i]
                pygame.draw.rect(screen, (0,255,0), (x[0], x[1], 40,40))
                snake_position2.append(x)
        while len(snake_position2) != length2:    # we get the momentarily coordination of the snake to a new list
            snake_position2.pop(0)

    def movement(self):                           # function to regulate direction
        global player2_movement_x, player2_movement_y, queue2
        if len(queue2) == 0:
            pass
        else:
            if queue2[0] == "UP":
                if player2_movement_y == 40:            # QUEUE - double clicking could make our snake kill itself
                    queue2.pop(0)
                else:
                    player2_movement_y = -40
                    player2_movement_x = 0
                    queue2.pop(0)
            elif queue2[0] == "DOWN":
                if player2_movement_y == -40:
                    queue2.pop(0)
                else:
                    player2_movement_y = +40
                    player2_movement_x = 0
                    queue2.pop(0)
            elif queue2[0] == "RIGHT":
                if player2_movement_x == -40:
                    queue2.pop(0)
                else:
                    player2_movement_x = 40
                    player2_movement_y = 0
                    queue2.pop(0)
            elif queue2[0] == "LEFT":
                if player2_movement_x == 40:
                    queue2.pop(0)
                else:
                    player2_movement_x = -40
                    player2_movement_y = 0
                    queue2.pop(0)

    def eat_apple(self):                                    # if apple is eaten, then we add to the score, length of our snake and we reset apple
        global a, player2_score, length2
        if self.head_checking in apple_position:
            apple_position.remove(self.head_checking)
            player2_score += 1
            length2 += 1
            a = 0

    def eat_poison(self):                                   # if poison is eaten, then we remove poison and we kill our snake
        global q
        if self.head_checking in poison_position:
            poison_position.remove(self.head_checking)
            q = 0

    def collision_border(self):                             # collision with border or with snake itself kills our snake
        global q
        if self.head_checking[0] > 960 or self.head_checking[0] < 0:
            q = 0
        if self.head_checking[1] > 720 or self.head_checking[1] < 0:
            q = 0
        if snake_position2.count(self.head) > 1:
            q = 0

    def collision_with_players(self):
        global q
        if self.head in snake_position1:
            q = 0
        if self.head in snake_position3:
            q = 0

"""
                                            ----
            SNAKE FOR PLAYER 3
                                            ----
"""

length3 = 1                 # length of the snake
a = 1                       # apple reaction
r = 1                       # poison reaction - snake is alive
queue3 = []                 # protection of snake comitting suicide
snake_road3 = []            # all coord of the snake's route
snake_position3 = []        # actual coords of the body - check if head collides with the rest of its body
class Snake3:
    def __init__(self, coord_x, coord_y):
        self.head = [coord_x, coord_y]
        self.head_checking = [coord_x + 20, coord_y + 20]
        self.alive = True

    def draw(self):                               # we take last coords of the snake's road and we draw exactly how many apples we have eaten
        snake_road3.append(self.head)
        if length3 > 0:
            for i in range(1, length3 + 1):
                x = snake_road3[-i]
                pygame.draw.rect(screen, (0,0,255), (x[0], x[1], 40,40))
                snake_position3.append(x)
        while len(snake_position3) != length3:    # we get the momentarily coordination of the snake to a new list
            snake_position3.pop(0)

    def movement(self):                           # function to regulate direction
        global player3_movement_x, player3_movement_y, queue3
        if len(queue3) == 0:
            pass
        else:
            if queue3[0] == "UP":
                if player3_movement_y == 40:            # QUEUE - double clicking could make our snake kill itself
                    queue3.pop(0)
                else:
                    player3_movement_y = -40
                    player3_movement_x = 0
                    queue3.pop(0)
            elif queue3[0] == "DOWN":
                if player3_movement_y == -40:
                    queue3.pop(0)
                else:
                    player3_movement_y = +40
                    player3_movement_x = 0
                    queue3.pop(0)
            elif queue3[0] == "RIGHT":
                if player3_movement_x == -40:
                    queue3.pop(0)
                else:
                    player3_movement_x = 40
                    player3_movement_y = 0
                    queue3.pop(0)
            elif queue3[0] == "LEFT":
                if player3_movement_x == 40:
                    queue3.pop(0)
                else:
                    player3_movement_x = -40
                    player3_movement_y = 0
                    queue3.pop(0)

    def eat_apple(self):                                    # if apple is eaten, then we add to the score, length of our snake and we reset apple
        global a, player3_score, length3
        if self.head_checking in apple_position:
            apple_position.remove(self.head_checking)
            player3_score += 1
            length3 += 1
            a = 0

    def eat_poison(self):                                   # if poison is eaten, then we remove poison and we kill our snake
        global r
        if self.head_checking in poison_position:
            poison_position.remove(self.head_checking)
            r = 0

    def collision_border(self):                             # collision with border or with snake itself kills our snake
        global r
        if self.head_checking[0] > 960 or self.head_checking[0] < 0:
            r = 0
        if self.head_checking[1] > 720 or self.head_checking[1] < 0:
            r = 0
        if snake_position3.count(self.head) > 1:
            r = 0

    def collision_with_players(self):
        global r
        if self.head in snake_position1:
            r = 0
        if self.head in snake_position2:
            r = 0

"""
                                        -----
            START END CODE SECTION
                                        -----
"""

def controls_player1():
    if event.type == pygame.KEYDOWN:
        if event.key == pygame.K_LEFT:          # movement with keyboard - arrows - for Player number 1
            queue1.append("LEFT")
        if event.key == pygame.K_RIGHT:
            queue1.append("RIGHT")
        if event.key == pygame.K_UP:
            queue1.append("UP")
        if event.key == pygame.K_DOWN:
            queue1.append("DOWN")

def controls_player2():
    if event.type == pygame.KEYDOWN:
        if event.key == pygame.K_a:          # movement with keyboard - WSAD - for Player number 2
            queue2.append("LEFT")
        if event.key == pygame.K_d:
            queue2.append("RIGHT")
        if event.key == pygame.K_w:
            queue2.append("UP")
        if event.key == pygame.K_s:
            queue2.append("DOWN")

def controls_player3():
    if event.type == pygame.KEYDOWN:
        if event.key == pygame.K_j:          # movement with keyboard - IKJL - for Player number 3
            queue3.append("LEFT")
        if event.key == pygame.K_l:
            queue3.append("RIGHT")
        if event.key == pygame.K_i:
            queue3.append("UP")
        if event.key == pygame.K_k:
            queue3.append("DOWN")

"""
                        ----
        MAIN
                        ----
"""
player1_score = 0
player2_score = 0           # this will represent the results at the end
player3_score = 0

print("Zvolte si počet hráčů:  (1 = jeden hráč, 2 = 2 hráči, 3 = 3hráči)")      # decide how many players are playing; 1 up to 3
while True:
    no_players = int(input())
    if no_players == 1 or no_players == 2 or no_players == 3:
        break
    else:
        print("Znovu a lépe :)")

if no_players == 1:             # we set up the game for only one player
    q = 0
    r = 0
    pygame.init()
    pygame.display.set_caption("Snake Game - multiplayer - by DZ")
    width = 960
    height = 720
    screen = pygame.display.set_mode((width, height))
    clock = pygame.time.Clock()
    A = Apple()                             # init of Apple
    P = Poison()

    player1_coord_x = 200                   # starting coordinations of 1. snake
    player1_coord_y = 200
    player1_movement_x = 40                 # starting speed and direction of 1. snake
    player1_movement_y = 0

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                sys.exit()
            controls_player1()

        screen.fill(pygame.Color(0, 0, 0))              # black screen, then our squares - board
        create_board()

        if p == 1:                                      # if our snake is alive, then we init him and his core functions
            S1 = Snake1(player1_coord_x, player1_coord_y)
            S1.draw()
            S1.movement()
            S1.eat_apple()
            S1.eat_poison()
            S1.collision_border()
        else:                                           # if the snake dies, his current position is set to none
            snake_position1.clear()
        S1.collision_with_players()
        A.check_if_apple_eaten()                        # core functions of Apple and Poison classes
        A.draw()
        P.draw()
        player1_coord_x = player1_coord_x + player1_movement_x                  # OUR MAIN SOURCE OF MOVEMENT !!!
        player1_coord_y = player1_coord_y + player1_movement_y
        i += 1
        if p == 0 and q == 0 and r == 0:
            break
        pygame.display.flip()
        clock.tick(5)                                                           # OUR MAIN SOURCE OF MOVEMENT !!!

if no_players == 2:                 # we set the game for 2 players
    r = 0
    pygame.init()
    pygame.display.set_caption("Snake Game - multiplayer - by DZ")
    width = 960
    height = 720
    screen = pygame.display.set_mode((width, height))
    clock = pygame.time.Clock()
    A = Apple()                     # init of Apple
    P = Poison()

    player1_coord_x = 200                   # starting coordinations of 1. snake
    player1_coord_y = 200
    player1_movement_x = 40                 # starting speed and direction of 1. snake
    player1_movement_y = 0

    player2_coord_x = 600
    player2_coord_y = 600                   # starting coordinations of 2. snake
    player2_movement_x = 0                  # starting speed and direction of 2. snake
    player2_movement_y = -40

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                sys.exit()
            controls_player1()
            controls_player2()
        screen.fill(pygame.Color(0, 0, 0))              # black screen, then our squares - board
        create_board()

        if p == 1:                                      # if 1. snake is alive, then we init him and his core functions
            S1 = Snake1(player1_coord_x, player1_coord_y)
            S1.draw()
            S1.movement()
            S1.eat_apple()
            S1.eat_poison()
            S1.collision_border()
        else:                                           # if the snake dies, his current position is set to none
            snake_position1.clear()
        if q == 1:                                      # if 2. snake is alive, then we init him and his core functions
            S2 = Snake2(player2_coord_x, player2_coord_y)
            S2.draw()
            S2.movement()
            S2.eat_apple()
            S2.eat_poison()
            S2.collision_border()
        else:                                           # if the snake dies, his current position is set to none
            snake_position2.clear()

        S1.collision_with_players()
        S2.collision_with_players()
        A.check_if_apple_eaten()                        # core functions of Apple and Poison classes
        A.draw()
        P.draw()
        player1_coord_x = player1_coord_x + player1_movement_x                  # OUR MAIN SOURCE OF MOVEMENT !!!
        player1_coord_y = player1_coord_y + player1_movement_y
        player2_coord_x = player2_coord_x + player2_movement_x
        player2_coord_y = player2_coord_y + player2_movement_y
        i += 1
        if p == 0 and q == 0 and r == 0:
            break
        pygame.display.flip()
        clock.tick(5)

if no_players == 3:                     # we set the game for 3 players
    pygame.init()
    pygame.display.set_caption("Snake Game - multiplayer - by DZ")
    width = 960
    height = 720
    screen = pygame.display.set_mode((width, height))
    clock = pygame.time.Clock()
    A = Apple()                     # init of Apple
    P = Poison()

    player1_coord_x = 200                   # starting coordinations of 1. snake
    player1_coord_y = 200
    player1_movement_x = 40                 # starting speed and direction of 1. snake
    player1_movement_y = 0

    player2_coord_x = 600                   # starting coordinations of 2. snake
    player2_coord_y = 600
    player2_movement_x = 0                  # starting speed and direction of 2. snake
    player2_movement_y = -40

    player3_coord_x = 480                   # starting coordinations of 3. snake
    player3_coord_y = 80
    player3_movement_x = -40                # startin speed and direction of 3. snake
    player3_movement_y = 0

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                sys.exit()
            controls_player1()
            controls_player2()
            controls_player3()
        screen.fill(pygame.Color(0, 0, 0))              # black screen, then we draw our board - squares
        create_board()

        if p == 1:                                      # if 1. snake is alive, then we init him and his core functions
            S1 = Snake1(player1_coord_x, player1_coord_y)
            S1.draw()
            S1.movement()
            S1.eat_apple()
            S1.eat_poison()
            S1.collision_border()
        else:                                           # if the snake dies, his current position is set to none
            snake_position1.clear()
        if q == 1:                                      # if 2. snake is alive, then we init him and his core functions
            S2 = Snake2(player2_coord_x, player2_coord_y)
            S2.draw()
            S2.movement()
            S2.eat_apple()
            S2.eat_poison()
            S2.collision_border()
        else:                                           # if the snake dies, his current position is set to none
            snake_position2.clear()
        if r == 1:                                      # if 3. snake is alive, then we init him and his core functions
            S3 = Snake3(player3_coord_x, player3_coord_y)
            S3.draw()
            S3.movement()
            S3.eat_apple()
            S3.eat_poison()
            S3.collision_border()
        else:                                           # if the snake dies, his current position is set to none
            snake_position3.clear()

        S1.collision_with_players()
        S2.collision_with_players()
        S3.collision_with_players()
        A.check_if_apple_eaten()                        # core functions of Apple and Poison classes
        A.draw()
        P.draw()
        player1_coord_x = player1_coord_x + player1_movement_x                  # OUR MAIN SOURCE OF MOVEMENT !!!
        player1_coord_y = player1_coord_y + player1_movement_y
        player2_coord_x = player2_coord_x + player2_movement_x
        player2_coord_y = player2_coord_y + player2_movement_y
        player3_coord_x = player3_coord_x + player3_movement_x                  # OUR MAIN SOURCE OF MOVEMENT !!!
        player3_coord_y = player3_coord_y + player3_movement_y
        i += 1
        if p == 0 and q == 0 and r == 0:
            break
        pygame.display.flip()
        clock.tick(5)

if p == 0 and q == 0 and r == 0:                # when the game ends, i created custom fireworks to celebrate with results
    pygame.init()
    width = 960
    height = 720
    screen = pygame.display.set_mode((width, height))
    clock = pygame.time.Clock()
    coord_x = 160
    coord_y = 680
    i = 0
    def particles(k,l,m,n):
        pygame.draw.polygon(screen, (255,0,0), [(k,l), (m,n)], width=3)

    player1_score = str(player1_score)
    player2_score = str(player2_score)          # preparations for printing results on the screen
    player3_score = str(player3_score)
    myfont = pygame.font.SysFont("Cosmic Sans MS", 32)
    textsurface1 = myfont.render("Player1: " + player1_score, False, (255,0,0))
    textsurface2 = myfont.render("Player2: " + player2_score, False, (255,0,0))
    textsurface3 = myfont.render("Player3: " + player3_score, False, (255,0,0))

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                sys.exit()
        screen.fill(pygame.Color(0, 0, 0))

        screen.blit(textsurface1, (520,440))
        screen.blit(textsurface2, (520,480))    # printing results on the screen
        screen.blit(textsurface3, (520,520))

        if i != 10:
            if coord_y != 260:
                particles(coord_x,coord_y, coord_x, coord_y-30)                     # fireworks itself, custom made
                coord_y -= 60
            else:
                particles(coord_x, coord_y-30, coord_x, coord_y-60)
                particles(coord_x-30, coord_y-90, coord_x-60, coord_y-120)
                particles(coord_x+30, coord_y-90, coord_x+60, coord_y-120)

                particles(coord_x-30, coord_y-30, coord_x-60, coord_y-60)
                particles(coord_x+30, coord_y-30, coord_x+60, coord_y-60)
                particles(coord_x-90, coord_y-80, coord_x-120, coord_y-80)
                particles(coord_x+90, coord_y-80, coord_x+120, coord_y-80)

                particles(coord_x-30, coord_y, coord_x-60, coord_y)
                particles(coord_x-90,coord_y,coord_x-120,coord_y+15)
                particles(coord_x+30,coord_y,coord_x+60,coord_y)
                particles(coord_x+90,coord_y,coord_x+120,coord_y+15)

                particles(coord_x-30, coord_y+30, coord_x-60, coord_y+60)
                particles(coord_x+30, coord_y+30, coord_x+60, coord_y+60)
                particles(coord_x-70, coord_y+90, coord_x-80, coord_y+120)
                particles(coord_x+70, coord_y+90, coord_x+80, coord_y+120)
        else:
            i = 0
            coord_x = 160
            coord_y = 680
        i += 1
        pygame.display.flip()
        clock.tick(5)

"""
Hopefully you enjoyed playing my version of the snake game!
"""


from math import sqrt
from random import uniform

import numpy as np
import openmc
import pytest

def test_time_filter_basics():
    f = openmc.TimeFilter([1.0, 2.0, 5.0, 9.0])
    np.testing.assert_allclose(f.bins[0], [1.0, 2.0])
    np.testing.assert_allclose(f.bins[1], [2.0, 5.0])
    np.testing.assert_allclose(f.bins[-1], [5.0, 9.0])
    assert len(f.bins) == 3

    # Make sure __repr__ works
    repr(f)

    # to_xml_element()
    elem = f.to_xml_element()
    assert elem.tag == 'filter'
    assert elem.attrib['type'] == 'time'

def time(particle, distance, E):
    """Return the time it takes a particle at a given energy to travel a certain
    distance"""
    if particle == 'neutron':
        mass = 939.56542052e6  # eV/c²
    elif particle == 'photon':
        mass = 0.0

    # Calculate speed via v = c * sqrt(1 - γ^-2)
    inv_gamma = mass / (E + mass)
    velocity = 2.99792458e10 * sqrt(1 - inv_gamma * inv_gamma)  # cm/s
    return distance / velocity

@pytest.fixture(params=['neutron', 'photon'])
def model(request):
    # Select random sphere radius, source position, and source energy
    r = uniform(0., 2.)
    x = uniform(2., 10.)
    E = uniform(0., 20.0e6)

    # Create model
    model = openmc.Model()
    mat = openmc.Material()
    mat.add_nuclide('Zr90', 1.0)
    mat.set_density('g/cm3', 1.0)
    model.materials.append(mat)
    inner_sphere = openmc.Sphere(r=r)
    outer_sphere = openmc.Sphere(r=10.0, boundary_type='vacuum')
    center_cell = openmc.Cell(fill=mat, region=-inner_sphere)
    outer_void = openmc.Cell(region=+inner_sphere & -outer_sphere)
    model.geometry = openmc.Geometry([center_cell, outer_void])
    model.settings = openmc.Settings()
    model.settings.run_mode = 'fixed source'
    model.settings.particles = 1000
    model.settings.batches = 20
    particle = request.param
    model.settings.source = openmc.Source(
        space=openmc.stats.Point((x, 0., 0.)),
        angle=openmc.stats.Monodirectional([-1., 0., 0.]),
        energy=openmc.stats.Discrete([E], [1.0]),
        particle=particle
    )

    # Calculate time it will take neutrons to reach sphere
    t0 = time(particle, x - r, E)

    # Create tally with time filter
    tally = openmc.Tally()
    tally.filters = [openmc.TimeFilter([0.0, t0, 2*t0])]
    tally.scores = ['total']
    model.tallies.append(tally)
    return model

@pytest.fixture(params=['neutron', 'photon'])
def model_surf(request):
    # Select random distance and source energy
    x = uniform(50., 100.)
    E = uniform(0., 20.0e6)

    # Create model
    model = openmc.Model()
    mat = openmc.Material()
    mat.add_nuclide('Zr90', 1.0)
    mat.set_density('g/cm3', 1.0)
    model.materials.append(mat)
    left = openmc.XPlane(-1., boundary_type='vacuum')
    black_surface = openmc.XPlane(x, boundary_type='vacuum')
    right = openmc.XPlane(x + 1)
    void_cell = openmc.Cell(region=+left & -black_surface)
    black_cell = openmc.Cell(region=+black_surface & -right)
    model.geometry = openmc.Geometry([void_cell, black_cell])
    model.settings = openmc.Settings()
    model.settings.run_mode = 'fixed source'
    model.settings.particles = 1000
    model.settings.batches = 20
    particle = request.param
    model.settings.source = openmc.Source(
        space=openmc.stats.Point((0., 0., 0.)),
        angle=openmc.stats.Monodirectional([1., 0., 0.]),
        energy=openmc.stats.Discrete([E], [1.0]),
        particle=particle
    )

    # Calculate time it will take neutrons to reach purely-absorbing surface
    t0 = time(particle, x, E)

    # Create tally with surface and time filters
    tally = openmc.Tally()
    tally.filters = [
        openmc.SurfaceFilter([black_surface]),
        openmc.TimeFilter([0.0, t0*0.999, t0*1.001, 100.0])
    ]
    tally.scores = ['current']
    model.tallies.append(tally)
    return model

def test_time_filter_volume(model, run_in_tmpdir):
    sp_filename = model.run()
    with openmc.StatePoint(sp_filename) as sp:
        t = sp.tallies[model.tallies[0].id]
        values = t.mean.ravel()

        # Before t0, the reaction rate should be zero
        assert values[0] == pytest.approx(0.0)

        # After t0, the reaction rate should be positive
        assert values[1] > 0.0

def test_time_filter_surface(model_surf, run_in_tmpdir):
    sp_filename = model_surf.run()
    with openmc.StatePoint(sp_filename) as sp:
        t = sp.tallies[model_surf.tallies[0].id]
        values = t.mean.ravel()
        print(values)

        # Before t0-ε, the current should be zero
        assert values[0] == 0.0

        # Between t0-ε and t0+ε, the current should be one
        assert values[1] == 1.0

        # After t0+ε, the current should be zero
        assert values[2] == 0.0


from sklearn.metrics import classification_report, f1_score, precision_score, recall_score, confusion_matrix, accuracy_score
import pandas as pd

def write_cm_report(Y, classes, predict, filename):
	print('Confusion Matrix')
	cm = confusion_matrix(classes, predict)
	file = open(filename, 'a+')
	file.write('cm \n')
	file.write(str(cm))
	file.write('\n')
	average_type = ['micro','macro']
	for i in average_type:
	    score = f1_score(classes, predict, average=i)
	    precision = precision_score(classes, predict, average=i)
	    recall = recall_score(classes, predict, average=i)
	    file.write('f1 score: {:.4f}, precision: {:.4f}, recall: {:.4f}'.format(score, precision, recall))
	    file.write('\n')
	file.write("Accuracy: %f" % accuracy_score(classes, predict))
	file.write('\n')

	nomes_classes = []
	for i in pd.DataFrame(Y.groupby('name')['name'].nunique().reset_index(name="unique"))[
	    'name']:
	    nomes_classes.append(str(i))
	file.write(str(classification_report(classes, predict, target_names=nomes_classes)))
	file.write('\n')
	file.close()


#!/usr/bin/env python
#auth
#version 0.1
#date 20181211
#mail:

import logging

def log(info):
    logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
                    datefmt='%a, %d %b %Y %H:%M:%S',
                    filename='/data/Python/Release/logs/release.log',
                    filemode='a')
    logging.info(info)

if __name__ == "__main__":
    log('This is a write test line')


import socket, serial
from scipy.fftpack import fft, ifft, fftshift
from scipy import signal
import matplotlib.pyplot as plt
import numpy as np
from tkinter import *
from PIL import ImageTk, Image
import os, time

class Window(Frame):

    def __init__(self, s, a, master=None):
        Frame.__init__(self, master)
        self.master = master
        self.init_window()
        self.sock = s
        self.data = []
        self.ard = a
        self.fft_size = 2048

    #Creation of init_window
    def init_window(self):
        self.master.title("TESTING GUI")
        self.pack(fill=BOTH, expand=1)

        def run_motors():
            """Runs motors based on inputs given"""
            def format_data(e_d, e_m, a_d, a_m,fn,axis):
                """Formats a message for send_data function"""
                # Message Protocol "99, Ele dir, Ele deg, Azi dir, Azi deg, function, axis, 99"
                # Message Struct   "99,    X   ,    XX  ,    X   ,   XX   ,    X    ,  X  , 99"
                e_m = int(e_m)
                a_m = int(a_m)

                # pad a zero to e_m and a_m if less than two digits
                if(e_m < 10):
                    e_m_str = "0" + str(e_m)
                else:
                    e_m_str = str(e_m)

                if(a_m < 10):
                    a_m_str = "0" + str(a_m)
                else:
                    a_m_str = str(a_m)

                formated_msg = "99"+str(e_d)+(e_m_str)+str(a_d)+(a_m_str)+str(fn)+str(axis)+"99"
                return formated_msg

            def send_data(data):
                    """Sends data to arduino"""
                    self.ard.flush
                    sys.stdout.flush()
                    strTemp = str(data)
                    strTemp_encoded = strTemp.encode("ascii")
                    print("Python value sent: ")
                    print(strTemp)
                    not_confirmed = True
                    i = 0
                    while not_confirmed:
                            self.ard.write(strTemp_encoded)
                            time.sleep(.2)
                            received_str = self.ard.readline().decode()
                            print(received_str)
                            if (received_str == strTemp):
                                    not_confirmed = False
                                    exit()
                                    self.ard.flush
                                    sys.stdout.flush()
                            i = 1 + i
                            if i > 10:
                                    exit()  #This will need to send a message to the GUI about somehtign not working
                                    self.ard.flush
                                    sys.stdout.flush()
            def drive_stepper(e_d, e_m, a_d, a_m, axis):
                """Drives steppers with input"""
                send_string = format_data(e_d, e_m, a_d, a_m,0,axis)
                send_data(send_string)

            az = az_val.get()
            el = el_val.get()
            # send the direction to motors
            try: # try except to catch invalid input
                if(el < 0):
                    el_dir = 0
                else:
                    el_dir = 1
                if(az < 0):
                    az_dir = 0
                else:
                    az_dir = 1
                print('Az {} El: {}'.format(str(az),str(el)))
                drive_stepper(el, el_dir, az, az_dir, 0) # 0 is azimuth, 1 is elevation
                drive_stepper(el, el_dir, az, az_dir, 1)
            except:
                pass

        title_label = Label(self, text="TEAM ROGUE", bg="black",fg = "white",font=("Helvetica", 14))
        title_label.place(x = 100, y = 0)

        RxButton = Button(self, text="Rx Data",command=self.Rx_data)
        RxButton.place(x=50,y=50)

        plotdataButton = Button(self, text="Plot Data",command=self.plot_data)
        plotdataButton.place(x=50,y=100)

        plotFFTButton = Button(self, text="Plot FFT of Data",command=self.plot_fft_data)
        plotFFTButton.place(x=50,y=150)

        plotWindowedButton = Button(self, text="Plot Windowed Data",command=self.plot_windowed_data)
        plotWindowedButton.place(x=50,y=200)

        az_label = Label(self, text="Az")
        az_label.place(x = 50, y = 250)
        az_val = Entry(self, width = 10)
        az_val.place(x = 80, y = 250)

        el_label = Label(self, text="El")
        el_label.place(x = 50, y = 280)
        el_val = Entry(self, width = 10)
        el_val.place(x = 80, y = 280)

        motorButton = Button(self, text="Move",command=run_motors)
        motorButton.place(x = 150,y = 265)

        motorReadButton = Button(self, text="Read Position",command=self.read_position)
        motorReadButton.place(x = 200,y = 265)

        binButton = Button(self, text="Select Bin",command=self.bin_select)
        binButton.place(x=50,y=310)

        phaseButton = Button(self, text="Extract Phase",command=self.phase_extraction)
        phaseButton.place(x=50,y=350)

        quitButton = Button(self, text="Exit",command=self.force_quit)
        quitButton.place(x=50, y=400)

    def force_quit(self):
        exit()

    def Rx_data(self):
        """Receives data to plot"""
        self.data.clear()
        n=0
        data_rx = s.recv(BUFFER_SIZE)
        while n < 1024:
            try:
                #print(data_rx[n])
                if(n == 0):
                    self.data.append(0)
                else:
                    self.data.append(float(data_rx[n]))
            except:
                pass
            n+=1
        print('Data Received.')
        #print(self.data)   # for testing

    def plot_data(self):
        """Plots data as function of time"""
        N = 1024
        T=1/5000000
        x = np.linspace(0.0,N*T,N)
        plt.plot(x,self.data)
        plt.show()

    def plot_windowed_data(self):
        """Plots data after being windowed"""
        N = 1024
        #sample spacing
        T = 1.0/N
        x = np.linspace(0.0,N*T,N)
        y = np.sin(50.0*2.0*np.pi*x) + 0.5*np.sin(80.0*2.0*np.pi*x)
        y = self.data
        yf = fft(y)
        #create window on data
        w = np.hanning(N)
        ywf = fft((y*w))
        #xf = np.linspace(0.0,1.0/(2.0*T), N/2)
        xf = np.linspace(0,1,len(ywf))
        #plot the data
        #plt.semilogy(xf[1:N//2], 2.0/N * np.abs(yf[1:N//2]), '-b')
        plt.semilogy(xf[1:N//2], 2.0/N * np.abs(ywf[1:N//2]), '-r')
        plt.legend(['FFT', 'FFT w. window'])
        plt.xlabel("Normalized frequency")
        plt.grid()
        plt.show()

    def plot_fft_data(self):
        """Plots data after the FFT"""
        N = 1024
        #sample spacing
        T = 1.0/N
        x = np.linspace(2400,2480,N)
        y = self.data
        yf = fft(y)
        #create window on data
        xf = np.linspace(0, 1, len(yf))
        #plot the data
        plt.semilogy(xf[1:N//2], 2.0/N * np.abs(yf[1:N//2]), '-b')
        plt.legend(['FFT'])
        plt.xlabel("Normalized frequency")
        plt.grid()
        plt.show()

    def phase_extraction(self):
        """Extract phases from selected bin"""
        def phasepopupmsg(phase_rad, phase_deg):
            """Pop up msg shows motor actual position"""
            Ppopup = Tk()
            Ppopup.geometry("300x150")
            Ppopup.wm_title("Phase Extraction")
            phase_msg = "Extracted Phase: " + str(phase_rad) + " rad"
            phase_msg2 = "Extracted Phase: " + str(phase_deg) + " deg"
            label = Label(Ppopup, text=phase_msg)
            label.pack(side="top", fill="x", pady=10)
            label2 = Label(Ppopup, text=phase_msg2)
            label2.pack(side="top", fill="x", pady=10)
            B = Button(Ppopup, text="Exit", command = Ppopup.destroy)
            B.pack()
        y = fft(self.data)
        mag_ar = np.absolute(y)
        mag_ar[0] = 0
        bin_val = np.argmax(mag_ar)
        phase_data_rad = np.angle(y[bin_val])
        phase_data_deg = np.angle(y[bin_val]) * 57.2958
        phasepopupmsg(phase_data_rad, phase_data_deg)

    def bin_select(self):
        """selects FFT bin for processing""" #------------currently not used; part of p_e()----
        def popupmsg(msg_a,msg_b):
            """Creates pop up message"""
            popup = Tk()
            popup.wm_title("Bin Selection")
            popup.geometry("300x150")
            labela = Label(popup, text=msg_a)#, font=NORM_FONT)
            labela.pack(side="top", fill="x", pady=10)
            labelb = Label(popup, text=msg_b)#, font=NORM_FONT)
            labelb.pack(side="top", fill="x", pady=10)
            B1 = Button(popup, text="Exit", command = popup.destroy)
            B1.pack()
        y = self.data
        #yf = fft(y,1024)
        w = np.hanning(1024)
        yf = fft((y))
        mag_ar = np.absolute(yf)
        mag_ar[0] = 0
        #print(mag_ar)  # for testing
        bin_val = np.argmax(mag_ar)
        selected_freq = bin_val * 2444 + 2419000000
        msg1 = "Selected Bin: " +str(bin_val)
        msg2 ="Selectred frequency: {} Hz".format(selected_freq)
        popupmsg(msg1,msg2)

    def read_position(self):
        """Reads position back from motors"""
        def motorpopupmsg(a,e):
            """Pop up msg shows motor actual position"""
            popup2 = Tk()
            popup2.geometry("200x150")
            popup2.wm_title("Bin Selection")
            az_msg = "Azimuth Degrees: " + str(a)
            el_msg = "Elevation Degrees: " + str(e)
            labela = Label(popup2, text=az_msg)
            labela.pack(side="top", fill="x", pady=10)
            labelb = Label(popup2, text=el_msg)
            labelb.pack(side="top", fill="x", pady=10)
            B = Button(popup2, text="Exit", command = popup2.destroy)
            B.pack()
        try:
            recv_msg = self.ard.readline()
            data = recv_msg.split(',')
            az = data[0]
            el = data[1]
            motorpopupmsg(az,el)
        except:
            motorpopupmsg('N/A','N/A')

if __name__ == '__main__':
    root = Tk()
    root.geometry("300x450")

    # find command to run GNU blocks for receiving
    host = '192.168.10.1' #check host number to match USRP's IP
    port = 12345
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind(('',port))
    BUFFER_SIZE = 1472

    try:
        ardunio = serial.Serial('COM3', 9600, timeout = 2)
    except:
        arduino = 0
    app = Window(s,arduino, root)
    root.mainloop()


import gym
import numpy as np
from gym.spaces import Box, Discrete, Tuple
from ray.rllib import MultiAgentEnv
from ray.rllib.env import EnvContext
from ray.rllib.utils import override
from ray.rllib.utils.typing import MultiAgentDict
from typing import Type
from prl.environment.Wrappers.utils import init_wrapped_env

def make_multi_agent_env(env_config) -> Type[MultiAgentEnv]:
    class _RLLibSteinbergerEnv(MultiAgentEnv):
        """Single Env holding multiple agents. Implements rllib.env.MultiAgentEnv to
        become vectorizable by rllib via `num_envs_per_worker` that can be increased via
        the algorithm config. """

        # see https://docs.ray.io/en/master/rllib/rllib-env.html
        def __init__(self, config: EnvContext):
            # config: EnvContext is not called here directly,
            # but it is called in ray.rllib.algorithm._get_env_id_and_creator
            # to set up ray remote workers
            self._n_players = env_config['n_players']
            self._starting_stack_size = env_config['starting_stack_size']
            self._agent_observation_mode = env_config['agent_observation_mode']
            self._blinds = env_config['blinds']
            self._env_cls = env_config['env_wrapper_cls']
            self.env_wrapped = self._single_env()
            self._num_agents = self._n_players  # keep naming consistency with rllib
            self.agents = env_config['agents']
            self.action_space = self.env_wrapped.action_space  # not batched, rllib wants that to be for single env
            # self.observation_space = self.agents[
            #     0].observation_space  # not batched, rllib wants that to be for single env
            # self.observation_space.dtype = np.float32
            # todo fix following monkeypatch:
            obs_space = Box(low=0.0, high=6.0, shape=(569,), dtype=np.float64)
            self.observation_space = obs_space
            if 'mask_legal_moves' in env_config:
                if env_config['mask_legal_moves']:
                    self.observation_space = gym.spaces.Dict({
                        'obs': obs_space,  # do not change key-name 'obs' it is internally used by rllib (!)
                        'legal_moves': Box(low=0, high=1, shape=(8,), dtype=int) # one-hot encoded [FOLD, CHECK_CALL, RAISE]
                    })

            self._agent_ids = list(self.agents.keys())

            MultiAgentEnv.__init__(self)

        def _single_env(self):
            return init_wrapped_env(self._env_cls,
                                    multiply_by=1,
                                    scale_rewards=False,
                                    stack_sizes=[self._starting_stack_size for _ in range(self._n_players)],
                                    blinds=self._blinds,
                                    agent_observation_mode=self._agent_observation_mode)

        @override(MultiAgentEnv)
        def reset(self):
            # return only obs, nothing else
            obs, _, _, info = self.env_wrapped.reset(config=None)
            next_to_act = self.env_wrapped.env.current_player.seat_id
            legal_moves = self.env_wrapped.get_legal_moves_extended()
            return {next_to_act: {'obs': obs, 'legal_moves': legal_moves}}

        @override(MultiAgentEnv)
        def step(self, action_dict):
            """When implementing your own MultiAgentEnv, note that you should only return those agent IDs in an
            observation dict, for which you expect to receive actions in the next call to step().

            From the rllib docs:
            # Env, in which two agents step in sequence (tuen-based game).
            # The env is in charge of the produced agent ID. Our env here produces
            # agent IDs: "player1" and "player2".
            env = TicTacToe()

            # Observations are a dict mapping agent names to their obs. Only those
            # agents' names that require actions in the next call to `step()` should
            # be present in the returned observation dict (here: one agent at a time).
            print(env.reset())
            # ... {
            # ...   "player1": [[...]],
            # ... }

            # In the following call to `step`, only those agents' actions should be
            # provided that were present in the returned obs dict:
            new_obs, rewards, dones, infos = env.step(actions={"player1": ...})

            # Similarly, new_obs, rewards, dones, etc. also become dicts.
            # Note that only in the `rewards` dict, any agent may be listed (even those that have
            # not(!) acted in the `step()` call). Rewards for individual agents will be added
            # up to the point where a new action for that agent is needed. This way, you may
            # implement a turn-based 2-player game, in which player-2's reward is published
            # in the `rewards` dict immediately after player-1 has acted.
            print(rewards)
            # ... {"player1": 0, "player2": 0}

            # Individual agents can early exit; The entire episode is done when
            # dones["__all__"] = True.
            print(dones)
            # ... {"player1": False, "__all__": False}

            # In the next step, it's player2's turn. Therefore, `new_obs` only container
            # this agent's ID:
            print(new_obs)
            # ... {
            # ...   "player2": [[...]]
            # ... }
            """

            observations, rewards, dones, infos = {}, {}, {}, {}

            # make sure only one player acted and step environment with its action
            have_played = list(action_dict.keys())
            has_played = have_played[0]
            assert len(have_played) == 1
            action = action_dict[has_played]
            obs, rews, done, info = self.env_wrapped.step(action)
            legal_moves = self.env_wrapped.get_legal_moves_extended()
            # assign returned observation to new player,
            # dones to old player and reward for all players
            next = self.env_wrapped.env.current_player.seat_id
            observations[next] = {'obs': obs, 'legal_moves': legal_moves}
            for i, v in enumerate(rews):
                rewards[i] = v / self._n_players  # normalize v because rllib stacks rewards per round
            dones[has_played] = done
            dones["__all__"] = done
            return observations, rewards, dones, infos

        @override(MultiAgentEnv)
        def render(self, mode='human'):
            pass

        @override(MultiAgentEnv)
        def observation_space_sample(self, agent_ids: list = None) -> MultiAgentDict:
            """The name 'agent_ids' is taken from rllib's core fn, although I think env_ids would be better"""
            if agent_ids is None:
                agent_ids = list(range(len(self.agents)))
            obs = {agent_id: self.observation_space.sample() for agent_id in agent_ids}

            return obs

        @override(MultiAgentEnv)
        def action_space_sample(self, agent_ids: list = None) -> MultiAgentDict:
            if agent_ids is None:
                agent_ids = list(range(len(self.agents)))
            actions = {agent_id: self.action_space.sample() for agent_id in agent_ids}
            return actions

        @override(MultiAgentEnv)
        def action_space_contains(self, x: MultiAgentDict) -> bool:
            if not isinstance(x, dict):
                return False
            return all(self.action_space.contains(val) for val in x.values())

        @override(MultiAgentEnv)
        def observation_space_contains(self, x: MultiAgentDict) -> bool:
            if not isinstance(x, dict):
                return False
            obs_dict = list(x.values())[0]
            if type(obs_dict) == dict:
                return self.observation_space['obs'].contains(obs_dict['obs'])
            else:
                return all(self.observation_space.contains(val) for val in x.values())

    return _RLLibSteinbergerEnv


import unittest
import unittest.mock as mock

from lib.cleanup_resource_handler import CleanupResourceHandler
from lib.route53 import Route53RecordSetLocator

class TestCleanupResourceHandler(unittest.TestCase):
    def test_handler_rejects_invalid_properties(self):
        handler = CleanupResourceHandler(route53_client=mock.Mock())
        with self.assertRaises(Exception):
            handler.handle_event({'RequestType': 'Delete', 'ResourceProperties': {'Invalid': 'Invalid'}}, {})

    def test_handling_delete(self):
        handler = CleanupResourceHandler(route53_client=mock.Mock(), monitor_interval=0)
        record_set_accessor = mock.Mock()
        record_set_accessor.delete = mock.Mock(return_value=True)  # True = Deleted

        exists_count = 0

        def exists_side_effect(*args):
            nonlocal exists_count
            exists_count += 1
            return True if exists_count < 3 else False

        record_set_accessor.exists = mock.Mock(side_effect=exists_side_effect)

        handler.record_set_accessor = record_set_accessor

        event = {
            'RequestType': 'Delete',
            'ResourceProperties': {
                'ServiceToken': 'Something',
                'HostedZoneId': 'ZONE',
                'RecordName': 'something.mydomain.com'
            }
        }

        # WHEN
        handler.handle_event(event, {})

        # THEN
        expected_locator = Route53RecordSetLocator(hosted_zone_id='ZONE', record_name='something.mydomain.com')
        record_set_accessor.delete.assert_called_with(locator=expected_locator)
        record_set_accessor.exists.assert_called()
        self.assertEqual(record_set_accessor.exists.call_count, 3)


import lcm

from gps import gps_t

def my_handler(channel, data):
    msg = gps_t.decode(data)
    print(" Received message on channel \"%s\"" % channel)
    print(" latitude = %s" % str(msg.lat))
    print(" latitude direction = %s" % msg.lat_dir)
    print(" longtitude = %s" % str(msg.lon))
    print(" longtitude direction = %s" % msg.lon_dir)
    print(" altitude = %s" % str(msg.alt))
    print(" utm_x = %s" % str(msg.utm_x))
    print(" utm_y = %s" % str(msg.utm_y))

lc = lcm.LCM()
subscription = lc.subscribe("GPS", my_handler)

try:
    while True:
        lc.handle()
except KeyboardInterrupt:
    pass

lc.unsubscribe(subscription)


import networkx as nx
import numpy as np
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
import time

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

import build_graph
from gnn import GNN
from util import normalize_adj

criterion = nn.CrossEntropyLoss()
d = torch.device('cpu')
frequency_as_feature = False
max_val_accuracy = 0
test_accuracy = 0
max_acc_epoch = 0
start_time = time.time()

h0, h1, h2, labels = None, None, None, None

class S2VGraph(object):
    def __init__(self, g, label, word_freq=None, node_features=None, positions=None):
        '''
            g: a networkx graph
            label: an integer graph label
            node_tags: a list of integer node tags
            node_features: a torch float tensor, one-hot representation of the tag that is used as input to neural nets
            edge_mat: a torch long tensor, contain edge list, will be used to create torch sparse tensor
            neighbors: list of neighbors (without self-loop)
        '''
        self.label = label
        self.g = g
        # self.word_freq = word_freq
        self.word_freq1 = word_freq[0]
        self.word_freq2 = word_freq[1]
        self.node_features = node_features
        self.positions = positions
        # self.node_features = row_norm(self.node_features)
        self.edge_mat = 0
        self.edges_weights = []

def print_distr(y, train_size):
    print('Class distributions')
    freq = [0 for i in range(len(set(y)))]
    for i in y[:train_size]:
        freq[i] += 1
    print('on train : ', freq)
    m = 1
    for i in freq:
        if i < m:
            m = i
    weights = [m / i if i != 0 else 1 for i in freq]
    global criterion
    # criterion = nn.CrossEntropyLoss(weight=torch.FloatTensor(weights).to(d))
    freq = [0 for i in range(len(set(y)))]
    for i in y[train_size:]:
        freq[i] += 1
    print('on test : ', freq)

def create_gaph(args):
    ls_adj, feature_list, word_freq_list, y, y_hot, train_size, word_vectors, positions_list = build_graph.build_graph(
        config=args.configfile)
    word_vectors = torch.from_numpy(word_vectors).float()

    print_distr(y, train_size)
    g_list = []
    max_words = 0
    for i, adj in enumerate(ls_adj):
        adj = normalize_adj(adj)
        g = nx.from_scipy_sparse_matrix(adj)
        lb = y[i]
        feat = feature_list[i]
        # if frequency_as_feature:
        #     feat = np.concatenate((feat, word_freq_list[i].toarray()), axis=1)
        #     # feat = feat * word_freq_list[i].toarray()
        if i == 10:
            print(word_freq_list[i])
        # s = sum(word_freq_list[i])
        # # s = 1
        # wf = [el / s for el in word_freq_list[i]]
        s = sum(word_freq_list[i][0])
        # s = 1
        wf1 = [el / s for el in word_freq_list[i][0]]
        s = sum(word_freq_list[i][1])
        # s = 1
        wf2 = [el / s for el in word_freq_list[i][1]]
        wf = (wf1, wf2)

        g_list.append(S2VGraph(g, lb, node_features=feat, word_freq=wf, positions=positions_list[i]))
        for ar in positions_list[i]:
            max_words = max(max_words, max(ar))

    max_words += 1

    zero_edges = 0
    for g in g_list:
        # edges = [list(pair) for pair in g.g.edges()]
        # edges_w = [w['weight'] for i, j, w in g.g.edges(data=True)]
        # edges.extend([[i, j] for j, i in edges])
        # edges_w.extend([w for w in edges_w])
        edges = []
        edges_w = []
        for i, j, wt in g.g.edges(data=True):
            w = wt['weight']
            edges.append([i, j])
            edges_w.append(w)
            if i != j:
                edges.append([j, i])
                edges_w.append(w)

        if len(edges) == 0:
            print('zero edge : ', len(g.g))
            zero_edges += 1
            edges = [[0, 0]]
            edges_w = [0]
        g.edge_mat = torch.tensor(edges).long().transpose(0, 1)
        g.edges_weights = torch.tensor(edges_w).float()
    print('total zero edge graphs : ', zero_edges)
    return g_list, len(set(y)), train_size, word_vectors, max_words

def train(args, model_e, device, graphs, optimizer, epoch, train_size, word_vectors):
    model_e.train()
    total_size = len(graphs)

    val_size = train_size // args.k_fold
    train_size = train_size - val_size
    test_size = total_size - train_size

    idx_train = np.random.permutation(train_size)
    loss_accum = 0
    for i in range(0, train_size, args.batch_size):
        selected_idx = idx_train[i:i + args.batch_size]
        batch_graph = [graphs[idx] for idx in selected_idx]
        # if len(selected_idx) == 0:
        #     continue
        output, pooled_h_ls = model_e(batch_graph, word_vectors)

        labels = torch.tensor([graph.label for graph in batch_graph]).long().to(device)

        # compute loss
        loss = criterion(output, labels)

        # backprop
        if optimizer is not None:
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        loss = loss.detach().cpu().numpy()
        loss_accum += loss

    print('Epoch : ', epoch, 'loss training: ', loss_accum, 'Time : ', int(time.time() - start_time))
    # print(model_e.pos)
    return loss_accum

# pass data to model with mini batch during testing to avoid memory overflow (does not perform backpropagation)
def pass_data_iteratively(args, model_e, graphs, minibatch_size, device, word_vectors):
    outputs = []
    hidden_rep_0 = []
    hidden_rep_1 = []
    full_idx = np.arange(len(graphs))
    for i in range(0, len(graphs), minibatch_size):
        sampled_idx = full_idx[i:i + minibatch_size]
        if len(sampled_idx) == 0:
            continue
        with torch.no_grad():
            output, pooled_h_ls = model_e([graphs[j] for j in sampled_idx], word_vectors)
        outputs.append(output)
        hidden_rep_0.append(pooled_h_ls[0].detach().cpu())
        hidden_rep_1.append(pooled_h_ls[1].detach().cpu())

    return torch.cat(outputs, 0), torch.cat(hidden_rep_0, 0).numpy(), torch.cat(hidden_rep_1, 0).numpy()

def test(args, model_e, device, graphs, train_size, epoch, word_vectors):
    model_e.eval()

    val_size = train_size // args.k_fold
    train_size = train_size - val_size

    output, hidden_rep_0, hidden_rep_1 = pass_data_iteratively(args, model_e, graphs, 100, device, word_vectors)

    output_train, train_graphs = output[:train_size], graphs[:train_size]
    output_val, val_graph = output[train_size:train_size + val_size], graphs[train_size:train_size + val_size]
    output_test, test_graphs = output[train_size + val_size:], graphs[train_size + val_size:]

    pred_train = output_train.max(1, keepdim=True)[1]
    labels_train = torch.tensor([graph.label for graph in train_graphs]).long().to(device)
    correct = pred_train.eq(labels_train.view_as(pred_train)).sum().cpu().item()
    acc_train = correct / float(len(train_graphs))

    pred_val = output_val.max(1, keepdim=True)[1]
    labels_val = torch.tensor([graph.label for graph in val_graph]).long().to(device)
    correct = pred_val.eq(labels_val.view_as(pred_val)).sum().cpu().item()
    acc_val = correct / float(len(val_graph))

    pred_test = output_test.max(1, keepdim=True)[1]
    labels_test = torch.tensor([graph.label for graph in test_graphs]).long().to(device)
    correct = pred_test.eq(labels_test.view_as(pred_test)).sum().cpu().item()
    acc_test = correct / float(len(test_graphs))

    if args.debug:
        print(time.time() - start_time, 's epoch : ', epoch)
    print("accuracy train: %f val: %f test: %f" % (acc_train, acc_val, acc_test), flush=True)
    global max_acc_epoch, max_val_accuracy, test_accuracy
    if acc_val > max_val_accuracy:
        max_val_accuracy = acc_val
        max_acc_epoch = epoch
        test_accuracy = acc_test
        if args.tsne:
            global h0, h1, h2, labels
            h0, h1 = hidden_rep_0[train_size + val_size:], hidden_rep_1[train_size + val_size:]
            h2 = output_test.detach().cpu().numpy()
            labels = labels_test.detach().cpu().numpy()

    print('max validation accuracy : ', max_val_accuracy, 'max acc epoch : ', max_acc_epoch, flush=True)
    if args.debug:
        print('epsilon : ', model_e.eps.detach().cpu().numpy(), 'w1 : ', model_e.w1.detach().cpu().numpy(), flush=True)

        # if epoch == 800:
        #     for i in range(len(test_graphs)):
        #         print('label : ', labels_test[i].cpu().item(), ' pred : ', pred_test[i].cpu().item())

    return acc_train, acc_test

def main():
    print('date and time : ', time.ctime())
    global max_acc_epoch, max_val_accuracy, test_accuracy
    acc_test = 0
    parser = argparse.ArgumentParser(
        description='PyTorch graph convolutional neural net for whole-graph classification')
    parser.add_argument('--device', type=int, default=0,
                        help='which gpu to use if any (default: 0)')
    parser.add_argument('--batch_size', type=int, default=16,
                        help='input batch size for training (default: 64)')
    parser.add_argument('--epochs', type=int, default=400,
                        help='number of epochs to train (default: 350)')
    parser.add_argument('--lr', type=float, default=0.0001,
                        help='learning rate (default: 0.01)')
    parser.add_argument('--seed', type=int, default=0,
                        help='random seed for splitting the dataset into 10 (default: 0)')
    parser.add_argument('--num_layers', type=int, default=3,
                        help='number of layers INCLUDING the input one (default: 5)')
    parser.add_argument('--num_mlp_layers', type=int, default=2,
                        help='number of layers for MLP EXCLUDING the input one (default: 2). 1 means linear model.')
    parser.add_argument('--hidden_dim', type=int, default=200,
                        help='number of hidden units (default: 64)')
    parser.add_argument('--final_dropout', type=float, default=0.5,
                        help='final layer dropout (default: 0.5)')
    parser.add_argument('--graph_pooling_type', type=str, default="average", choices=["sum", "average"],
                        help='Pooling for over nodes in a graph: sum or average')
    parser.add_argument('--neighbor_pooling_type', type=str, default="sum", choices=["sum", "average", "max"],
                        help='Pooling for over neighboring nodes: sum, average or max')
    parser.add_argument('--learn_eps', action="store_true",
                        help='Whether to learn the epsilon weighting for the center nodes. Does not affect training '
                             'accuracy though.')
    parser.add_argument('--configfile', type=str, default="param",
                        help='configuration file')
    parser.add_argument('--filename', type=str, default="",
                        help='output file')
    parser.add_argument('--k_fold', type=int, default=0,
                        help='k_fold')
    parser.add_argument('--early_stop', type=int, default=30,
                        help='early_stop')
    parser.add_argument('--debug', action="store_true",
                        help='run in debug mode')
    parser.add_argument('--tsne', action="store_true",
                        help='tsne')
    parser.add_argument('--num_heads', type=int, default=2,
                        help='number of hidden units (default: 64)')
    parser.add_argument('--weight_decay', type=float, default=0, help='weight decay (default: 0.3)')

    args = parser.parse_args()

    print(args)

    # set up seeds and gpu device
    torch.manual_seed(0)
    np.random.seed(0)
    device = torch.device("cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu")
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(0)
    print('device : ', device, flush=True)

    global d
    d = device
    all_graphs, num_classes, train_size, word_vectors, max_words = create_gaph(args)

    acc_detais = []
    k_start = 0
    if args.k_fold == 0:
        k_start = 4
        args.k_fold = 5
    val_size = train_size // args.k_fold
    for k in range(k_start, args.k_fold):
        start = k * val_size
        end = start + val_size
        graphs = all_graphs[:start] + all_graphs[end: train_size] + all_graphs[start:end] + all_graphs[train_size:]

        # ge = [None for i in range(args.num_layers)]

        # model_c = ClusterNN(num_classes, graphs[0].node_features.shape[1], args.hidden_dim, args.num_layers,
        #                     args.num_mlp_layers_c).to(device)
        model_e = GNN(args.num_layers, args.num_mlp_layers, word_vectors.shape[1], args.hidden_dim,
                      num_classes,
                      args.final_dropout, args.learn_eps, args.graph_pooling_type, args.neighbor_pooling_type,
                      device, max_words, args.num_heads, word_vectors).to(device)

        optimizer = optim.Adam(model_e.parameters(), lr=args.lr, weight_decay=args.weight_decay)
        # optimizer_c = optim.Adam(model_c.parameters(), lr=args.lr_c)
        # optimizer = optim.SGD(model_e.parameters(), lr=args.lr, momentum=0.9)
        # optimizer_c = optim.SGD(model_c.parameters(), lr=args.lr_c, momentum=0.9)
        # scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.5)
        # cl = None

        print(time.time() - start_time, 's Training starts', flush=True)
        for epoch in range(1, args.epochs + 1):
            avg_loss = train(args, model_e, device, graphs, optimizer, epoch, train_size, word_vectors)
            acc_train, acc_test = test(args, model_e, device, graphs, train_size, epoch, word_vectors)
            print("")

            if epoch > max_acc_epoch + args.early_stop \
                    and epoch > args.early_stop:
                break

        if h0 is not None and args.tsne:
            tsne = TSNE()
            h0_e = tsne.fit_transform(h0)
            h1_e = tsne.fit_transform(h1)
            h2_e = tsne.fit_transform(h2)

            plt.figure()
            plt.scatter(h0_e[:, 0], h0_e[:, 1], c=labels)
            plt.savefig(args.configfile + str(test_accuracy) + 'h0.png')

            plt.figure()
            plt.scatter(h1_e[:, 0], h1_e[:, 1], c=labels)
            plt.savefig(args.configfile + str(test_accuracy) + 'h1.png')

            plt.figure()
            plt.scatter(h2_e[:, 0], h2_e[:, 1], c=labels)
            plt.savefig(args.configfile + str(test_accuracy) + 'h2.png')

        print('=' * 200)
        print('K : ', k, 'Time : ', abs(time.time() - start_time))
        print('max validation accuracy : ', max_val_accuracy * 100)
        print('max acc epoch : ', max_acc_epoch)
        print('test accuracy : ', test_accuracy * 100)
        print('latest_test_accuracy : ', acc_test * 100)
        print('=' * 200 + '\n')
        acc_detais.append((max_val_accuracy, test_accuracy, max_acc_epoch, acc_test))
        max_val_accuracy = 0
        test_accuracy = 0
        max_acc_epoch = 0

    print(args)
    print('total size : ', len(all_graphs), '\n')
    print('=' * 71 + 'Summary' + '=' * 71)
    avg = [0] * 3
    for k in range(len(acc_detais)):
        avg[0] += acc_detais[k][0]
        avg[1] += acc_detais[k][1]
        avg[2] += acc_detais[k][3]
        print('k : ', k,
              '\tval_accuracy : ', acc_detais[k][0] * 100,
              '\ttest_accuracy : ', acc_detais[k][1] * 100,
              '\tmax_acc epoch : ', acc_detais[k][2],
              '\tlatest_test_accuracy : ', acc_detais[k][3] * 100)

    # print('\navg : ',
    #       '\tval_accuracy : ', avg[0] / len(acc_detais) * 100,
    #       '\ttest_accuracy : ', avg[1] / len(acc_detais) * 100,
    #       '\tlatest_test_accuracy : ', avg[2] / len(acc_detais) * 100)

if __name__ == '__main__':
    main()