from lib_prime import primes, is_prime
from dataclasses import dataclass
from heapq import heappop, heappush
from collections import namedtuple
from collections import defaultdict
from math import sqrt

P = namedtuple("P", ["x", "y"])

State = namedtuple("State", ["f", "g", "red", "space"])


@dataclass(order=True)
class State:
    g: int
    red: P
    space: P


def s(m: int, n: int) -> int:
    """What am I doing? It was necessary to find the pattern."""

    def h(s: State) -> int:
        return abs(m - s.space.x) + abs(n - s.space.y)

    def f(s: State) -> int:
        return s.g + h(s)

    def done(s: State) -> bool:
        return s.red.x == m - 1 and s.red.y == n - 1

    # assert m < n, then we could just take result times 2
    init = State(0, P(0, 0), P(m - 1, n - 1))
    states = [(0, init)]
    g_costs = {(init.red, init.space): init.g}

    while states:
        _, s = heappop(states)
        if done(s):
            return s.g

        for dx, dy in [(-1, 0), (0, 1), (1, 0), (0, -1)]:
            nx, ny = s.space.x + dx, s.space.y + dy
            if not (0 <= nx < m and 0 <= ny < n):
                continue
            nstate = State(s.g + 1, s.red, P(nx, ny))
            if s.red.x == nx and s.red.y == ny:
                nstate.red = P(s.space.x, s.space.y)
            if (
                not (nstate.red, nstate.space) in g_costs
                or g_costs[(nstate.red, nstate.space)] > nstate.g
            ):
                g_costs[nstate.red, nstate.space] = nstate.g
                heappush(states, (f(nstate), nstate))
    return -1


def solve_naiv(p_max: int) -> int:
    t = 0
    p2_max = p_max * p_max
    ps = primes(p_max)
    p2s = set([p * p for p in ps])
    for m in range(2, p2_max):
        for n in range(m, p2_max):
            r = s(m, n)
            if r in p2s:
                t += 2
            if r > p2_max:
                break
    return t


def solve(p_max: int) -> int:
    t = 2
    r_base = 0
    delta = 1
    three = 3
    for p in range(3, p_max + 1, 2):
        if is_prime(p):
            t += r_base * 2
        r_base += delta
        three -= 1
        if three == 0:
            delta += 1
            three = 3
    return t


def euler_313():
    assert solve_naiv(5) == 4
    assert solve(20) == 86
    assert solve(100) == 5482
    return solve(10**6)


if __name__ == "__main__":
    solution = euler_313()
    print(f"e313.py: {solution}")
    # assert solution == 2057774861813004