from gauss_turing import Program, Board, Command, Robot
import numpy as np
import json

levels = {
    0: lambda cmds: len(cmds) <= 3,
    1: lambda cmds: len(cmds) == 3
}

def rate_level(robot_path, cmds, board):
    path_length = len(set(robot_path))
    n_cmds = len(cmds)
    n_cmd_colors = len(set([tuple(c.color) for c in cmds]))

    difficulty = (path_length - 1) * (n_cmds + n_cmd_colors)

    # place coins on the robot path to create a solution
    if difficulty > 0:
        n_coins = np.random.randint(1, min(path_length, 5)) - 1

        # put one coin on last tile visited
        coins = [robot_path[-1]]

        # distribute other coins randomly on the path
        # path without first and list tile
        unique_tiles = list(set(robot_path) - {robot_path[-1]} - {robot_path[0]})

        for _ in range(n_coins):
            c = np.random.randint(0, len(unique_tiles))
            new_coin = unique_tiles.pop(c)
            coins.append(new_coin)
        pass
    else:
        coins = []

    return difficulty, coins

def generate_level(dimx, dimy, max_steps=100):
    n_cmds = np.random.randint(2, 6)
    assert n_cmds <= 5

    # generate random board without any coins
    board = Board(dimx, dimy, n_coins=0)

    cmds = []
    actions = list(sorted(Command.valid_actions - {'-'}))
    # generate random commands
    for i in range(n_cmds):
        action = np.random.choice(actions)
        color = Board.valid_colors[np.random.randint(len(Board.valid_colors))]
        cmds.append(Command(action, color))

    # generate robot at random position
    rx = np.random.randint(0, dimx-1)
    ry = np.random.randint(0, dimy-1)
    orientation = np.random.choice(['>','v','<','^'])
    r = Robot(rx, ry, orientation)

    prg = Program(r, board, cmds)
    continue_running = True
    state = 'running'
    prg_counter_old = prg.prg_counter

    robot_path = [(r.x, r.y)]
    step = 0

    while continue_running and step < max_steps:
        #print(f"prg_counter = {prg.prg_counter} - robot: {r} - state: {state}")
        state = prg.step(state, check_victory=False)

        robot_path.append((r.x, r.y))

        stuck = prg.prg_counter == prg_counter_old

        prg_counter_old = prg.prg_counter

        if state == 'game_over' or stuck:
            continue_running = False

        step += 1

    last_pos = robot_path[-1]
    if not ((0 <= last_pos[0] < dimx) and (0 <= last_pos[1] < dimy)):
        # remove last entry of path if robot leaves the board
        robot_path.pop(-1)

    difficulty, coins = rate_level(robot_path, cmds, board)

    # put coins on the board
    for coin in coins:
        board.tiles[coin[1], coin[0]]['star'] = True
    n_coins = len(coins)


    return difficulty, board, n_coins, set(robot_path), (rx, ry, orientation), cmds

if __name__ == "__main__":
    np.random.seed(2)

    levels = {}

    for i in range(100):
        diff, board, n_coins, robot_path, init_robot_pos, solution = generate_level(7, 4)
        if diff > 0:
            print("difficulty: ", diff, "n_coins: ", n_coins, "path length: ", len(robot_path))
            if diff in levels:
                if n_coins > levels[diff]['n_coins'] and len(robot_path) > levels[diff]['path_length']:
                    levels[diff] = {'board': board, 'init_robot_pos': init_robot_pos, 'solution': solution,
                                    'n_coins': n_coins, 'path_length': len(robot_path)}
            else:
                levels[diff] = {'board': board, 'init_robot_pos': init_robot_pos, 'solution': solution,
                                'n_coins': n_coins, 'path_length': len(robot_path)}

    level_info = {}
    for l, data in levels.items():
        np.save(f'levels/{l}.npy', data['board'].tiles)
        sol = [(cmd.action, tuple(map(int, cmd.color))) for cmd in data['solution']]
        level_info[l] = {'init_robot_pos': data['init_robot_pos'], 'solution': sol,
                                    'n_coins': int(data['n_coins']), 'path_length': int(data['path_length']),
                         'file': f'levels/{l}.npy'}

    with open('levels/level_info.json', 'w') as f:
        json.dump(level_info, f, indent=4)
    pass