lab2

#  4, 8, 1,       1, 2, 3,
#  0, 3, 6,  ->   8, 0, 4,
#  2, 7, 5        7, 6, 5

import sys
import enum
import heapq
import psutil
import os
from time import process_time

START_TIME = 0  # время запуска программы
TIME_STOP = 0 # время окончания программы
TREE = None  # дерево решения
DEBUG = True
''' Если True, то каждый шаг придётся нажимать enter'''
H1_H2 = True
''' Выбор эвристической функции'''

class Tree:
    ''' Класс представления дерева '''
    # Все его методы - статические, т.к. дерево - единственное
    __nodes = []
    '''Все узлы'''
    __hashes = {}

    #def __init__(self):
        # node = Node(get_initial_state(), None, None, 0, 0)
        # self.__nodes = {0: [node]}

    def add_node(self, new_node: "Node"):
        ''' Добавить узел в дерево '''
        # если состояние повторяется, узел не добавляется
        if (not self.hasState(new_node.current_state)):
            self.__nodes.append(new_node)
            self.__hashes[state_hash(new_node.current_state)] = new_node

    def getNode(self, node: int) -> list["Node"]:
        '''Получить список узлов по глубине'''
        return list(self.__nodes[node])

    def getNodeByState(self, state: list) -> "Node":
        return self.__hashes[state_hash(state)]
    
    def hasState(self, state: list) -> bool:
        return state_hash(state) in self.__hashes

    def print_node(self, node: "Node"):
        ''' Вывод узла на экран '''
        parent_id: int = None
        if(node.parent_node):
            parent_id = node.parent_node.node_id
        node_prev_action: str = None
        if(node.previous_action):
            node_prev_action = node.previous_action.name
        print(f"id = {node.node_id}, parent_id = {parent_id}, " +
            f"action = {node_prev_action}, " +
            f"cost = {node.path_cost}, current_state: ")
        self.print_state(node.current_state)
        print("")

    def print_state(self, state: list):
        ''' Вывод состояния в три ряда'''
        for i in range(9):
            if (i % 3 == 0 and i != 0):
                 print("")
            print(state[i] if state[i] != 0 else " ", end=" ")

    def print_path(self, node: "Node", isReversed = False):
        ''' Вывод пути на экран '''
        path = []
        current_node = node

        while current_node.parent_node:
            path.append(current_node)
            current_node = current_node.parent_node
        path.append(current_node)
        if (isReversed):
            path = path[::-1]
        for path_node in path:
            self.print_node(path_node)
            print("\t^\n\t|\n")


class Action(enum.Enum):
    ''' Перечисление действий над полем '''
    UP = 1
    DOWN = 2
    RIGHT = 3
    LEFT = 4

class Node:
    ''' Класс представления узла '''
    current_state: list = None
    ''' Начальное состояние'''
    parent_node: "Node" = None
    ''' Указатель на родительский узел'''
    previous_action: Action = None
    '''Действие, которое было применено к родительскому узлу для формирования данного узла'''
    path_cost: int = 0
    '''Стоимость пути от начального состояния до данного узла g(n)'''
    #depth: int = 0
    '''Количество этапов пути от начального состояния (глубина)'''
    node_id: int = 0
    nodes_count = 0

    def __init__(self, state: list, parent: "Node", action: Action, cost: int):
        self.current_state = state
        self.parent_node = parent
        self.previous_action = action
        self.path_cost = cost
        #self.depth = depth
        self.node_id = Node.nodes_count
        Node.nodes_count += 1

    @classmethod
    def get_nodes_count(cls) -> int:
        ''' Статический метод класса, возвращающий количество узлов '''
        return cls.nodes_count + 1

def get_initial_state() -> list:
    ''' Получение начального состояния игры (Вариант 7) '''
    return [4, 8, 1,
            0, 3, 6,
            2, 7, 5 ]

def get_finish_state() -> list:
    ''' Получение конечного состояния игры (Вариант 7) '''
    return [1, 2, 3,
            8, 0, 4,
            7, 6, 5 ]

# O(1)
def check_final(current_state: list) -> bool:
    ''' Проверка, является ли данное состояние конечным '''
    return current_state == get_finish_state()

# O(1)
def state_hash(state: list) -> int:
    ''' Хэширование состояния (листа)'''
    hash = 7
    for i in state:
        hash = 31*hash + i
    return hash

# O(1)
def heuristics(node: "Node") -> int:
    if(not H1_H2):
        cur_value = node.path_cost + h1(node.current_state)
    else:
        cur_value = node.path_cost + h2(node.current_state)
    return cur_value

# O(1)
def h1(state: list) -> int:
    '''Возвращает кол-во фишек не на своем месте'''
    state_final = get_finish_state()
    res = 0
    for i in range(9):
        if(state_final[i] != state[i]):
            res+=1
    return res

# O(1)
def h2(state: list) -> int:#метод Манхеттена
    '''Возвращает сумму Манхеттенских расстояний'''
    state_final = get_finish_state()
    res = 0
    for i in range(9):
        state_x, state_y = i // 3, i % 3
        where = state_final.index(state[i])
        final_x, final_y = where // 3, where % 3
        res += abs(final_x-state_x) + abs(final_y-state_y)
    return res

