import java.time.LocalDateTime
import java.util.*
import kotlin.math.abs
import kotlin.random.Random

typealias Tile = Pair<Int, Int>
typealias Board = Array<IntArray>
typealias Path = List<Board>

typealias Genome = List<Pair<Board, Tile>>

class Puzzle {
    companion object {
        const val MATRIX_SIZE = 3
        val START_POINT = Tile(0, 1)
        val END_POINT = Tile(0, 0)
        private const val MIN_STATES = 31
        const val POPULATION_SIZE = 300
    }

    private inner class TableNode(
        val moves: Int,
        val path: Path,
        val currentState: Board,
        val blankTile: Tile
    ) : Comparable<TableNode> {
        val f: Int
            get() = moves + currentState.manhattanDistance()

        override fun compareTo(other: TableNode): Int {
            return this.f.compareTo(other.f)
        }
    }

    private val startMatrix: Board = arrayOf(
        intArrayOf(8, 0, 6),
        intArrayOf(5, 4, 7),
        intArrayOf(2, 3, 1)
    )

    private val goalMatrix: Board = arrayOf(
        intArrayOf(0, 1, 2),
        intArrayOf(3, 4, 5),
        intArrayOf(6, 7, 8)
    )

    private fun Board.manhattanDistance(): Int {
        var manhattanDistanceSum = 0
        for (row in 0 until MATRIX_SIZE)
            for (col in 0 until MATRIX_SIZE) {
                val value = this[row][col]
                if (value != 0) {
                    val targetX: Int = value / MATRIX_SIZE
                    val targetY: Int = value % MATRIX_SIZE
                    val dx = row - targetX
                    val dy = col - targetY
                    manhattanDistanceSum += abs(dx) + abs(dy)
                }
            }
        return manhattanDistanceSum
    }

    private fun Board.copy() = Array(size) { get(it).clone() }

    private fun Board.newState(blankSpace: Tile, newTile: Tile): Board {
        val newState = this.copy()
        newState[blankSpace.first][blankSpace.second] = newState[newTile.first][newTile.second].also {
            newState[newTile.first][newTile.second] = newState[blankSpace.first][blankSpace.second]
        }
        return newState
    }

    private fun Board.misplacedTiles(): Int {
        var result = 0
        for (row in 0 until MATRIX_SIZE)
            for (col in 0 until MATRIX_SIZE)
                if (this[row][col] != goalMatrix[row][col])
                    result++
        return result - 1
    }

    private fun Tile.isInBounds(): Boolean {
        return first in 0 until MATRIX_SIZE &&
                second in 0 until MATRIX_SIZE
    }

    private val usedTables = mutableListOf<Board>()
    private fun MutableList<Board>.isNewState(element: Board): Boolean {
        return !this.map { it.contentDeepEquals(element) }.contains(true)
    }

    private fun PriorityQueue<TableNode>.addAndMarkAsPassed(node: TableNode) {
        usedTables.add(node.currentState)
        add(node)
    }

    fun solveWithAStar(): Pair<Int, Path> {
        usedTables.clear()
        val priorityQueue = PriorityQueue<TableNode>()
        priorityQueue.addAndMarkAsPassed(TableNode(0, listOf<Board>(startMatrix), startMatrix, START_POINT))
        while (priorityQueue.isNotEmpty()) {
            val currentNode = priorityQueue.peek()
            if (currentNode.currentState.manhattanDistance() == 0)
                return Pair(currentNode.moves, currentNode.path)
            priorityQueue.remove()
            Direction.values().forEach { dir ->
                val newTile = dir.nextTile(currentNode.blankTile)
                if (newTile.isInBounds()) {
                    val newBoard = currentNode.currentState.newState(currentNode.blankTile, newTile)
                    if (usedTables.isNewState(newBoard)) {
                        val path = currentNode.path.toMutableList()
                        path.add(newBoard)
                        priorityQueue.addAndMarkAsPassed(TableNode((currentNode.moves + 1), path, newBoard, newTile))
                    }
                }
            }
        }
        return Pair(0, mutableListOf())
    }

