import numpy as np
import pygame
import os
import threading

from webapp import app

from event_server_comm import move_grid

BLACK = np.array([0, 0, 0], dtype=np.uint8)
WHITE = np.array([255, 255, 255], dtype=np.uint8)
GRAY = np.array([200, 200, 200], dtype=np.uint8)
RED = np.array([255, 0, 0], dtype=np.uint8)
BLUE = np.array([0, 0, 255], dtype=np.uint8)
YELLOW = np.array([255, 255, 0], dtype=np.uint8)
GREEN = np.array([0, 255, 0], dtype=np.uint8)

pygame.init()

pygame.font.init()  # you have to call this at the start,
# if you want to use this module.
myfont = pygame.font.SysFont('Comic Sans MS', 55)
myfont_small = pygame.font.SysFont('Comic Sans MS', 45)

P0_text = myfont.render('P0', False, tuple(BLACK))

tiledt = np.dtype([('color', np.uint8, 3), ('star', np.bool)])



class Board:
    valid_colors = [GRAY, RED, BLUE]

    def __init__(self, dim_x, dim_y, n_coins=2, file=None):
        if file is None:
            self.tiles = np.zeros((dim_y, dim_x), dtype=tiledt)
            for x in range(dim_x):
                for y in range(dim_y):
                    self.tiles[y, x]['color'] = Board.valid_colors[np.random.randint(len(Board.valid_colors))]

            coins_distributed = n_coins == 0
            while not coins_distributed:
                coinx = np.random.randint(0, dim_x)
                coiny = np.random.randint(0, dim_y)
                self.tiles[coiny,coinx]['star'] = True
                coins_distributed = sum([t['star'] for t in self.tiles.flatten()]) == n_coins
        else:
            self.tiles = np.load(file)

    def render(self, scale_fac):
        dimy, dimx = self.tiles.shape
        board_surf = pygame.Surface((dimx * scale_fac, dimy * scale_fac))

        star_surf = pygame.Surface((scale_fac, scale_fac), pygame.SRCALPHA)
        pygame.draw.circle(star_surf, YELLOW, (int(0.5 * scale_fac), int(0.5 * scale_fac)), int(0.25 * scale_fac))

        for y in range(self.tiles.shape[0]):
            for x in range(self.tiles.shape[1]):
                pygame.draw.rect(board_surf, tuple(self.tiles[y, x]['color']),
                                 (x * scale_fac, y * scale_fac, scale_fac, scale_fac), 0)
                pygame.draw.rect(board_surf, (0, 0, 0),
                                 (x * scale_fac, y * scale_fac, scale_fac, scale_fac), 1)
                if self.tiles[y, x]['star']:
                    board_surf.blit(star_surf, (x * scale_fac, y * scale_fac, scale_fac, scale_fac))

        return board_surf

    def get_xdims(self):
        return self.tiles.shape[1]

    def get_ydims(self):
        return self.tiles.shape[0]

    # def __repr__(self):
    #     s = ''
    #     for y in range(self.tiles.shape[0]):
    #         for x in range(self.tiles.shape[1]):
    #             if (x,y) == (self.robot.x, self.robot.y):
    #                 s += self.robot.orientation
    #             else:
    #                 s += '.'
    #         s += '\n'
    #     return s


class Robot:
    orientations = ['^', 'left', 'down', 'right']
    resulting_orientation = {
        '^': {'left': '<', 'right': '>'},
        '>': {'left': '^', 'right': 'v'},
        'v': {'left': '>', 'right': '<'},
        '<': {'left': 'v', 'right': '^'},
    }


    def __init__(self, x, y, orientation, use_real_robot=False):
        self.x = x
        self.y = y
        self.orientation = orientation
        self.position_changed = False
        self.use_real_robot = use_real_robot

    def get_forward_coordinates(self):
        # get the coordinates of the neighboring tile in the given direction
        if self.orientation == '^':
            return self.y - 1, self.x
        elif self.orientation == '>':
            return self.y, self.x + 1
        elif self.orientation == 'v':
            return self.y + 1, self.x
        elif self.orientation == '<':
            return self.y, self.x - 1
        else:
            raise Exception("error: undefined direction")

    def get_angle(self):
        angle = {'>': 0, '^': np.pi/2, '<': np.pi, 'v': 3*np.pi/2}[self.orientation]
        return np.rad2deg(angle)

    def render(self, scale_fac):
        robot_surf = pygame.Surface((scale_fac, scale_fac), pygame.SRCALPHA)
        pygame.draw.lines(robot_surf, (0, 0, 0), True, [(0.75 * scale_fac, 0.5 * scale_fac),
                                                        (0.25 * scale_fac, 0.25 * scale_fac),
                                                        (0.25 * scale_fac, 0.75 * scale_fac)], 3)
        robot_surf = pygame.transform.rotate(robot_surf, self.get_angle())
        return robot_surf

    def update_pos(self, dimx, dimy):
        if self.use_real_robot:
            move_grid(self.x, self.y, self.orientation, dimx, dimy)
        self.position_changed = False

    def __repr__(self):
        return f"({self.y}, {self.x}) - {self.orientation}"


