"""
Surprisingly there are only three numbers that can be written as the sum of fourth powers of their digits:

    1634 = 1^4 + 6^4 + 3^4 + 4^4
    8208 = 8^4 + 2^4 + 0^4 + 8^4
    9474 = 9^4 + 4^4 + 7^4 + 4^4

As 1 = 1^4 is not a sum it is not included.

The sum of these numbers is 19316.

Find the sum of all the numbers that can be written as the sum of fifth powers of their digits.
"""

import time
from itertools import combinations_with_replacement


def timer(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(f"{func.__name__} on ({kwargs}) taken: {end - start: .6f} seconds")
        return result

    return wrapper


def fifth_power(n: int, power: int = 5, strict: bool = False) -> tuple[bool, int]:
    sumx = sum(int(digit) ** power for digit in str(n))
    if strict:
        isok = (sumx == n) and (power == len(str(sumx)))
        return isok, sumx
    return sumx == n, sumx


def top_number(n: int) -> int:
    beg = (9**n) * n
    beg_len = len(str(beg))
    last = n
    while True:
        if beg_len >= n:
            beg = (9**n) * beg_len
            last = beg_len
            beg_len = len(str(beg))
            if beg_len == last:
                break
        else:
            break
    return int("9" * beg_len)


def is_match(num: int, comb: list[str]) -> bool:
    if num in [1, 0]:
        return False
    numstr = str(num).zfill(len(comb))
    numstr = sorted(numstr)
    return numstr == sorted(comb)


@timer
def main_just_compute(power: int, strict: bool = False) -> None:
    res = []
    top = top_number(power)
    for i in range(2, top + 1):
        is_sum, sumx = fifth_power(i, power, strict)
        if is_sum:
            print(sumx, end="\t")
            res.append(sumx)
    print(f"\nThe sum is {sum(res):,d}")


@timer
def main_combinations(power: int) -> None:
    top_len = len(str(top_number(power)))
    res = set()
    combs = combinations_with_replacement("0123456789", top_len)
    for comb in combs:
        tmp = sum(map(lambda x: int(x) ** power, comb))
        if is_match(tmp, list(comb)):
            print(tmp, end="\t")
            res.add(tmp)
    print(f"\nThe sum is {sum(res):,d}")


if __name__ == "__main__":
    the_power = 5
    main_just_compute(power=the_power, strict=False)
    print("=" * 50)
    main_combinations(power=the_power)