from collections import deque
import heapq

class Tile:
    def __init__(self, description: str):
        self.description = description  
        self.is_player = False
        self.is_ball = False
        self.is_goal = False
        self.is_obstacle = False
        if len(description) > 1:
            for letter in range(1, len(description)):
                if description[letter] == 'P':
                    self.is_player = True
                if description[letter] == 'B':
                    self.is_ball = True
                if description[letter] == 'G':
                    self.is_goal = True
        if description[0] == 'X':
            self.penalty = 0
            self.is_obstacle = True
        else:
            self.penalty = int(description[0])

    def place_player(self):
        self.is_player = True
        self.description = self.description + "P"
    def remove_player(self):
        self.is_player = False
        self.description = self.description.replace('P', '')

    def place_ball(self):
        self.is_ball = True
        self.description = self.description + "B"
    def remove_ball(self):
        self.is_ball = False
        self.description = self.description.replace('B', '')

    def place_obstacle(self):
        self.is_obstacle = True
        self.description = "X"
        if self.is_ball or self.is_player: return True
        else: return False
    def remove_obstacle(self):
        self.is_obstacle = False
        self.description = self.description.replace('X', '0')
    
    def describe(self):
        return self.description
    
class hockey_field:
    def __init__(self, field_map, total_cost: int = 0, obstacle_next_move: str = "Left", past_moves: list = []):
        self.map_dims = (len(field_map), len(field_map[0]))
        self.map = []
        self.player_position = None
        self.obstacle_positions = []
        self.ball_positions = []
        self.goal_positions = []
        self.obstacle_next_move = obstacle_next_move
        self.total_cost = total_cost
        self.is_invalid = is_invalid
        self.total_goal_count = 0
        self.successful_goals = 0
        self.past_moves = past_moves

        for row in range(self.map_dims[0]):
            map_row = []
            for column in range(self.map_dims[1]):
                tile = Tile(field_map[row][column])
                map_row.append(tile)
                if tile.is_player:
                    self.player_position = (row, column)
                if tile.is_obstacle:
                    self.obstacle_positions.append((row, column))
                if tile.is_goal:
                    self.goal_positions.append((row, column))
                    self.total_goal_count += 1
                if tile.is_ball:
                    self.ball_positions.append((row, column))
                if tile.is_goal and tile.is_ball:
                    self.successful_goals += 1
            self.map.append(map_row)
        self.is_win = False
        if self.successful_goals == self.total_goal_count:
            self.is_win = True

    def __lt__(self, other):
        return self.total_cost < other.total_cost
    
    def __eq__(self, other):
        return self.total_cost == other.total_cost
    
    def move_player(self, action: str):
        if action == "R":
            move_player = (0, 1)
            move_ball = (0, 2)
        if action == "D":
            move_player = (1, 0)
            move_ball = (2, 0)
        if action == "U":
            move_player = (-1, 0)
            move_ball = (-2, 0)
        if action == "L":
            move_player = (0, -1)
            move_ball = (0, -2)
        player_position = self.player_position
        next_player_position = [x + y for x, y in zip(player_position, move_player)]
        next_ball_position = [x + y for x, y in zip(player_position, move_ball)]

        if next_player_position[0] < 0 or next_player_position[0] >= self.map_dims[0]: self.is_invalid = True
        if next_player_position[1] < 0 or next_player_position[1] >= self.map_dims[1]: self.is_invalid = True
        if next_ball_position[0] < 0 or next_ball_position[0] >= self.map_dims[0]: self.is_invalid = True
        if next_ball_position[1] < 0 or next_ball_position[1] >= self.map_dims[1]: self.is_invalid = True

        if self.is_invalid: return

        if self.map[next_player_position[0]][next_player_position[1]].is_ball and self.map[next_player_position[0]][next_player_position[1]].is_goal: self.is_invalid = True
        self.map[next_ball_position[0]][next_ball_position[1]].is_ball: self.is_invalid = True

        if self.is_invalid: return

        self.map[player_position[0]][player_position[1]].remove_player()
        self.map[next_player_position[0]][next_player_position[1]].place_player()
        self.total_cost += self.map[next_player_position[0]][next_player_position[1]].penalty
        self.player_position = (next_player_position[0], next_player_position[1])

        if self.map[next_player_position[0]][next_player_position[1]].is_ball:
            self.map[next_player_position[0]][next_player_position[1]].remove_ball()
            self.ball_positions.remove((next_player_position[0], next_player_position[1]))
            self.map[next_ball_position[0]][next_ball_position[1]].place_ball()
            self.ball_positions.append((next_ball_position[0], next_ball_position[1]))

            if self.map[next_ball_position[0]][next_ball_position[1]].is_goal: 
                self.successful_goals += 1
                if self.successful_goals == self.total_goal_count:
                    self.is_win = True

    def move_obstacles(self):
        if self.obstacle_next_move == "Right":
            move_obstacle = (0, 1)
            self.obstacle_next_move = "Up"
        elif self.obstacle_next_move == "Down":
            move_obstacle = (1, 0)
            self.obstacle_next_move = "Right"
        elif self.obstacle_next_move == "Up":
            move_obstacle = (-1, 0)
            self.obstacle_next_move = "Left"
        elif self.obstacle_next_move == "Left":
            move_obstacle = (0, -1)
            self.obstacle_next_move = "Down"
        
        for obstacle in self.obstacle_positions:
            next_obstacle_position = [x + y for x, y in zip(obstacle, move_obstacle)]
            self.map[obstacle[0]][obstacle[1]].remove_obstacle()
            if self.map[next_obstacle_position[0]][next_obstacle_position[1]].place_obstacle():
                self.is_invalid = True

    def step(self, action: str):
        self.move_player(action)
        self.move_obstacles()
    
    def successor_function(self):
        legal_moves = ['R', 'D', "L", "U"]
        next_states = []
        for move in legal_moves:
            next_step = hockey_field(self.show_map(), self.total_cost, self.obstacle_next_move, self.past_moves + [move])
            next_step.step(move)
            if not next_step.is_invalid:
                next_states.append(next_step)
        return next_states
    
    def heuristic(self):
        player_distance = []
        ball_distance = []
        for ball_position in self.ball_positions:
            p_distance = abs(self.player_position[0] - ball_position[0]) + abs(self.player_position[1] - ball_position[1])
            player_distance.append(p_distance)
            for goal_position in self.goal_positions:
                b_distance = abs(ball_position[0] - goal_position[0]) + abs(ball_position[1] - goal_position[1])
                ball_distance.append(b_distance)
        return (sum(player_distance)/len(player_distance) + 10 * sum(ball_distance)/len(ball_distance))

    def show_map(self):
        field_map = []
        for row in self.map:
            field_row = []
            for tile in row:
                field_row.append(tile.describe())
            field_map.append(field_row)
        return field_map
    
