from lib_prime import is_prime_rabin_miller
from typing import Optional
from functools import cache


def compute_deltas() -> list[list[int]]:
    S = [
        set([0, 1, 2, 4, 5, 6]),  # 0
        set([2, 5]),  # 1
        set([0, 2, 3, 4, 6]),  # 2
        set([0, 2, 3, 5, 6]),  # 3
        set([1, 2, 3, 5]),  # 4
        set([0, 1, 3, 5, 6]),  # 5
        set([0, 1, 3, 4, 5, 6]),  # 6
        set([0, 1, 2, 5]),  # 7
        set([0, 1, 2, 3, 4, 5, 6]),  # 8
        set([0, 1, 2, 3, 5, 6]),  # 9
    ]

    xs = [[0] * 10 for _ in range(10)]
    for a in range(10):
        for b in range(10):
            xs[a][b] = len(S[a] - S[b]) + len(S[b] - S[a])
    return xs


NUM_TO_SEGMENTS = [6, 2, 5, 5, 4, 5, 6, 4, 7, 6]
DELTA_SEGMENTS = compute_deltas()


@cache
def digital_root(n: int) -> int:
    root = 0
    while n != 0:
        d = n % 10
        n = n // 10
        root += d
    return root


@cache
def total_segments(n: int) -> int:
    count = 0
    n_next = 0
    while n != 0:
        d = n % 10
        n = n // 10
        count += NUM_TO_SEGMENTS[d]
        n_next += d
    return count


@cache
def transition_segments(a: int, b: int) -> int:
    count = 0
    while a != 0 or b != 0:
        assert a >= b
        da, db = a % 10, b % 10
        if b == 0:
            # We are just turning off the remaining digits
            count += NUM_TO_SEGMENTS[da]
        else:
            count += DELTA_SEGMENTS[da][db]
        a, b = a // 10, b // 10
    return count


@cache
def sam_count(n: int) -> int:
    count = total_segments(n)
    count *= 2  # turn off n
    n_next = digital_root(n)
    if n == n_next:
        return count
    return count + sam_count(n_next)


@cache
def max_count(n_current: Optional[int], n: int) -> int:
    if n_current is None:
        # First number, only turn number on
        count = total_segments(n)
        n_next = digital_root(n)
        assert n != n_next
        return count + max_count(n, n_next)
    elif n_current < 10:
        # Last number, only turn number off
        assert n_current == n
        return total_segments(n)

    # This is the case where we have to compute the actual delta
    count = transition_segments(n_current, n)
    n_next = digital_root(n)
    return count + max_count(n, n_next)


def euler_315():
    assert sam_count(137) == 40
    assert max_count(None, 137) == 30

    t = 0
    for n in range(10**7, 2 * 10**7):
        if is_prime_rabin_miller(n):
            t += sam_count(n) - max_count(None, n)
    return t


if __name__ == "__main__":
    solution = euler_315()
    print(f"e315.py: {solution}")
    assert solution == 13625242