#########################
#   Matthew Callens     #
#   mcallen1@umbc.edu   #
#                       #
#   CS 471 Homework 2   #
#                       #
#   Usage in README.md  #
#########################

import random
import sys
import heapq
from timeit import default_timer
from copy import copy, deepcopy
from math import sqrt


class Node:
    def __init__(self, val, par):
        self.val = val
        self.parent = par
        self.h = 0
        self.g = 0
        self.variant = random.randint(1, 500)

    def __find(self, item):
        for i, s in enumerate(self.val):
            try:
                tmp = s.index(item)
                return (i, tmp)
            except ValueError:
                continue

    def collect(self):
        print('Collecting solution path...')
        sequence = list()

        def backtrack(node):
            if not node.parent == None:
                backtrack(node.parent)
            sequence.append(node)

        backtrack(self)
        return sequence

    def __get_moves(self):
        loc = self.__find(' ')
        row, column = loc
        available = list()

        bound = len(self.val) - 1

        if row > 0:
            available.append((row - 1, column))
        if row < bound:
            available.append((row + 1, column))
        if column > 0:
            available.append((row, column - 1))
        if column < bound:
            available.append((row, column + 1))

        return available

    def create_children(self):
        legal = self.__get_moves()
        blank = self.__find(' ')
        blank_row, blank_col = blank

        def swap_tiles(a, b):
            new_state = deepcopy(self.val)
            new_state[a[0]][a[1]] = self.val[b[0]][b[1]]
            new_state[b[0]][b[1]] = self.val[a[0]][a[1]]
            child = Node(new_state, self)
            child.g = self.g + 1
            return child

        return list(map(lambda move: swap_tiles((blank_row, blank_col), move), legal))

    def __hash__(self):
        x = self.g + self.h + self.variant
        if not self.parent == None:
            x += self.parent.g + self.parent.h
        return x

    def __str__(self):
        display = ''
        for i in range(len(self.val)):
            for j in range(len(self.val[i])):
                display += '{} '.format(self.val[i][j])
            display += '\n'
        return display

    def __lt__(self, other):
        return (self.g + self.h) < (other.g + other.h)

    def __eq__(self, other):
        if other == None: return False

        for i in range(len(self.val)):
            if not self.val[i] == other.val[i]:
                return False
        return True


def create_2d(l, n):
    jump = int(sqrt(n))
    tmp = list()
    i = 0
    while i < n:
        tmp.append(l[i:i + jump])
        i += jump
    return tmp


def initialize():
    n = int(sys.argv[1])
    random.seed(sys.argv[2])
    state = list(range(1, n)) + [' ']
    goal = copy(state)
    random.shuffle(state)
    return create_2d(state, n), create_2d(goal, n)


def manhattan(node, goal, size):
    total = 0
    for i in range(len(node.val)):
        for j in range(len(node.val[i])):
            if node.val[i][j] == ' ' or node.val[i][j] == goal.val[i][j]: continue
            goal_row = (node.val[i][j] - 1) // size
            goal_col = (node.val[i][j] - 1) % size
            total += abs(goal_row - i) + abs(goal_col - j)
    return total


def is_solvable(state):
    n = (len(state) ** 2)
    inversions = 0
    tmp = [j for i in state for j in i]

    for i in range(n - 1):
        for j in range(i + 1, n):
            if not tmp[i] == ' ' and not tmp[j] == ' ' and tmp[i] > tmp[j]:
                inversions += 1
    return inversions % 2 == 0


def a_star(start, goal, size):
    open_nodes = []
    heapq.heappush(open_nodes, (0, start))

    costs = dict()
    costs[start] = 0

    while open_nodes:
        curr_tup = heapq.heappop(open_nodes)
        curr = curr_tup[1]

        if curr == goal:
            print('Success!')
            return curr.collect()

        moves = curr.create_children()

        for m in moves:
            new_cost = costs[curr] + m.g
            if m not in costs or new_cost < costs[m]:
                m.h = manhattan(m, goal, size)
                costs[m] = new_cost + m.h
                heapq.heappush(open_nodes, (costs[m], m))

    return None


def main():
    state_list, goal_list = initialize()

    start = Node(state_list, None)
    goal = Node(goal_list, None)

    print('Start State:\n{}'.format(start))

    if not is_solvable(start.val):
        print('Failed!\nThis starting state is not solvable!')
        return
    else:
        print('Initial state is solvable. Calculating...')

    size = len(goal.val)

    start_timer = default_timer()
    result = a_star(start, goal, size)
    stop_timer = default_timer()

    if result == None:
        print('Failed!')
        print("No solution found for this state.")
    else:
        for r in result:
            print('{}\n'.format(r))
        print('{} steps for solution'.format(len(result) - 1))

    print('Runtime -> {}'.format(stop_timer - start_timer))


if __name__ == "__main__":
    main()