    fun solveWithBeamSearch(beamWidth: Int): Pair<Int, Path> {
        usedTables.clear()
        val priorityQueue = PriorityQueue<TableNode>()
        priorityQueue.addAndMarkAsPassed(TableNode(0, listOf<Board>(startMatrix), startMatrix, START_POINT))
        while (priorityQueue.isNotEmpty()) {
            val currentNode = priorityQueue.peek()
            if (currentNode.currentState.manhattanDistance() == 0)
                return Pair(currentNode.moves, currentNode.path)
            priorityQueue.remove()
            val options = PriorityQueue<TableNode>()
            Direction.values().forEach { dir ->
                val newTile = dir.nextTile(currentNode.blankTile)
                if (newTile.isInBounds()) {
                    val newBoard = currentNode.currentState.newState(currentNode.blankTile, newTile)
                    if (usedTables.isNewState(newBoard)) {
                        val path = currentNode.path.toMutableList()
                        path.add(newBoard)
                        options.addAndMarkAsPassed(TableNode((currentNode.moves + 1), path, newBoard, newTile))
                    }
                }
            }
            options.addAll(priorityQueue)
            priorityQueue.clear()
            repeat(beamWidth) {
                if (options.isEmpty()) {
                    return@repeat
                }
                priorityQueue.add(options.remove())
            }
        }
        return Pair(0, mutableListOf())
    }


    private fun Genome.fitness(): Int {
        return this.mapIndexed { index, element -> element.first.manhattanDistance() + index }.sum()
    }

    private fun getSplitIndexes(father: Genome, mother: Genome): Pair<Int, Int>? {
        for (index in 1 until father.size)
            for (mIndex in 1 until mother.size) {
                if (father[index].first.contentDeepEquals(mother[mIndex].first))
                    return Pair(index, mIndex)
            }
        return null
    }

    private fun crossOver(father: Genome, mother: Genome): List<Genome> {
        val splitIndex = getSplitIndexes(father, mother)
        splitIndex?.let {
            val fatherFirstHalf = father.slice(0..splitIndex.first)
            val fatherSecondHalf = father.drop(splitIndex.first + 1)
            val motherFirstHalf = mother.slice(0..splitIndex.second)
            val motherSecondHalf = mother.drop(splitIndex.second + 1)
            return listOf<Genome>(
                fatherFirstHalf.plus(motherSecondHalf),
                motherFirstHalf.plus(fatherSecondHalf)
            )
        }
        return mutableListOf()
    }

    private fun Genome.maleMutate(): Genome {
        if (4 == Random.nextInt(0, 100)) {  // 1%
            val newTile = randomDirection(this[lastIndex - 1].second)
            if (newTile.isInBounds()) {
                val newBoard = last().first.newState(last().second, newTile)
                this.mapIndexed { i, _ -> if (i == lastIndex) (Pair(newBoard, newTile)) }
                return this
            }
        }
        return this
    }

    private fun Genome.femaleMutate(): Genome {
        if (4 == Random.nextInt(0, 100)) {  // 1%
            val newTile = randomDirection(this[1].second)
            if (newTile.isInBounds()) {
                val newBoard = first().first.newState(first().second, newTile)
                this.mapIndexed { i, _ -> if (i == 0) (Pair(newBoard, newTile)) }
                return this
            }
        }
        return this
    }

    private fun generateMalePopulation(): MutableList<Genome> {
        val population = mutableListOf<Genome>()
        repeat(POPULATION_SIZE / 2) {
            population.add(generateGenomeFromBeginning())
        }
        return population
    }

    private fun generateFemalePopulation(): MutableList<Genome> {
        val population = mutableListOf<Genome>()
        repeat(POPULATION_SIZE / 2) {
            population.add(generateGenomeFromEnd())
        }
        return population
    }


    private fun randomDirection(point: Tile): Tile {
        val choices = mutableListOf<Tile>()
        Direction.values().forEach {
            val newTile: Tile = it.nextTile(point)
            if (newTile.isInBounds())
                choices.add(newTile)
        }
        return choices[Random.nextInt(choices.size)]
    }

    private fun generateGenomeFromBeginning(): Genome {
        val possibleSolution = mutableListOf(Pair(startMatrix, START_POINT))
        val usedPaths = mutableListOf<Board>(startMatrix)
        var loop = 0
        while (possibleSolution.size < MIN_STATES) {
            var flag = true
            val currentNode = possibleSolution.last()
            while (flag) {
                val newTile = randomDirection(currentNode.second)
                val newBoard = currentNode.first.newState(currentNode.second, newTile)
                if (usedPaths.isNewState(newBoard)) {
                    usedPaths.add(newBoard)
                    possibleSolution.add(Pair(newBoard, newTile))
                    loop = 0
                    flag = false
                }
                if (loop > Direction.values().size)
                    return generateGenomeFromBeginning()
                loop++
            }
        }
        return possibleSolution
    }

