import math
import concurrent.futures
from collections import defaultdict


def attempt_1(a, b, c):
    """ For the triangle with the sites a, b, c, calculate
        the coordinates of C where we define A to be at (0, 0),
        and B at (c, 0). """
    assert a + b > c
    assert a + c > b
    assert b + c > a
    S2 = [s**2 for s in range(120000)]
    S4 = [s**4 for s in range(120000)]

    a2 = S2[a]
    b2 = S2[b]
    c2 = S2[c]

    a4 = S4[a]
    b4 = S4[b]
    c4 = S4[c]

    # I've derived this myself but unfortunately I was completely on the wrong
    # track.
    x = -3 * (a4 + b4 + c4) + 6 * (a2*b2 + a2*c2 + b2*c2)
    y = math.isqrt(x)
    # y = math.sqrt(-3 * (a4 + b4 + c4) + 6 * (a2*b2 + a2*c2 + b2*c2))
    if not y * y == x:
        return "no int"

    d = math.isqrt((c2 + b2 + a2 + y) // 2)
    if d * d * 2 == c2 + b2 + a2 + y:
        return d
    else:
        return "no int"


def is_square(n):
    i = math.isqrt(n)
    return i * i == n


def check_range(start, end, limit):
    # When you realize that the angles in the middle are all 120 degress, you
    # can derive the following formula with the law of cosines:
    # c**2 = p**2 + p * r + r**2
    rs = set()
    for q in range(start, end):
        for r in range(1, q):
            if q + r > limit:
                break
            a2 = q**2 + q * r + r**2
            if not is_square(a2):
                continue
            for p in range(1, r):
                if q + r + p > limit:
                    break
                b2 = p**2 + p * q + q**2
                if not is_square(b2):
                    continue

                c2 = p**2 + p * r + r**2
                if is_square(c2):
                    d = p + q + r
                    # assert 3*(a**4 + b**4 + c**4 + d**4) == (a**2 + b**2 + c**2 + d**2)**2
                    print(math.sqrt(a2), math.sqrt(b2), math.sqrt(c2), d)
                    if d < limit:
                        rs.add(d)
    return rs


def euler_143_brute_foce():
    m = 120_000
    range_size = 1000
    s = set()
    with concurrent.futures.ProcessPoolExecutor() as executor:
        futures = []
        for start in range(1, m, range_size):
            end = min(start + range_size, m)
            futures.append(executor.submit(check_range, start, end, m))
        for future in concurrent.futures.as_completed(futures):
            s |= future.result()
    return sum(s)


def euler_143():
    pairs = defaultdict(set)
    limit = 120_000

    # Instead of bruteforcing, we can calculate 120 degree triangles directly.
    #
    # https://en.wikipedia.org/wiki/Integer_triangle#Integer_triangles_with_a_120%C2%B0_angle
    #
    # We then map one 120 angle adjacent site to another one.
    for m in range(1, math.isqrt(limit)):
        for n in range(1, m):
            a = m**2 + m * n + n**2
            b = 2*m*n + n**2
            c = m**2 - n**2
            assert a**2 == b**2 + b * c + c**2

            k = 1
            while k * (b + c) < limit:
                pairs[k * b].add(k * c)
                pairs[k * c].add(k * b)
                k += 1

    # Which ultimately allows us to construct all the triangles by iterating
    # over the sides and looking up the respective next side.
    xs = set()
    for q in pairs:
        for r in pairs[q]:
            for p in pairs[r]:
                if q in pairs[p]:
                    s = q + r + p
                    if s < limit:
                        xs.add(s)

    return sum(xs)


if __name__ == "__main__":
    solution = euler_143()
    print("e143.py: " + str(solution))
    assert solution == 30758397