YAEP/composeApp/src/commonMain/kotlin/domain/generator.kt

package domain

import domain.PuzzleSolution.MULTIPLE_SOLUTIONS
import domain.PuzzleSolution.NO_SOLUTION
import domain.PuzzleSolution.SOLVABLE
import kotlin.random.Random

fun generateGame(size: Int = 6): Game {
    // Generate a random puzzle instance.
    val classes = ItemClass.randomClasses(size)

    val grid: List<List<Item<ItemClass<*>>>> = classes.map {
        it.randomItems(size).map { item -> Item(item) }
    }

    // Build a set C of all possible clues that pertain to this puzzle instance.
    // (There are a finite and in fact quite small number of possible clues: for example
    // if there are 5 houses, there are 5 possible clues of the form "Person A lives in
    // house B", 8 possible clues of the form "Person A lives next to house B", and so on.)
    var clues = getAllClues(grid).shuffled()

    var i = 0

    // If i >= n, we are done. The set of clues is minimal.
    while (i < clues.size) {
        // Run your solver on the reduced set of clues and count the number of possible solutions.
        val temp = clues - clues[i]
        if (solve(grid.toGrid(), temp) == SOLVABLE) {
            // If there is exactly one solution, update clues and reset index.
            clues = temp
            i = 0
        }
        i++
    }
    return Game(grid.toGrid(), clues)

    // (You can speed this up by removing clues in batches rather than one at a time, but it makes the algorithm more complicated to describe.)
}

// FIXME: I need to better include the options into the solver (rule violations checks)
private fun solve(
    grid: Grid,
    clues: Collection<Clue>
): PuzzleSolution {
    // Start with a grid where each cell is a list of all possible items.

    // First, set the positions of the items that are already known.
    clues.filterIsInstance<PositionClue<ItemClass<*>>>().forEach { position ->
        val row = grid[position.item.itemType.companion]
        val newSelections = mutableListOf<GameCell<ItemClass<*>>>()
        row.forEachIndexed { index, gameCell ->
            if (index == position.index) {
                gameCell.selection = position.item
                newSelections.add(gameCell)
            }
        }
        newSelections.forEach { row.cleanupOptions(it) }
    }

    // For each clue, remove any items that violate the clue.
    // If any cell has only one item left, remove that item from all other cells.
    // Repeat until no more items can be removed.
    val otherClues = clues.filter { it !is PositionClue<*> }
    var removedOptions: Boolean
    do {
        removedOptions = false
        grid.forEach { row ->
            removedOptions = row.tryOptionsForClues(grid, otherClues) || removedOptions
        }
    } while (removedOptions)

    // If any cell has no items left, the puzzle has no solution.
    if (grid.flatMap { it }.any { cell -> cell.options.isEmpty() }) {
        return NO_SOLUTION
    }

    // If all cells have exactly one item left, the puzzle is solved.
    if (grid.flatMap { it }.all { it.selection != null }) return SOLVABLE

    // If there are still cells with multiple items, pick one and try each item in turn, then go back to step 2.
    // TODO: Does this need to be implemented? We would need to try all possible options to check if there are multiple solutions.
    return MULTIPLE_SOLUTIONS
}

private enum class PuzzleSolution {
    NO_SOLUTION,
    SOLVABLE,
    MULTIPLE_SOLUTIONS
}

fun <C : ItemClass<C>> GameRow<C>.tryOptionsForClues(grid: Grid, clues: List<Clue>): Boolean {
    var removedOptions = false
    filter { cell -> cell.selection == null }.forEach { cell ->
        val removed = cell.options.removeIf { option ->
            cell.selection = option
            clues.any { it.isRuleViolated(grid) }
        }
        cell.selection = null
        if (removed) {
            cleanupOptions(cell)
        }
        removedOptions = removedOptions || removed
    }
    return removedOptions
}

fun <C : ItemClass<C>> GameRow<C>.cleanupOptions(cell: GameCell<C>) {
    if (cell.options.size == 1) {
        cell.selection = cell.options.first()
        forEach { otherCell ->
            if (otherCell != cell && otherCell.selection == null) {
                otherCell.options.remove(cell.selection)
                cleanupOptions(otherCell)
            }
        }
    }
}

private fun getAllClues(rows: List<List<Item<ItemClass<*>>>>): MutableSet<Clue> {
    val clues = mutableSetOf<Clue>()
//    rows.forEach { row ->
//        row.forEachIndexed { i, item ->
//            clues.add(PositionClue(item, i))
//        }
//    }

    rows.forEach { columns ->
        columns.forEachIndexed { j, item ->
            // Clue: Neighbours
            if (j > 0) {
                rows.map { it[j - 1] }.forEach {
                    // We don't want to give away the order with this clue,
                    // but it needs to be displayed consistently.
                    if (Random.nextBoolean()) {
                        clues.add(NeighbourClue(item, it))
                    } else {
                        clues.add(NeighbourClue(it, item))
                    }
                }
            }

            // Clue: Order
            if (j > 0) {
                rows.flatMap { it.take(j - 1) }.forEach {
                    clues.add(OrderClue(it, item))
                }
            }

            // Clue: Triplet
            if (j > 1) {
                val lefts = rows.map { it[j - 2] }
                val middles = rows.map { it[j - 1] }

                lefts.forEach { left ->
                    middles.forEach { middle ->
                        // We don't want to give away the order with this clue,
                        // but it needs to be displayed consistently.
                        if (Random.nextBoolean()) {
                            clues.add(TripletClue(left, middle, item))
                        } else {
                            clues.add(TripletClue(item, middle, left))
                        }
                    }
                }
            }

            // Clue: Same Column
            rows.map { it[j] }.forEach {
                if (it != item) {
                    clues.add(SameRowClue(item, it))
                }
            }
        }
    }
    return clues
}