import numpy as np
import random
import pygame

import os

BLACK = np.array([0, 0, 0], dtype=np.uint8)
WHITE = np.array([255, 255, 255], 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([('x', np.uint8), ('y', np.uint8), ('color', np.uint8, 3), ('star', np.bool)])


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

    def __init__(self, dim_x, dim_y):
        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]['x'] = x
                self.tiles[y, x]['y'] = y
                self.tiles[y, x]['color'] = random.choice(Board.valid_colors)

    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 __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):
        self.x = x
        self.y = y
        self.orientation = orientation

    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 __repr__(self):
        return f"({self.y}, {self.x}) - {self.orientation}"


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

    def __init__(self, action=None, color=WHITE):
        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=WHITE)]
        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):
        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 == WHITE) 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
            elif cmd.action in {'left', 'right'}:
                self.robot.orientation = Robot.resulting_orientation[self.robot.orientation][cmd.action]
            elif cmd.action == 'P0':
                self.prg_counter = 0
        else:
            print("color not matching -> skipping command")

            # get all events
            ev = pygame.event.get()
            # proceed events
            for event in ev:
                # handle MOUSEBUTTONUP
                if event.type == pygame.MOUSEBUTTONUP:
                    pos = pygame.mouse.get_pos()
                    if pos[0] >= 325 and pos[0] <= 400 and pos[1] >= 600 and pos[1] <= 650:
                        print(f"clicked at pos = {pos}")
                        self.state = 'input'

        # 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")
            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 all([not t['star'] for t in self.board.tiles.flatten()]):
            print("YOU WON")
            return 'won'

        # by default we continue in the running state
        return 'running'

    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(WHITE)
            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):
        self.robot = Robot(x=robotx, y=roboty, orientation='v')
        self.board = Board(dimx, dimy)
        self.board.tiles[3,3]['star'] = True
        self.board.tiles[3,2]['star'] = True

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

        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, RED)
        self.won_text = myfont.render('YOU WON', False, GREEN)
        self.run_text = myfont.render('RUN', False, tuple(BLACK))
        self.stop_text = myfont_small.render('STOP', False, tuple(BLACK))

        # save initial state
        self.initial_pos = (self.robot.x, self.robot.y, self.robot.orientation)
        self.inital_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 = 0
            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))

        if not self.beamer_mode:
            # if we are not in beamer mode we also show the current program / inputs

            # 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()))
            self.screen.blit(prg_surf, self.prg.screen_rect)

            # 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()))
            self.screen.blit(inp_surf, self.programmer.screen_rect)

            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))
            self.screen.blit(btn_surf, self.btn_rect)

        # render messages
        if self.state == 'game_over':
            self.screen.blit(self.game_over_text, (50, 00))
            pygame.display.update()
            pygame.time.wait(1500)
            self.state = 'reset'
        elif self.state == 'won':
            self.screen.blit(self.won_text, (50, 00))
            pygame.display.update()
            pygame.time.wait(1500)
            self.state = 'reset'

        pygame.display.update()

    def process_inputs(self):
        # proceed events
        for event in pygame.event.get():
            # handle MOUSEBUTTONUP
            if 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 event.type == pygame.KEYDOWN:
                if event.key == pygame.K_x:
                    if not self.beamer_mode:
                        # switch to beamer mode
                        os.environ['SDL_VIDEO_WINDOW_POS'] = '1920, 280'
                        self.scale_fac = 180
                        self.screen = pygame.display.set_mode((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((self.board.tiles.shape[1] * self.scale_fac,
                                                          self.board.tiles.shape[0] * self.scale_fac + 5 * self.scale_fac))
                        self.beamer_mode = False
                elif event.key == pygame.K_r:
                    # run program
                    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.board.tiles = self.inital_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()
            else:
                print("unknown state")
                return

            self.render()
            pygame.time.wait(100)

if __name__ == "__main__":
    random.seed(42)
    game = Game(dimx=7, dimy=4, robotx=3, roboty=1)
    game.run()