import numpy as np
import random
import pygame
import time
import copy

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)
P0_text = myfont.render('P0', False, (0, 0, 0))
game_over_text = myfont.render('GAME OVER', False, RED)
run_text = myfont.render('RUN', False, tuple(BLACK))

random.seed(0)

dimx = 5
dimy = 10

scale_fac = 50
screen = pygame.display.set_mode((dimx*scale_fac+200,dimy*scale_fac+300))

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, robot):
        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)
        self.robot = robot

    def render(self):
        board_surf = pygame.Surface((dimx * scale_fac, dimy * scale_fac))
        robot_surf = pygame.Surface((scale_fac, scale_fac), pygame.SRCALPHA)
        # Use the local coordinate system of the surface to draw the lines.
        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)
        # I rotate it so that the arrow is pointing to the right (0° is right).
        robot_surf = pygame.transform.rotate(robot_surf, self.robot.get_angle())

        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))
                if (x, y) == (self.robot.x, self.robot.y):
                    board_surf.blit(robot_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 __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):
        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 Program:
    def __init__(self, robot, board, cmds):
        self.cmds = cmds
        self.robot = robot
        self.board = board

        self.available_inputs = [Command('forward'), Command('left'), Command('right'), Command('P0'),
                            Command('-', color=RED), Command('-', color=BLUE), Command('-', color=WHITE)]

    def input_program(self):
        selected_cmd = 0
        self.render(selected_cmd)

        start = False

        while not start:
            # 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] >= 50 and pos[0] <= 50 + 250 and pos[1] >= 600 and pos[1] <= 650:
                        print(f"clicked at pos = {pos}")
                        selected_cmd = (pos[0] - 50)//50
                    if pos[0] >= 50 and pos[0] <= 50 + 350 and pos[1] >= 700 and pos[1] <= 750:
                        print(f"clicked at pos = {pos}")
                        chosen_input = (pos[0] - 50)//50
                        chosen_input_cmd = self.available_inputs[chosen_input]
                        if selected_cmd < len(self.cmds):
                            edited_cmd = self.cmds[selected_cmd]
                            if chosen_input_cmd.action is not '-':
                                edited_cmd.action = chosen_input_cmd.action
                            else:
                                edited_cmd.color = chosen_input_cmd.color
                        else:
                            self.cmds.append(copy.copy(chosen_input_cmd))
                    if pos[0] >= 325 and pos[0] <= 400 and pos[1] >= 600 and pos[1] <= 650:
                        print(f"clicked at pos = {pos}")
                        return
            self.render(selected_cmd)



    def run(self):
        running = True
        prg_counter = 0

        self.render(prg_counter)

        while running:
            cmd = self.cmds[prg_counter]
            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]

            if np.all(cmd.color == WHITE) or np.all(cmd.color == tile['color']):
                # matching color -> execute command
                if cmd.action == 'forward':
                    ynew, xnew = r.get_forward_coordinates()
                    r.x = xnew
                    r.y = ynew
                elif cmd.action in {'left', 'right'}:
                    r.orientation = Robot.resulting_orientation[self.robot.orientation][cmd.action]
                elif cmd.action == 'P0':
                    prg_counter = 0
            else:
                print("color not matching -> skipping command")

            self.render(prg_counter)

            if (not (0 <= r.x < self.board.tiles.shape[1])) or not (0 <= r.y < self.board.tiles.shape[0]):
                print("GAME OVER")
                screen.blit(game_over_text, (50, 00))
                pygame.display.update()
                running = False

            time.sleep(0.1)



    def render(self, prg_counter=None):
        dx = 50
        dy = 50
        screen.fill(tuple(BLACK))
        board_surf = self.board.render()
        screen.blit(board_surf, (dx, dy, dx + dimx * scale_fac, dy + dimy * scale_fac))

        prg_surf = self.render_prg(prg_counter)
        screen.blit(prg_surf, (dx, board_surf.get_height()+2*dy, prg_surf.get_width(), prg_surf.get_height()))

        inp_surf = self.render_inputs()
        screen.blit(inp_surf, (dx, board_surf.get_height()+4*dy, inp_surf.get_width(), inp_surf.get_height()))

        run_surf = pygame.Surface((80, 50))
        run_surf.fill(tuple(GREEN))
        run_surf.blit(run_text, (0,10))
        screen.blit(run_surf, (325, 600, run_surf.get_height(), run_surf.get_width()))

        pygame.display.update()

    def render_inputs(self):
        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()
            inp_surf.blit(cmd_surf, (i * scale_fac, 0, scale_fac, scale_fac))
        return inp_surf

    def render_prg(self, prg_counter=None):
        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()
            else:
                cmd_surf = pygame.Surface((50,50))
                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

r = Robot(x=1, y=1, orientation='v')
b = Board(dimx,dimy, r)
b.tiles[3,4]['star'] = True
print(b)


cmds = [Command('forward'), Command('left', color=RED), Command('left', color=BLUE), Command('P0')]

prg = Program(r, b, cmds)
prg.input_program()
prg.run()