class Command:
    valid_actions = {'forward', 'left', 'right', 'P0', '-'}

    def __init__(self, action=None, color=GRAY):
        if not (action in Command.valid_actions and any([np.all(color == c) for c in Board.valid_colors])):
            raise ValueError("invalid values for command")
        self.action = action
        self.color = color

    def __repr__(self):
        return f"{self.action}: {self.color}"

    def render(self, scale_fac):
        cmd_surf = pygame.Surface((scale_fac, scale_fac))
        cmd_surf.fill(tuple(self.color))

        arrow_surf = pygame.Surface((300, 300), pygame.SRCALPHA)

        pygame.draw.polygon(arrow_surf, (0, 0, 0),
                            ((0, 100), (0, 200), (200, 200), (200, 300), (300, 150), (200, 0), (200, 100)))
        arrow_surf = pygame.transform.scale(arrow_surf, (int(0.9*scale_fac), int(0.9*scale_fac)))

        if self.action == 'forward':
            arrow_surf = pygame.transform.rotate(arrow_surf, 90)
        elif self.action == 'left':
            arrow_surf = pygame.transform.rotate(arrow_surf, 180)

        if self.action in {'left', 'forward', 'right'}:
            cmd_surf.blit(arrow_surf, (0.05*scale_fac,0.05*scale_fac,0.95*scale_fac,0.95*scale_fac))
        elif self.action == 'P0':
            cmd_surf.blit(P0_text, (0.05*scale_fac,0.05*scale_fac,0.95*scale_fac,0.95*scale_fac))

        return cmd_surf

class Programmer:
    def __init__(self, prg):
        self.prg = prg
        self.available_inputs = [Command('forward'), Command('left'), Command('right'), Command('P0'),
                                 Command('-', color=RED), Command('-', color=BLUE), Command('-', color=GRAY)]
        self.command_to_edit = 0
        self.screen_rect = None

    def render(self, scale_fac):
        """Render surface with possible inputs for the robot.

        :return: surface of the input commands
        """
        inp_surf = pygame.Surface((len(self.available_inputs) * scale_fac, 1 * scale_fac))
        for i, inp in enumerate(self.available_inputs):
            cmd_surf = inp.render(scale_fac)
            inp_surf.blit(cmd_surf, (i * scale_fac, 0, scale_fac, scale_fac))
        return inp_surf

    def update_selected_command(self, pos):
        print(f"clicked at pos = {pos}")
        xoffset = pos[0] - self.screen_rect.x
        selected_input_index = xoffset * len(self.available_inputs) // self.screen_rect.width
        selected_input = self.available_inputs[selected_input_index]

        edited_command = self.prg.cmds[self.command_to_edit]

        print("command before edit = ", edited_command)
        if selected_input.action != '-':
            edited_command.action = selected_input.action
        else:
            edited_command.color = selected_input.color
        print("command after edit = ", edited_command)

