import itertools
from typing import Any, Iterator

from ..base import BaseSolver

# see http://www.se16.info/js/lands2.htm for the explanation of 'atoms' (or elements)
#
# see also https://www.youtube.com/watch?v=ea7lJkEhytA (video link from AOC) and this
# CodeGolf answer https://codegolf.stackexchange.com/a/8479/42148

# fmt: off
ATOMS: list[tuple[str, tuple[int, ...]]] = [
    ("22", (0, )),  # 0
    ("13112221133211322112211213322112", (71, 90, 0, 19, 2, )),  # 1
    ("312211322212221121123222112", (1, )),  # 2
    ("111312211312113221133211322112211213322112", (31, 19, 2, )),  # 3
    ("1321132122211322212221121123222112", (3, )),  # 4
    ("3113112211322112211213322112", (4, )),  # 5
    ("111312212221121123222112", (5, )),  # 6
    ("132112211213322112", (6, )),  # 7
    ("31121123222112", (7, )),  # 8
    ("111213322112", (8, )),  # 9
    ("123222112", (9, )),  # 10
    ("3113322112", (60, 10, )),  # 11
    ("1113222112", (11, )),  # 12
    ("1322112", (12, )),  # 13
    ("311311222112", (66, 13, )),  # 14
    ("1113122112", (14, )),  # 15
    ("132112", (15, )),  # 16
    ("3112", (16, )),  # 17
    ("1112", (17, )),  # 18
    ("12", (18, )),  # 19
    ("3113112221133112", (66, 90, 0, 19, 26, )),  # 20
    ("11131221131112", (20, )),  # 21
    ("13211312", (21, )),  # 22
    ("31132", (22, )),  # 23
    ("111311222112", (23, 13, )),  # 24
    ("13122112", (24, )),  # 25
    ("32112", (25, )),  # 26
    ("11133112", (29, 26, )),  # 27
    ("131112", (27, )),  # 28
    ("312", (28, )),  # 29
    ("13221133122211332", (62, 19, 88, 0, 19, 29, )),  # 30
    ("31131122211311122113222", (66, 30, )),  # 31
    ("11131221131211322113322112", (31, 10, )),  # 32
    ("13211321222113222112", (32, )),  # 33
    ("3113112211322112", (33, )),  # 34
    ("11131221222112", (34, )),  # 35
    ("1321122112", (35, )),  # 36
    ("3112112", (36, )),  # 37
    ("1112133", (37, 91, )),  # 38
    ("12322211331222113112211", (38, 0, 19, 42, )),  # 39
    ("1113122113322113111221131221", (67, 39, )),  # 40
    ("13211322211312113211", (40, )),  # 41
    ("311322113212221", (41, )),  # 42
    ("132211331222113112211", (62, 19, 42, )),  # 43
    ("311311222113111221131221", (66, 43, )),  # 44
    ("111312211312113211", (44, )),  # 45
    ("132113212221", (45, )),  # 46
    ("3113112211", (46, )),  # 47
    ("11131221", (47, )),  # 48
    ("13211", (48, )),  # 49
    ("3112221", (60, 49, )),  # 50
    ("1322113312211", (62, 19, 50, )),  # 51
    ("311311222113111221", (66, 51, )),  # 52
    ("11131221131211", (52, )),  # 53
    ("13211321", (53, )),  # 54
    ("311311", (54, )),  # 55
    ("11131", (55, )),  # 56
    ("1321133112", (56, 0, 19, 26, )),  # 57
    ("31131112", (57, )),  # 58
    ("111312", (58, )),  # 59
    ("132", (59, )),  # 60
    ("311332", (60, 19, 29, )),  # 61
    ("1113222", (61, )),  # 62
    ("13221133112", (62, 19, 26, )),  # 63
    ("3113112221131112", (66, 63, )),  # 64
    ("111312211312", (64, )),  # 65
    ("1321132", (65, )),  # 66
    ("311311222", (66, 60, )),  # 67
    ("11131221133112", (67, 19, 26, )),  # 68
    ("1321131112", (68, )),  # 69
    ("311312", (69, )),  # 70
    ("11132", (70, )),  # 71
    ("13112221133211322112211213322113", (71, 90, 0, 19, 73, )),  # 72
    ("312211322212221121123222113", (72, )),  # 73
    ("111312211312113221133211322112211213322113", (31, 19, 73, )),  # 74
    ("1321132122211322212221121123222113", (74, )),  # 75
    ("3113112211322112211213322113", (75, )),  # 76
    ("111312212221121123222113", (76, )),  # 77
    ("132112211213322113", (77, )),  # 78
    ("31121123222113", (78, )),  # 79
    ("111213322113", (79, )),  # 80
    ("123222113", (80, )),  # 81
    ("3113322113", (60, 81, )),  # 82
    ("1113222113", (82, )),  # 83
    ("1322113", (83, )),  # 84
    ("311311222113", (66, 84, )),  # 85
    ("1113122113", (85, )),  # 86
    ("132113", (86, )),  # 87
    ("3113", (87, )),  # 88
    ("1113", (88, )),  # 89
    ("13", (89, )),  # 90
    ("3", (90, )),  # 91
]
# fmt: on

STARTERS = [
    "1",
    "11",
    "21",
    "1211",
    "111221",
    "312211",
    "13112221",
    "1113213211",
    "31131211131221",
]


def look_and_say_length(s: str, n: int) -> int:
    if n == 0:
        return len(s)

    if s in STARTERS:
        return look_and_say_length(
            "".join(f"{len(list(g))}{k}" for k, g in itertools.groupby(s)), n - 1
        )

    counts = {i: 0 for i in range(len(ATOMS))}
    idx = next(i for i, (a, _) in enumerate(ATOMS) if s == a)
    counts[idx] = 1

    for _ in range(n):
        c2 = {i: 0 for i in range(len(ATOMS))}
        for i in counts:
            for j in ATOMS[i][1]:
                c2[j] += counts[i]
        counts = c2

    return sum(counts[i] * len(a[0]) for i, a in enumerate(ATOMS))


class Solver(BaseSolver):
    def solve(self, input: str) -> Iterator[Any] | None:
        yield look_and_say_length(input, 40)
        yield look_and_say_length(input, 50)