package dev.bogwalk.batch6

import kotlin.math.floor
import kotlin.math.sqrt
import dev.bogwalk.util.maths.isHexagonalNumber
import dev.bogwalk.util.maths.isTriangularNumber
import dev.bogwalk.util.maths.isPentagonalNumber

/**
 * Problem 61: Cyclical Figurate Numbers
 *
 * https://projecteuler.net/problem=61
 *
 * Goal: Given a set of N numbers representing a polygonal, find the sum of a set of N 4-digit
 * numbers that respect the cyclical properties described below. If multiple such sets exist,
 * return all their sums in sorted order.
 *
 * Constraints: 3 <= N <= 6, from {3,4,5,6,7,8}
 *
 * Cyclic Set: In this context, this means the last 2 digits of each number is the first 2 digits
 * of the next number (including the last number with the first).
 *
 * Figurate/Polygonal numbers are generated by the formulae ->
 *
 *      Triangle: P_3n = n(n + 1)/2 -> 1, 3, 6, 10, 15, ...
 *      Square: P_4n = n^2 -> 1, 4, 9, 16, 25, ...
 *      Pentagonal: P_5n = n(3n - 1)/2 -> 1, 5, 12, 22, 35, ...
 *      Hexagonal: P_6n = n(2n - 1) -> 1,6, 15, 28, 45, ...
 *      Heptagonal: P_7n = n(5n - 3)/2 -> 1, 7, 18, 34, 55, ...
 *      Octagonal: P_8n = n(3n - 2) -> 1, 8, 21, 40, 65, ...
 *
 *      e.g. {8128, 2882, 8281} is the only cyclical set of 3 4-digit numbers to represent
 *      polygonal types -> {P_3-127, P_5-44, P_4-91}.
 *
 * e.g.: N = {3, 4, 5}
 *       cyclical set = {8128, 2882, 8281}
 *       sum = 19291
 */

class CyclicalFigurateNumbers {
    /**
     * Solution uses recursion to search for cyclic chains using the smallest polygonal set
     * numbers as starters (e.g. there will be fewer octagonal numbers that triangle numbers to
     * start chains).
     *
     * A separate stack of the remaining polygon numbers to find is cached so the search can
     * continue if a chain that cannot close its first and last elements is reached. This is done
     * by remembering the last polygon added to the chain & adding it back to the stack it was
     * just popped from if the chain becomes invalid.
     *
     * Rather than searching for a final figurate that closes the chain, the expected number is
     * checked to see if it is in the generated polygon list.
     *
     * @return list of all cyclic lists that represent the polygons requested.
     */
    fun cyclicalFigurates(polygons: Set<Int>): List<List<Int>> {
        val cycles = mutableListOf<List<Int>>()
        var notFound = polygons.toMutableList()
        val allPolygons = getPolygons(polygons)

        fun nextCyclical(prefix: Int, lastFound: Int, chain: MutableList<Int>): MutableList<Int> {
            if (notFound.size == 1) {
                val expected = chain.last() % 100 * 100 + chain.first() / 100
                if (expected in allPolygons[notFound.first()]) {
                    notFound.removeLast()
                    chain.add(expected)
                } else {
                    chain.removeLast()
                    notFound.add(lastFound)
                }
            } else {
                for (p in polygons) {
                    if (p in notFound) {
                        for (number in allPolygons[p]) {
                            if (number in chain) continue
                            if (prefix < number / 100) break
                            if (prefix == number / 100) {
                                chain.add(number)
                                notFound.remove(p)
                                val ans = nextCyclical(number % 100, p, chain)
                                if (notFound.isEmpty()) return ans
                            }
                        }
                    }
                }
                chain.removeLast()
                notFound.add(lastFound)
            }
            return chain
        }

        val smallest = notFound.maxOrNull()!!
        for (starter in allPolygons[smallest]) {
            notFound.remove(smallest)
            val cycle = nextCyclical(starter % 100, smallest, mutableListOf(starter))
            if (cycle.size == polygons.size) {
                cycles.add(cycle)
                notFound = polygons.toMutableList()
            }
        }
        return cycles
    }

    /**
     * Generates all 4-digit figurate/polygonal numbers for the polygons specified.
     *
     * N.B. Numbers with a 0 in the 3rd position are filtered out as these will not allow the
     * cyclic property to occur. e.g. 2701 (P_6_37) -> 01xx, which is not a 4-digit number.
     *
     * @param [polygons] set of Integers in [3, 8] representing triangle,...octagonal polygons.
     * @return nested lists of all requested 4-digit figurate numbers, found at the index
     * matching their polygonal value. e.g. All triangle numbers between 1000 and 9999 stored at
     * index 3.
     */
    private fun getPolygons(polygons: Set<Int>): List<List<Int>> {
        val formulae = listOf(
            Long::isTriangularNumber, ::isSquareNumber, Long::isPentagonalNumber,
            Long::isHexagonalNumber, ::isHeptagonalNumber, ::isOctagonalNumber
        )
        val allPolygons = List(9) { p ->
            if (p in polygons) {
                (1010 until 10_000).filter { num ->
                    num.toString()[2] != '0' && formulae[p-3](num.toLong()) != null
                }
            } else emptyList()
        }
        return allPolygons
    }

    /**
     * Returns the corresponding term of the number if square, or null.
     *
     * Derivation solution is based on the formula:
     *
     *      n^2 = sN, in quadratic form becomes:
     *
     *      0 = n^2 + 0 - sN, with a, b, c = 1, 0, -sN
     *
     * putting these values in the quadratic formula becomes:
     *
     *      n = (0 +/- sqrt(0 + 4sN))/2
     *
     * so the inverse function, positive solution becomes:
     *
     *      n = sqrt(sN)
     */
    private fun isSquareNumber(sN: Long): Int? {
        val n = sqrt(1.0 * sN)
        return if (n == floor(n)) n.toInt() else null
    }

    /**
     * Returns the corresponding term of the number if heptagonal, or null.
     *
     * Derivation solution is based on the formula:
     *
     *      n(5n - 3)/2 = hN, in quadratic form becomes:
     *
     *      0 = 5n^2 - 3n - 2hN, with a, b, c = 5, -3, (-2hN)
     *
     * putting these values in the quadratic formula becomes:
     *
     *      n = (3 +/- sqrt(9 + 40hN))/10
     *
     * so the inverse function, positive solution becomes:
     *
     *      n = 0.1 * (3 + sqrt(9 + 40hN))
     */
    private fun isHeptagonalNumber(hN: Long): Int? {
        val n = 0.1 * (3 + sqrt(9 + 40.0 * hN))
        return if (n == floor(n)) n.toInt() else null
    }

    /**
     * Returns the corresponding term of the number if octagonal, or null.
     *
     * Derivation solution is based on the formula:
     *
     *      n(3n - 2) = oN, in quadratic form becomes:
     *
     *      0 = 3n^2 - 2n - oN, with a, b, c = 3, -2, -oN
     *
     * putting these values in the quadratic formula becomes:
     *
     *      n = (2 +/- sqrt(4 + 12oN))/6
     *
     * so the inverse function, positive solution becomes:
     *
     *      n = (1 + sqrt(1 + 3oN))/3
     */
    private fun isOctagonalNumber(oN: Long): Int? {
        val n = (1 + sqrt(1 + 3.0 * oN)) / 3.0
        return if (n == floor(n)) n.toInt() else null
    }
}