class Program:
    def __init__(self, robot, board, cmds):
        self.cmds = cmds
        self.robot = robot
        self.board = board
        self.prg_counter = 0
        self.screen_rect = None

    def step(self, state='running', check_victory=True):
        if self.prg_counter >= len(self.cmds):
            return 'game_over'
        cmd = self.cmds[self.prg_counter]
        self.prg_counter += 1

        # current position
        x = self.robot.x
        y = self.robot.y

        # current tile the robot is on
        tile = self.board.tiles[y, x]

        # apply next instruction of the program
        if np.all(cmd.color == GRAY) or np.all(cmd.color == tile['color']):
            # matching color -> execute command
            if cmd.action == 'forward':
                ynew, xnew = self.robot.get_forward_coordinates()
                self.robot.x = xnew
                self.robot.y = ynew
                self.robot.position_changed = True
                #self.robot.update_pos(self.board.get_xdims(), self.board.get_ydims())
            elif cmd.action in {'left', 'right'}:
                self.robot.orientation = Robot.resulting_orientation[self.robot.orientation][cmd.action]
                #self.robot.update_pos(self.board.get_xdims(), self.board.get_ydims())
                self.robot.position_changed = True
            elif cmd.action == 'P0':
                self.prg_counter = 0
        else:
            #print("color not matching -> skipping command")
            pass

        # update state for new robot position
        if (not (0 <= self.robot.x < self.board.tiles.shape[1])) or not (0 <= self.robot.y < self.board.tiles.shape[0]):
            # robot leaves the board -> GAME OVER
            print("GAME OVER")
            self.robot.update_pos(self.board.get_xdims(), self.board.get_ydims())
            return 'game_over'

        # robot collects star on new tile
        tile = self.board.tiles[self.robot.y, self.robot.x]
        if tile['star']:
            tile['star'] = False

        if check_victory and all([not t['star'] for t in self.board.tiles.flatten()]):
            self.robot.update_pos(self.board.get_xdims(), self.board.get_ydims())
            print("YOU WON")
            return 'won'

        # by default we continue in the running state
        return state

    def render(self, scale_fac, prg_counter_override=None):
        """Render the current program. This will render all commands and highlight the next command to execute
        (determined by self.prg_counter).

        The prg_counter_override can be used to highlight a different command instead. This is used during input mode
        to highlight the command the user will edit.

        :param scale_fac:
        :param prg_counter_override:
        :return:
        """
        prg_counter = self.prg_counter
        if prg_counter_override is not None:
            prg_counter = prg_counter_override
        prg_surf = pygame.Surface((5 * scale_fac, 1 * scale_fac))
        for i in range(5):
            if i < len(self.cmds):
                cmd = self.cmds[i]
                cmd_surf = cmd.render(scale_fac)
            else:
                cmd_surf = pygame.Surface((scale_fac,scale_fac))
                cmd_surf.fill(GRAY)
            if prg_counter is not None and i == prg_counter:
                pygame.draw.rect(cmd_surf, tuple(GREEN), (0, 0, scale_fac, scale_fac), 5)
            prg_surf.blit(cmd_surf, (i * scale_fac, 0, scale_fac, scale_fac))

        return prg_surf

    def get_clicked_command(self, pos):
        print(f"clicked at pos = {pos}")
        xoffset = pos[0] - self.screen_rect.x
        clicked_cmd = xoffset * len(self.cmds) // self.screen_rect.width
        print("clicked command = ", clicked_cmd)
        return clicked_cmd


