import json

# Load original data with corrected file path
with open("graphs/prime_perfect_data.json", "r") as f:
    data = json.load(f)

# New list to store inferred relationships
inferred_relationships = []

# Analyze existing relationships
for link in data["relationships"]:
    source = link["source"]
    target = link["target"]
    relationship = link["relationship"]

    # Check for multiplicative relationships
    if relationship == "multiplied_to_generate":
        factor = target / source
        if factor.is_integer():
            inferred_relationships.append({
                "source": source,
                "target": target,
                "relationship": f"multiplied_by_{int(factor)}"
            })

    # Check for next in sequence (simple addition pattern)
    elif relationship == "next_in_sequence":
        diff = target - source
        inferred_relationships.append({
            "source": source,
            "target": target,
            "relationship": f"add_{int(diff)}"
        })

    # Check for generating Mersenne primes
    elif relationship == "generates_Mersenne_prime":
        if (source > 1) and (target == 2**source - 1):
            inferred_relationships.append({
                "source": source,
                "target": target,
                "relationship": "mersenne_equation"
            })

    # Check for generating perfect numbers
    elif relationship == "generates_perfect_number":
        if (source > 1) and (target == 2**(source - 1) * (2**source - 1)):
            inferred_relationships.append({
                "source": source,
                "target": target,
                "relationship": "perfect_equation"
            })

# Combine original relationships with inferred relationships
combined_data = {
    "nodes": data["nodes"],
    "relationships": data["relationships"] + inferred_relationships
}

# Save to new JSON file in the graphs directory
with open("graphs/inferred_prime_perfect_data.json", "w") as f:
    json.dump(combined_data, f, indent=4)

print("Inferred relationships added to graphs/inferred_prime_perfect_data.json")