from itertools import combinations

def build_ramsey_r555_cnf(n=52):
    """
    Build CNF for R(5,5,5) Ramsey problem with n vertices.
    Returns clauses and number of variables.
    """
    k = 5
    num_colors = 3
    
    # Generate all edges in complete graph K_n
    edges = [(i, j) for i in range(n) for j in range(i + 1, n)]
    edge_to_idx = {e: i for i, e in enumerate(edges)}
    
    def get_var(edge_idx, color):
        return edge_idx * num_colors + color + 1
    
    clauses = []
    
    # Constraint 1: Each edge must have exactly one color
    for i in range(len(edges)):
        # At least one color
        clauses.append([get_var(i, c) for c in range(num_colors)])
        # At most one color (no two colors on same edge)
        for c1, c2 in combinations(range(num_colors), 2):
            clauses.append([-get_var(i, c1), -get_var(i, c2)])
    
    # Constraint 2: No monochromatic K5 in any color
    for clique in combinations(range(n), k):
        clique_edges = [edge_to_idx[(u, v)] for u, v in combinations(clique, 2)]
        for color in range(num_colors):
            # At least one edge in this clique must NOT be this color
            clauses.append([-get_var(e, color) for e in clique_edges])
    
    return clauses, len(edges) * num_colors

if __name__ == "__main__":
    clauses, num_vars = build_ramsey_r555_cnf()
    print(f"R(5,5,5) CNF construction:")
    print(f"Variables: {num_vars:,}")
    print(f"Clauses: {len(clauses):,}")
    print(f"Expected K5 cliques: {len(list(combinations(range(52), 5))):,}")