class Game:
    def __init__(self, dimx, dimy, robotx, roboty, orientation, use_real_robot=False):
        self.robot = Robot(x=robotx, y=roboty, orientation=orientation, use_real_robot=use_real_robot)
        #self.board = Board(dimx, dimy)
        self.board = Board(dimx, dimy, file='levels/56.npy')

        # TODO fix number of commands at 5
        self.cmds = [Command('forward'), Command('left', color=RED), Command('left', color=BLUE), Command('P0'), Command('-')]
        self.state = 'reset'

        self.prg = Program(self.robot, self.board, self.cmds)

        self.programmer = Programmer(self.prg)

        #self.scale_fac = 180
        self.scale_fac = 125
        self.beamer_mode = False
        self.screen = pygame.display.set_mode((int(self.board.tiles.shape[1] * self.scale_fac * 1.1),
                                               int((self.board.tiles.shape[0] + 2) * self.scale_fac * 1.2)))

        self.game_over_text = myfont.render('GAME OVER', False, WHITE)
        self.won_text = myfont.render('YOU WON', False, BLACK)
        self.run_text = myfont.render('RUN', False, tuple(BLACK))
        self.stop_text = myfont_small.render('STOP', False, tuple(BLACK))
        self.step_text = myfont_small.render('STEP', False, tuple(BLACK))
        self.prg_text = myfont_small.render('CURRENT PROGRAM', False, tuple(GREEN))


        # save initial state
        self.initial_pos = (self.robot.x, self.robot.y, self.robot.orientation)
        self.initial_board_tiles = self.board.tiles.copy()

    def render(self):
        """Render the game screen.

        This will render the board and the robot.
        Depending on the display mode it will also render the current program and the input commands for the robot.

        :param prg_counter:
        :return:
        """
        if self.beamer_mode:
            dx = self.xoffset
            dy = 0
        else:
            dx = int(0.05 * self.screen.get_width())
            dy = int(0.05 * self.screen.get_height())
        self.screen.fill(tuple(BLACK))

        # render the board
        board_surf = self.board.render(self.scale_fac)

        # render robot onto the board surface
        robot_surf = self.robot.render(self.scale_fac)
        board_surf.blit(robot_surf, (self.robot.x * self.scale_fac, self.robot.y * self.scale_fac, self.scale_fac, self.scale_fac))
        self.screen.blit(board_surf, (dx, dy, dx + self.board.tiles.shape[1] * self.scale_fac, dy + self.board.tiles.shape[0] * self.scale_fac))

        # render program
        if self.state == 'input':
            # in input mode we highlight the command which is selected for edit
            prg_surf = self.prg.render(self.scale_fac, prg_counter_override=self.programmer.command_to_edit)
        else:
            # in other modes we render the current program counter
            prg_surf = self.prg.render(self.scale_fac)
        prg_surf = pygame.transform.scale(prg_surf, (self.screen.get_width() * 2 // 3, self.scale_fac * 2 // 3))
        self.prg.screen_rect = pygame.Rect(
            (dx, board_surf.get_height() + 2 * dy, prg_surf.get_width(), prg_surf.get_height()))
        # render input fields and buttons
        inp_surf = self.programmer.render(self.scale_fac)
        inp_surf = pygame.transform.scale(inp_surf, (self.screen.get_width() * 2 // 3, self.scale_fac * 2 // 3))
        self.programmer.screen_rect = pygame.Rect(
            (dx, board_surf.get_height() + prg_surf.get_height() + 3 * dy, inp_surf.get_width(), inp_surf.get_height()))
        btn_surf = pygame.Surface((3 * self.scale_fac // 2, self.scale_fac))
        self.btn_rect = pygame.Rect((2 * dx + prg_surf.get_width(), board_surf.get_height() + 2 * dy,
                                     btn_surf.get_height(), btn_surf.get_width()))
        if self.state == 'input':
            btn_surf.fill(tuple(GREEN))
            btn_surf.blit(self.run_text, (0, 10))
        elif self.state == 'running':
            btn_surf.fill(tuple(RED))
            btn_surf.blit(self.stop_text, (0, 10))
        elif self.state == 'stepping':
            btn_surf.fill(tuple(YELLOW))
            btn_surf.blit(self.step_text, (0, 10))

        if not self.beamer_mode:
            # if we are not in beamer mode we render program and inputs below the board
            self.screen.blit(prg_surf, self.prg.screen_rect)

            self.screen.blit(inp_surf, self.programmer.screen_rect)

            self.screen.blit(btn_surf, self.btn_rect)
        else:
            prg_surf = pygame.transform.scale(prg_surf, (dx, dx//5))
            # in beamer mode we render the program to the left of the board to appear on the laptop
            self.screen.blit(prg_surf, (0,100))

            prg_descr_surb = pygame.Surface((500, 100))
            #prg_descr_surb.blit(self.prg_text)
            self.screen.blit(self.prg_text, (50, 50))
            mode_text = myfont_small.render(f'STATE: {self.state}', False, tuple(GREEN))
            self.screen.blit(mode_text, (50, dx//5 + 350))

        # render messages
        if self.state == 'game_over':
            game_over_surf = pygame.Surface(((self.screen.get_width() - dx) // 2, self.screen.get_height() // 2))
            game_over_surf.fill(tuple(BLACK))
            game_over_surf.blit(self.game_over_text, ((game_over_surf.get_width() - self.game_over_text.get_width()) // 2,
                                                      (game_over_surf.get_height() - self.game_over_text.get_height()) // 2))
            self.screen.blit(game_over_surf, (dx + (self.screen.get_width() - dx) // 4, self.screen.get_height() // 4))
            pygame.display.update()
            pygame.time.wait(1500)
            self.state = 'reset'
        elif self.state == 'won':
            won_surf = pygame.Surface(((self.screen.get_width() - dx) // 2, self.screen.get_height() // 2))
            won_surf.fill(tuple(WHITE))
            won_surf.blit(self.won_text,
                                ((won_surf.get_width() - self.won_text.get_width()) // 2,
                                 (won_surf.get_height() - self.won_text.get_height()) // 2))
            self.screen.blit(won_surf, (dx + (self.screen.get_width() - dx) // 4, self.screen.get_height() // 4))
            pygame.display.update()
            pygame.time.wait(1500)
            self.state = 'reset'

        pygame.display.flip()

    def process_inputs(self):
        # proceed events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.state = 'quit'
            # handle MOUSEBUTTONUP
            elif event.type == pygame.MOUSEBUTTONUP:
                pos = pygame.mouse.get_pos()

                # select command to edit by the programmer
                if self.prg.screen_rect is not None:
                    if self.prg.screen_rect.collidepoint(pos):
                        print(f"clicked at pos = {pos}")
                        self.programmer.command_to_edit = self.prg.get_clicked_command(pos)

                # set the selected command to a new value
                if self.programmer.screen_rect is not None:
                    if self.programmer.screen_rect.collidepoint(pos):
                        self.programmer.update_selected_command(pos)

                # clicked RUN/STOP button
                if self.btn_rect is not None and self.btn_rect.collidepoint(*pos):
                    print(f"clicked at pos = {pos}")
                    if self.state == 'running':
                        self.state = 'input'
                    elif self.state == 'input':
                        self.state = 'running'
                    elif self.state == 'stepping':
                        self.state = self.prg.step(self.state)
            elif event.type == pygame.KEYUP:
                if event.key == pygame.K_x:
                    if not self.beamer_mode:
                        # switch to beamer mode
                        self.xoffset = 1000
                        os.environ['SDL_VIDEO_WINDOW_POS'] = f'{1920-self.xoffset}, 280'
                        self.scale_fac = 180
                        self.screen = pygame.display.set_mode((self.xoffset + self.board.tiles.shape[1] * self.scale_fac,
                                                          self.board.tiles.shape[0] * self.scale_fac),
                                                         pygame.NOFRAME)
                        self.beamer_mode = True
                    else:
                        # switch to normal mode
                        os.environ['SDL_VIDEO_WINDOW_POS'] = '0, 0'
                        self.scale_fac = 125
                        self.screen = pygame.display.set_mode((int(self.board.tiles.shape[1] * self.scale_fac * 1.1),
                                               int((self.board.tiles.shape[0] + 2) * self.scale_fac * 1.2)))
                        self.beamer_mode = False
                elif event.key == pygame.K_s:
                    # run program
                    self.state = 'running'
                elif event.key == pygame.K_SPACE:
                    if self.state != 'stepping':
                        self.state = 'stepping'
                    else:
                        self.state = self.prg.step(self.state)
                elif event.key == pygame.K_r:
                    self.state = 'reset'
                elif event.key == pygame.K_n:
                    self.initial_board_tiles = Board(self.board.get_xdims(), self.board.get_ydims()).tiles.copy()
                    self.state = 'reset'
            elif event.type == pygame.USEREVENT:
                for i, cmd in enumerate(event.cmds):
                    self.cmds[i].action = cmd.action
                    self.cmds[i].color = np.array(cmd.color, dtype=np.uint8)
                self.reset()
                self.state = 'running'
        return self.state

    def reset(self):
        self.prg.prg_counter = 0
        self.robot.x = self.initial_pos[0]
        self.robot.y = self.initial_pos[1]
        self.robot.orientation = self.initial_pos[2]
        self.robot.update_pos(self.board.get_xdims(), self.board.get_ydims())
        self.board.tiles = self.initial_board_tiles.copy()
        return 'input'

    def run(self):
        running = True
        while running:
            self.state = self.process_inputs()
            if self.state == 'input':
                pass
            elif self.state == 'running':
                self.state = self.prg.step()
            elif self.state == 'reset':
                self.state = self.reset()
            elif self.state == 'quit':
                running = False
            elif self.state == 'stepping':
                pass
            elif self.state == 'game_over' or self.state == 'won':
                pass
            else:
                print("unknown state")
                return

            self.render()
            if self.robot.position_changed:
                self.robot.update_pos(self.board.get_xdims(), self.board.get_ydims())

            pygame.time.wait(100)

if __name__ == "__main__":
    # launch webapp in thread
    webserver_thread = threading.Thread(target=app.run, kwargs={'host': '0.0.0.0', 'port': 5000})
    webserver_thread.start()

    seed = 4
    np.random.seed(seed)
    game = Game(dimx=7, dimy=4, robotx=5, roboty=1, orientation='>', use_real_robot=False)
    game.run()

# TODOs
# - in stepping mode (s) it is possible to drive outside of the map
# - when no P0 command is present the program counter will overflow