class hockey_solver:
    def __init__(self, field: hockey_field):
        self.hockey_field = field
    def BFS(self):
        initial_state = self.hockey_field
        fringe = deque([initial_state])
        while fringe:
            n = len(fringe)
            for _ in range(n):
                curr_field = fringe.popleft()
                depth = len(curr_field.past_moves)
                if curr_field.is_win or depth > 10:
                    print("BFS Answer:")
                    print("Steps:", curr_field.past_moves)
                    print("Total Cost:", curr_field.total_cost)
                    print("Total Depth:", len(curr_field.past_moves))
                    return curr_field

                next_steps = curr_field.successor_function()
                for next_step in next_steps:
                    fringe.append(next_step)
    
    def DFS(self):
        initial_state = self.hockey_field
        search_queue = deque([initial_state])
        while search_queue:
            curr_field = search_queue.pop()
            depth = len(curr_field.past_moves)
            if curr_field.is_win or depth >= 50:
                print("DFS Answer:")
                print("Steps:", curr_field.past_moves)
                print("Total Cost:", curr_field.total_cost)
                print("Total Depth:", len(curr_field.past_moves))
                return curr_field
            
            next_steps = curr_field.successor_function()
            for next_step in next_steps:
                search_queue.append(next_step)
    
    def IDS(self):
        initial_state = self.hockey_field
        for depth in range(12):
            search_queue = deque([initial_state])
            curr_depth = 0
            while search_queue:
                curr_field = search_queue.pop()
                curr_depth = len(curr_field.past_moves)
                if curr_field.is_win or curr_depth > 10:
                    print("IDS Answer:")
                    print("Steps:", curr_field.past_moves)
                    print("Total Cost:", curr_field.total_cost)
                    print("Total Depth:", len(curr_field.past_moves))
                    return curr_field
                if curr_depth < depth:
                    next_steps = curr_field.successor_function()
                    for next_step in next_steps:
                        search_queue.append(next_step)
    
    def UCS(self):
        initial_state = self.hockey_field
        priority_queue = [initial_state]
        while priority_queue:
            curr_field = heapq.heappop(priority_queue)
            depth = len(curr_field.past_moves)
            if curr_field.is_win or depth > 50:
                print("UCS Answer:")
                print("Steps:", curr_field.past_moves)
                print("Total Cost:", curr_field.total_cost)
                print("Total Depth:", len(curr_field.past_moves))
                return curr_field
                
            next_steps = curr_field.successor_function()
            for next_step in next_steps:
                heapq.heappush(priority_queue, next_step)
    
    def A_star(self):
        initial_state = self.hockey_field
        priority_queue = [(0, initial_state)]
        while priority_queue:
            _, curr_field = heapq.heappop(priority_queue)
            depth = len(curr_field.past_moves)
            if curr_field.is_win or depth > 50:
                print("A* Answer:")
                print("Steps:", curr_field.past_moves)
                print("Total Cost:", curr_field.total_cost)
                print("Total Depth:", len(curr_field.past_moves))
                return curr_field
                
            next_steps = curr_field.successor_function()
            for next_step in next_steps:
                field_score = next_step.heuristic() + next_step.total_cost
                heapq.heappush(priority_queue, (field_score, next_step))
    
    def Best_first(self):
        initial_state = self.hockey_field
        priority_queue = [(0, initial_state)]
        while priority_queue:
            _, curr_field = heapq.heappop(priority_queue)
            depth = len(curr_field.past_moves)
            if curr_field.is_win or depth > 50:
                print("Best First Search Answer:")
                print("Steps:", curr_field.past_moves)
                print("Total Cost:", curr_field.total_cost)
                print("Total Depth:", len(curr_field.past_moves))
                return curr_field
                
            next_steps = curr_field.successor_function()
            for next_step in next_steps:
                field_score = next_step.heuristic()
                heapq.heappush(priority_queue, (field_score, next_step))
    
    def IDA(self):
        initial_state = self.hockey_field
        for max_distance in range(100):
            priority_queue = [(0, initial_state)]
            while priority_queue:
                _, curr_field = heapq.heappop(priority_queue)
                depth = len(curr_field.past_moves)
                if curr_field.is_win or depth > 50:
                    print("IDA Answer:")
                    print("Steps:", curr_field.past_moves)
                    print("Total Cost:", curr_field.total_cost)
                    print("Total Depth:", len(curr_field.past_moves))
                    return curr_field

                curr_field_score = curr_field.heuristic() + curr_field.total_cost
                if curr_field_score < max_distance:
                    next_steps = curr_field.successor_function()
                    for next_step in next_steps:
                        field_score = next_step.heuristic() + next_step.total_cost
                        heapq.heappush(priority_queue, (field_score, next_step))

