import re
import sys

from tools.aoc import AOCDay
from tools.itertools import combinations_of_sum, len_combinations_of_sum
from typing import Any, Iterator
from tqdm.auto import tqdm


def clean_springs(springs: str, groups: tuple[int, ...]) -> str:
    hashes = []
    start = None
    c_size = 0
    for i, c in enumerate(springs):
        if c == "#":
            if start is None:
                start = i
            c_size += 1
        else:
            if start is not None:
                hashes.append((start, c_size))
                start = None
                c_size = 0

    if start is not None:
        hashes.append((start, c_size))

    # print("cleaning", springs, "with hashes", hashes, "and groups", groups)
    found = {}
    for i, s in enumerate(groups):
        for hash_start, hash_len in hashes:
            if sum(groups[: i + 1]) + i + 1 < hash_start < len(springs) - 1 - (sum(groups[i:]) + len(groups) - i):
                if hash_len == s:
                    # print(
                    #    f"PLACING HASH {i=} {groups[:i+1]} -> {sum(groups[:i + 1]) + i} <= {hash_start} <= {len(springs)=} - ({groups[i:]} -> {sum(groups[i:]) + len(groups) -  i})"
                    # )
                    found[i] = (hash_start, s)
                    if hash_start == 0:
                        if len(springs) > hash_len:
                            springs = "#" * hash_len + "." + springs[hash_start + hash_len + 1 :]
                        else:
                            springs = "#" * hash_len
                    elif hash_start + hash_len == len(springs) - 1:
                        springs = springs[: -(hash_len + 1)] + "." + "#" * hash_len
                    else:
                        springs = (
                            springs[: hash_start - 1]
                            + "."
                            + "#" * hash_len
                            + "."
                            + springs[hash_start + hash_len + 1 :]
                        )
                    # print("found index", i, "value", groups[i], "at", hash_start, "new springs", springs,)
                break

    # print("cleaned springs:", springs)
    return springs


def get_open_groups(springs: str) -> list[str]:
    open_groups = re.split(r"\.+", springs)
    while "" in open_groups:
        open_groups.remove("")
    return open_groups


def get_placement_options(groups: tuple[int, ...], open_groups: list[str]) -> Iterator[tuple]:
    print(
        f"get_placement_options({len(groups)}, {len(open_groups)}) returns",
        len_combinations_of_sum(len(groups), len(open_groups)),
    )
    t = tqdm(total=len_combinations_of_sum(len(groups), len(open_groups)), leave=False, file=sys.stdout)
    for options in combinations_of_sum(total_sum=len(groups), length=len(open_groups)):
        t.update(1)
        idx = 0
        possible = True
        for i, x in enumerate(options):
            group_str = open_groups[i]
            if sum(groups[idx : idx + x]) + x - 1 > len(group_str) or ("#" in group_str and x == 0):
                possible = False
                break  # not possible
            idx += x

        if possible:
            yield options

    t.close()


def brute_force_group(group_str: str, groups: tuple[int, ...]) -> int:
    count = 0
    p = re.compile(group_str.replace("?", "."))
    g_len = len(groups)
    t = tqdm(
        total=len_combinations_of_sum(len(group_str) - sum(groups), len(groups)),
        leave=False,
        file=sys.stdout,
        postfix="brute_force",
    )
    for x in combinations_of_sum(total_sum=len(group_str) - sum(groups), length=len(groups) + 1):
        t.update()
        test_str = ""
        valid = True
        for i, c in enumerate(x):
            test_str += " " * c
            if 0 < i < g_len and c == 0:
                valid = False
                break

            if i < len(groups):
                test_str += groups[i] * "#"

        if valid and p.match(test_str):
            count += 1

    t.close()
    return count


def get_arrangements(springs: str, groups: tuple[int, ...]) -> int:
    if sum(groups) + len(groups) - 1 == len(springs):
        return 1

    open_groups = get_open_groups(springs)
    # print(f"{groups} => {springs} =>", list((x, len(x)) for x in open_groups))

    arrangements = 0
    for placement_option in get_placement_options(groups, open_groups):
        # print("PLACEMENT OPTION", placement_option)
        sub_arrangements = 1
        group_idx = 0
        for i, x in enumerate(placement_option):
            if x == 0:
                continue
            numbers_to_place = groups[group_idx : group_idx + x]
            group_idx += x
            group_str = open_groups[i]
            cleaned_group_str = clean_springs(group_str, numbers_to_place)
            # print(f"CLEAN ({i}, {x}, {group_idx})", group_str, "=>", cleaned_group_str, "from", numbers_to_place)
            if cleaned_group_str != group_str:
                f = get_arrangements(cleaned_group_str, numbers_to_place)
                # print("RECURSE CALL", cleaned_group_str, numbers_to_place, "=>", f)
                sub_arrangements *= f
            else:
                f = brute_force_group(group_str, numbers_to_place)
                # print("MISSING CASE:", group_str, numbers_to_place, "=>", f)
                sub_arrangements *= f
        # print("SUB", placement_option, "=>", sub_arrangements)
        arrangements += sub_arrangements

    return arrangements


class Day(AOCDay):
    inputs = [
        [
            (101, "input12_debug"),
            (21, "input12_test"),
            (6981, "input12"),
        ],
        [
            (525152, "input12_test"),
            (None, "input12"),
        ],
    ]

    def parse_input(self, part2: bool = False) -> list[tuple[str, tuple[int, ...]]]:
        records = []
        for line in self.getInput():
            springs, groupd = line.split()
            groups = tuple(map(int, groupd.split(",")))
            if part2:
                springs = "?".join([springs] * 5)
                groups *= 5
            records.append((springs, groups))

        return records

    def part1(self) -> Any:
        ans = 0
        for springs, groups in self.parse_input():
            f = get_arrangements(springs, groups)
            print("RESULT", springs, "=>", f)
            ans += f

        return ans

    def part2(self) -> Any:
        ans = 0
        # explore moving windows; or some other way to make use of DP
        for i, (springs, groups) in enumerate(self.parse_input(part2=True)):
            f = get_arrangements(springs, groups)
            print(i + 1, "/", len(self.input), "=>", f)
            ans += f

        return ans


if __name__ == "__main__":
    day = Day(2023, 12)
    day.run(verbose=True)