/*
 * Implements AStar algorithm to solve the puzzle
 * 4/16/20
 */
package src

import java.util.*
import kotlin.math.abs

/**
 *
 * @author <NAME>
 */
class PuzzleSolver(private val boardSize: Int) {
    private val goalState: PuzzleState

    init {
        // creates the goal state for the specified board size
        var num = 1
        val goalPuzzle = Array(boardSize){ IntArray(boardSize) }
        for (i in 0 until boardSize) {
            goalPuzzle[i] = IntArray(boardSize)
            for (j in 0 until boardSize) {
                goalPuzzle[i][j] = num
                num++
            }
        }
        goalPuzzle[boardSize - 1][boardSize - 1] = 0
        goalState = PuzzleState(goalPuzzle)
    }

    private fun heuristic(puzzleState: PuzzleState): Int {
        val puzzle = puzzleState.puzzle
        var h = 0
        // compute the heuristic against the goal state (we calculate the goal values rather than read since it is faster)
        for (i in 0 until boardSize) {
            for (j in 0 until boardSize) {
                if (puzzle[i][j] != 0) {
                    h += manhattanDistance(i, j, puzzle[i][j] / boardSize, puzzle[i][j] % boardSize)
                }
            }
        }
        return h
    }

    fun findSolution(initialState: PuzzleState): ArrayList<PuzzleState> {
        // creates the closed set to stores nodes we've already expanded
        val closedSet = HashSet<PuzzleState>(10000)
        closedSet.add(initialState)

        // creates the open set to choose the closest state to the solution at each iteration
        val openSet = PriorityQueue<PuzzleState>(10000)
        openSet.add(initialState)

        // iterate until we find a solution or are out of puzzle states
        while (!openSet.isEmpty()) {
            // pop the closest state off the priority queue
            val currentState = openSet.poll()
            closedSet.add(currentState)

            // check if we've reached the goal state, if so we can return the solution
            if (currentState == goalState) {
                return reconstructPath(currentState)
            }

            // expand our search by getting the neighbor states for the current state
            val neighborStates = neighborStates(currentState)

            // for each neighbor, if it isn't already visited update the FScore for ranking and add it to the open set
            for (neighborState in neighborStates) {
                if (!closedSet.contains(neighborState)) {
                    neighborState.f = neighborState.g + heuristic(neighborState)
                    openSet.add(neighborState)
                }
            }
        }
        return ArrayList()
    }

    private fun neighborStates(current: PuzzleState): ArrayList<PuzzleState> {
        val neighborStates = ArrayList<PuzzleState>()
        val upPuzzle = current.shiftUp()
        if (upPuzzle != null) {
            neighborStates.add(upPuzzle)
        }

        val downPuzzle = current.shiftDown()
        if (downPuzzle != null) {
            neighborStates.add(downPuzzle)
        }

        val rightPuzzle = current.shiftRight()
        if (rightPuzzle != null) {
            neighborStates.add(rightPuzzle)
        }

        val leftPuzzle = current.shiftLeft()
        if (leftPuzzle != null) {
            neighborStates.add(leftPuzzle)
        }
        return neighborStates
    }

    private fun reconstructPath(bottomLeaf: PuzzleState?): ArrayList<PuzzleState> {
        var current = bottomLeaf
        val list = ArrayList<PuzzleState>()
        // traverse tree up through parents until we reach the root
        while (current != null) {
            list.add(current)
            current = current.parent
        }
        // reverse the list since we traversed from goal state to initial
        list.reverse()
        return list
    }

    fun generateRandomSolvable(): PuzzleState {
        val moves = Utils.rand(15, 20)
        var currentState = goalState
        for (i in 0 until moves) {
            val neighborStates = neighborStates(currentState)
            val move = Utils.rand(0, neighborStates.size - 1)
            currentState = neighborStates[move]
        }
        currentState.unlink()
        return currentState
    }

    fun isSolvable(puzzleState: PuzzleState): Boolean {
        val invCount = countInversions(puzzleState)
        return if (puzzleState.getBoardSize() % 2 == 0) {
            // even board size
            val zeroStart = puzzleState.find0Position()
            (zeroStart + invCount) % 2 != 0
        } else {
            // odd board size
            invCount % 2 == 0
        }
    }

    private fun countInversions(puzzleState: PuzzleState): Int {
        val puzzle = IntArray(boardSize * boardSize)

        // copy the 2D puzzle arrays to a laid out 1D array
        var counter = 0
        for (i in 0 until boardSize) {
            for (j in 0 until boardSize) {
                puzzle[counter] = puzzleState.puzzle[i][j]
                counter++
            }
        }
        // counts the number of inversions in the puzzle
        // an inversion is any pair i and j where i < j but i appears after j on the solved state
        var inversions = 0
        for (i in puzzle.indices) {
            if (puzzle[i] == 0) {
                continue
            }
            for (j in i + 1 until puzzle.size) {
                if (puzzle[j] == 0) {
                    continue
                }
                if (puzzle[i] > puzzle[j]) {
                    inversions++
                }
            }
        }
        return inversions
    }

    companion object {
        fun manhattanDistance(row: Int, col1: Int, row2: Int, col2: Int): Int {
            return abs(row2 - row) + abs(col2 - col1)
        }
    }
}