my_map = [["1P", "1B", "1", "0G"], ["0", "X", "2", "1"], ["0", "0", "1", "1"]]
my_field = hockey_field(my_map)
testing_field = ["1P 1 1 1 0 X 1 1 1 1".split(' '), "0 X 1 1 0 0 0 1 0 X".split(' '), "0 0 1 2B 2 2 2B 1 0 0".split(' '),
                 "1 1 0 X X 2 2 1 1G 1".split(' '), "1 1 0 0 0 2 1 1 1 1".split(' '), "1 1 1 1 1 1G 1 1 1 1".split(' ')]
#for field in my_field.successor_function():
#    print(field.show_map(), field.total_cost)
testing_field = hockey_field(testing_field)
print(testing_field.show_map())
solver = hockey_solver(testing_field)
#solver.BFS()
#solver.DFS()
#solver.IDS()
#solver.UCS()
answer_field = solver.Best_first()
print(answer_field.show_map())
#solver.Best_first()
#solver.IDA()
#tile = Tile("0GB")
#print(tile.is_goal, tile.is_ball, tile.penalty)

this code is supposed to be an excercise on various search algorithms. B means ball and G means goal. X means obstacle, and numbers mean cost. balls should go to goals. non of my algorithms have been working properly since they keep falling into a perpetual left right loop in 0 cost tiles. fix it for me and point out what my mistake was withought changing the input of the problem which is the testing_field variable.