def print_info(iterations: int, time: float):
    print(f"Итого узлов: {Node.get_nodes_count()}")
    print(f"Итого итераций: {iterations}")
    print(f"Потраченное время процессора: {time*1000} миллисекунд")
    print(f"Памяти использовано: {psutil.Process(os.getpid()).memory_info().rss} байтов")

# O(1)
def state_swap(new_states: dict, current_state: list, i: int, j: int, action: Action):
    state = list(current_state)
    state[i], state[j] = state[j], state[i]
    new_states[action] = state

# O(1)
def get_new_states(current_state: list) -> dict[Action, list[Node]]:
    ''' Получение новых состояний поля '''
    #В словаре ключ - действие для получения состояния, значение - само состояние
    new_states = {}
    pos = current_state.index(0)
    # up
    if pos not in (0, 1, 2):
        state_swap(new_states, current_state, pos, pos-3, Action.UP)
    # down
    if pos not in (6, 7, 8):
        state_swap(new_states, current_state, pos, pos+3, Action.DOWN)
    # right
    if pos not in (2, 5, 8):
        state_swap(new_states, current_state, pos, pos+1, Action.RIGHT)
    # left
    if pos not in (0, 3, 6):
        state_swap(new_states, current_state, pos, pos-1, Action.LEFT)
    return new_states

# O(1)
def get_neighbours(node: Node) -> list[Node]:
    new_states_dict = get_new_states(node.current_state)
    neighbours = []
    for new_state_action in new_states_dict:
        new_state = new_states_dict[new_state_action]
        if (TREE.hasState(new_state)):
            neighbours.append(TREE.getNodeByState(new_state))
        else:
            new_node = Node(new_state, node, new_state_action, node.path_cost+1)
            neighbours.append(new_node)
            TREE.add_node(new_node)
    return neighbours

# O(C^N), C -  константа, N - глубина найденного узла
def A_star():
    open_list = set() #список граничащих узлов
    # следующая ячейка выбирается согласно эвристической функции
    close_list = set() # список узлов, больше не учавствующие в вычислениях
    heap = [] # представление кучи
    Found = False

    current_node = Node(get_initial_state(), None, None, 0)
    TREE.add_node(current_node)
    close_list.add(current_node)
    neighbours = get_neighbours(current_node)
    for neighbour in neighbours:
        open_list.add(neighbour.node_id)
        neighbour_h = heuristics(neighbour)
        heapq.heappush(heap, (neighbour_h, neighbour.node_id, neighbour))
    iteration_count = 0
    step_i = 0

    while(not Found):
        iteration_count+=1
        heap_lowest = heapq.heappop(heap) #возвращает и удаляет наименьший элемент из кучи
        current_node: Node = heap_lowest[2] #берется 3й элемент, т.к. см. heappush выше
        open_list.remove(current_node.node_id)
        close_list.add(current_node.node_id)
        if(DEBUG):
            print(f"Длина открытого списка = {len(open_list)}")
            print(f"Длина закрытого списка = {len(close_list)}")
            print(f"Был выбран узел из-за минимального значения функции f={heap_lowest[0]}")
            print("Текущий узел : ")
            TREE.print_node(current_node)
        if (check_final(current_node.current_state)):
            if(DEBUG):
                print("Ответ найден!")
                TREE.print_node(current_node)
            Found = True
            TIME_STOP = process_time()
            TREE.print_path(current_node)
            print_info(iteration_count, TIME_STOP - START_TIME)
            break
        neighbours = get_neighbours(current_node)
        if(DEBUG):
            print("Соседи текущего узла:")
            for neighbour in neighbours:
                if(neighbour.node_id in close_list):
                    print("Этот сосед находится в закрытом списке и не будет рассматриваться")
                elif(neighbour.node_id in open_list):
                    print("Этот сосед находится в открытом списке и при необходимости будет для него будет пересчитана функция")
                else:
                    print("Этот сосед новый и ранее не рассматривался. Он будет добавлен в список открытых")
                TREE.print_node(neighbour)
                print(f"Функция f = {heuristics(neighbour)}")
        for neighbour in neighbours:
            if (neighbour.node_id in open_list):
                old_g = neighbour.path_cost
                new_g = current_node.path_cost + 1
                if (new_g < old_g):
                    neighbour.path_cost = new_g
                    neighbour.parent_node = current_node
            else:
                if(neighbour.node_id not in close_list):
                    open_list.add(neighbour.node_id)
                    neighbor_i_h = heuristics(neighbour)
                    heapq.heappush(heap, (neighbor_i_h, neighbour.node_id, neighbour))
        if(DEBUG):
            print(f"Текущий шаг: {step_i}. Нажмите Enter... ")
            input()
            print("\n")
        step_i += 1

if __name__ == '__main__':
    START_TIME = process_time()
    TREE = Tree()
    A_star()