    private fun generateGenomeFromEnd(): Genome {
        val possibleSolution = mutableListOf(Pair(goalMatrix, END_POINT))
        val usedPaths = mutableListOf(goalMatrix)
        var loop = 0
        while (possibleSolution.size < MIN_STATES) {
            var flag = true
            val currentNode = possibleSolution.first()
            while (flag) {
                val newTile = randomDirection(currentNode.second)
                val newBoard = currentNode.first.newState(currentNode.second, newTile)
                if (usedPaths.isNewState(newBoard)) {
                    usedPaths.add(newBoard)
                    possibleSolution.add(0, Pair(newBoard, newTile))
                    loop = 0
                    flag = false
                }
                if (loop > Direction.values().size)
                    return generateGenomeFromEnd()
                loop++
            }
        }
        return possibleSolution
    }

    private fun selection(male: Genome, females: List<Genome>): List<Int> {
        val connections = mutableListOf<Int>()
        females.forEachIndexed { index, female ->
            if (getSplitIndexes(male, female) != null)
                connections.add(index)
        }
        return connections
    }

    fun solveWithGeneticAlgorithm(): Genome {
        var noBetter = 0
        val malePopulation = generateMalePopulation()
        val solutions: MutableList<Genome> = mutableListOf()
        val femalePopulation = generateFemalePopulation()
        var max: Genome = malePopulation.last()
        var previousMax: Genome
        val newMalePopulation = mutableListOf<Genome>()
        val newFemalePopulation = mutableListOf<Genome>()

        while (!(malePopulation.last().last().first.contentDeepEquals(goalMatrix) && malePopulation.last().size == MIN_STATES)) {
            if (noBetter > 10)
                return malePopulation.find { genome -> genome.last().first.contentDeepEquals(goalMatrix) }
                    ?: solutions.minBy { it.size } ?: max
            else {
                newMalePopulation.clear()
                newFemalePopulation.clear()
                for (index in 0 until malePopulation.size) {
                    val connections = selection(malePopulation[index], femalePopulation)
                    connections.forEach {
                        val children = crossOver(malePopulation[index], femalePopulation[it])
                        newMalePopulation.add(children[0])
                        newFemalePopulation.add(children[1])
                        solutions.add(children[0])
                    }
                }
                malePopulation.asSequence().drop(newMalePopulation.size).plus(newMalePopulation)
                    .sortedBy { it.fitness() }
                    .map { it.maleMutate() }.toMutableList()
                femalePopulation.asSequence().drop(newFemalePopulation.size).plus(newFemalePopulation)
                    .sortedBy { it.fitness() }
                    .map { it.femaleMutate() }.toMutableList()
                previousMax = max
                solutions.maxBy { it.size }?.let { max = it }
                if (previousMax == max)
                    noBetter++
            }
        }
        return max
    }

    fun printSolution(solution: Pair<Int, Path>) {
        println("Number of steps: ${solution.first}")
        solution.second.forEach {
            println(Arrays.deepToString(it).replace("],", "],\n"))
            println("-----------")
//            Thread.sleep(2_000)
        }
    }

    enum class Direction {
        LEFT, RIGHT, UP, DOWN;

        fun nextTile(point: Tile): Tile {
            return when (this) {
                LEFT -> Tile(point.first, point.second - 1)
                RIGHT -> Tile(point.first, point.second + 1)
                UP -> Tile(point.first - 1, point.second)
                DOWN -> Tile(point.first + 1, point.second)
            }
        }
    }
}

fun main() {
    val puzzle = Puzzle()

//    println("With A* search:")
//    puzzle.printSolution(puzzle.solveWithAStar())

//    println("With Beam search:")
//    puzzle.printSolution(puzzle.solveWithBeamSearch(3))

    println("With GA:")
    val solution = puzzle.solveWithGeneticAlgorithm()
    puzzle.printSolution(Pair(solution.size - 1, solution.map { it.first }))
}