pythagorean gaussian primes

import math
import sys
from sympy import primerange, isprime

# Function to check if a number can be written as a sum of two squares
def is_gaussian_integer(P):
    solutions = []
    for a in range(1, int(math.sqrt(P / 2))+1):
        b_squared = P - a**2
        b = int(math.sqrt(b_squared))
        if a**2 + b**2 == P:
            solutions.append((a, b))
    return solutions

# Function to check if a number can form a Pythagorean triplet
def is_pythagorean_triplet(P):
    solutions = []
    for c in range(1, int(math.sqrt(P**2 / 2))+1):
        d_squared = P**2 - c**2
        d = int(math.sqrt(d_squared))
        if c**2 + d**2 == P**2:
            solutions.append((c, d))
    return solutions

# Main function to process both primes and composites, check both conditions
def find_gaussian_pythagorean_numbers(limit):
    result = []

    for P in range(1, limit + 1):

        # Check if P is a Gaussian prime
        gaussian_solutions = is_gaussian_integer(P)

        if gaussian_solutions:
            # Check if P^2 forms a Pythagorean triplet
            pythagorean_solutions = is_pythagorean_triplet(P)
            if pythagorean_solutions:
                result.append({
                    'number': P,
                    'is_prime': isprime(P),
                    'gaussian': gaussian_solutions,
                    'pythagorean_triplet': pythagorean_solutions
                })

    return result

result = find_gaussian_pythagorean_numbers(251)

integerStr=""
primeStr=""

for entry in result:
    prime_composite = "Prime" if entry['is_prime'] else "Composite"
    print(f"{entry['number']} {prime_composite}")
    print(f"  Pythagorean triplets (a^2 + b^2 = P^2): {entry['pythagorean_triplet']}")
    print(f"  Sum of two squares (c^2 + d^2 = P): {entry['gaussian']}")

    integerStr+=str(entry['number'])+", "
    if entry['is_prime']:
        primeStr+=str(entry['number'])+", "

print("\n",end="")

print(f"OEIS A004431 {integerStr}\n")

print(f"OEIS A002144 {primeStr}\n")