# This script computes the Cayley graph for an N-quandle (if finite)
# This version replicates the Mathematica program
# It stores the directed edges of the graph as a list of triples (start, end, label)

#Date: August, 2026

import time
import math
import os
from tabulate import tabulate
import networkx as nx
from networkx.drawing.nx_agraph import graphviz_layout
import matplotlib.pyplot as plt 
from pyvis.network import Network
from collections import defaultdict #for collapse function

ALPHABET = list('abcdefghijklmnopqrstuvwxyz')

# Collapse the set of edges
# Using union-find algorithm from ChatGPT
def collapse(edges, vertices):

    # ---------- Union-Find ----------
    
    parent = {v: v for v in vertices}

    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]   # path compression
            x = parent[x]
        return x

    def union(a, b):
        ra, rb = find(a), find(b)

        if ra == rb:
            return False

        # keep smaller representative
        if ra < rb:
            parent[rb] = ra
        else:
            parent[ra] = rb

        return True

    # ---------- Group edges ----------

    changed = True

    while changed:
        changed = False

        outgoing = defaultdict(list)
        incoming = defaultdict(list)

        # rebuild using CURRENT representatives
        for u, v, label in edges:
            u = find(u)
            v = find(v)

            outgoing[(u, label)].append(v)
            incoming[(v, label)].append(u)

        # merge outgoing duplicates
        for group in outgoing.values():
            base = group[0]
            for x in group[1:]:
                if union(base, x):
                    changed = True

        # merge incoming duplicates
        for group in incoming.values():
            base = group[0]
            for x in group[1:]:
                if union(base, x):
                    changed = True



    # ---------- Rebuild graph ----------

    new_edges = {
        (find(u), find(v), label)
        for u, v, label in edges
    }

    new_vertices = {
        find(v)
        for v in vertices
    }
    #print(len(new_vertices))

    return new_vertices, new_edges
    
# Add secondary relations at a vertex, then collapse redundant edges.
def add_relations(vertex, relations, Vertices, Edges):
    for rel in relations:
        v1 = vertex
        v2 = max(Vertices) + 1 # new vertex
        for w in rel[:len(rel)-1]:
            if w > 0:  # add all edges with positive labels
                Edges.add((v1, v2, w))
            else:
                Edges.add((v2, v1, -w))
            Vertices.add(v2) # add new vertex to list of vertices
            v1 = v2
            v2 = v2 + 1
        # final edge returns to original vertex
        if rel[len(rel)-1] > 0:  # add all edges with positive labels
            Edges.add((v1, vertex, rel[len(rel)-1]))
        else:
            Edges.add((vertex, v1, -rel[len(rel)-1]))
    Vertices, Edges = collapse(Edges, Vertices)
        
    return Vertices, Edges

# gen is the number of generators
# init is the list of initial relations a^{g_1g_2...g_k} = b, in the form [a, g_1, ..., g_k, b]
# sec is the list of secondary relations x^{g_1...g_k} = x, in the form [g_1, ..., g_k]
# Each generator is represented by an integer from 1 to gen
def q_graph(gen, init, sec):
    Vertices = set({}) # set of vertices.  Vertices are represented as positive integers.
    Edges = set({}) # set of edges. Each edge is a tuple (start, end, label)

    start = time.time()

    # Add loops at each generator
    # if computing a rack, rather than a quandle, only add the vertices
    for g in range(1,gen+1):
        Edges.add((g, g, g))
        Vertices.add(g)

    # Add the initial relations
    for rel in init:
        v1 = rel[0]
        v2 = max(Vertices)+1 # new vertex
        for w in rel[1:len(rel)-2]:
            if w > 0: # add all edges with positive labels
                Edges.add((v1, v2, w))
            else:
                Edges.add((v2, v1, -w))
            Vertices.add(v2)
            v1 = v2
            v2 = v2+1
        # add the final edge
        w = rel[len(rel)-2]
        v2 = rel[len(rel)-1]
        if w > 0:  # add all edges with positive labels
            Edges.add((v1, v2, w))
        else:
            Edges.add((v2, v1, -w))

    # Add the secondary relations to each vertex
    completed = set({})
    count = 1 # keep track of number of completed vertices

    while completed != Vertices:
        next_vertex = min(Vertices-completed)
        Vertices, Edges = add_relations(next_vertex, sec, Vertices, Edges)
        completed.add(next_vertex)
        completed.intersection_update(Vertices) # remove completed vertices that were collapsed

        if len(completed) > 500*count:
             print(len(Vertices),'*',len(completed), '*', time.time()-start, "seconds")
             count = count+1

    print(len(Vertices),'*',len(completed))
    print("runtime =", time.time()-start, "seconds")
    
    return Vertices, Edges

# This version only places the listed generators
# It is used when you only want to compute one component of the Cayley graph
def q_graph_gen(gen, init, sec):
    Vertices = set({})  # set of vertices.  Vertices are represented as positive integers.
    Edges = set({})  # set of edges. Each edge is a tuple (start, end, label)

    start = time.time()

    # Add loops at each generator
    # if computing a rack, rather than a quandle, only add the vertices
    for g in gen:
        Edges.add((g, g, g))
        Vertices.add(g)

    # Add the initial relations
    for rel in init:
        v1 = rel[0]
        v2 = max(Vertices) + 1  # new vertex
        for w in rel[1:len(rel) - 2]:
            if w > 0:  # add all edges with positive labels
                Edges.add((v1, v2, w))
            else:
                Edges.add((v2, v1, -w))
            Vertices.add(v2)
            v1 = v2
            v2 = v2 + 1
        # add the final edge
        w = rel[len(rel) - 2]
        v2 = rel[len(rel) - 1]
        if w > 0:  # add all edges with positive labels
            Edges.add((v1, v2, w))
        else:
            Edges.add((v2, v1, -w))

    # Add the secondary relations to each vertex
    completed = set({})
    count = 1  # keep track of number of completed vertices

    while completed != Vertices:
        next_vertex = min(Vertices - completed)
        Vertices, Edges = add_relations(next_vertex, sec, Vertices, Edges)
        completed.add(next_vertex)
        completed.intersection_update(Vertices)  # remove completed vertices that were collapsed

        if len(completed) > 500 * count:
            print(len(Vertices), '*', len(completed), '*', time.time() - start, "seconds")
            count = count + 1

    print(len(Vertices), '*', len(completed))
    print("runtime =", time.time() - start, "seconds")

    return Vertices, Edges

def generate_graph(Vertices, Edges, path, filename, count):

    Edges = [(edge[0], edge[1], ALPHABET[edge[2]-1]) for edge in Edges]

    net = Network()

    for vertex in Vertices:
        net.add_node(vertex, size=10, label='', title='')
    for edge in Edges:
        net.add_edge(edge[0], edge[1], label=edge[2])

    net.show(os.path.join(path,'html',filename+'.html'))

def draw_labeled_digraph(vertices, edges):
    """
    Draw a directed graph.

    Parameters
    ----------
    vertices : set[int]
        Set of vertex IDs.

    edges : set[tuple[int, int, int]]
        Set of labeled directed edges (a,b,c)
        where a -> b and c is the edge label.
    """

    # Create directed graph
    G = nx.DiGraph()

    # Add vertices
    G.add_nodes_from(vertices)

    # Add edges with labels
    for a, b, c in edges:
        G.add_edge(a, b, label=str(c))

    # Compute node positions
    pos = nx.spring_layout(G)

    # Draw nodes and arrows
    nx.draw(
        G,
        pos,
        with_labels=True,
        node_size=10,
        node_color="lightblue",
        arrows=True,
        arrowsize=20,
        font_size=8
    )

    # Draw edge labels
    edge_labels = nx.get_edge_attributes(G, "label")
    nx.draw_networkx_edge_labels(
        G,
        pos,
        edge_labels=edge_labels
    )

    plt.show()

def draw_interactive_digraph(vertices, edges,
                             filename="graph.html"):
    """
    Draw an interactive directed graph.

    Parameters
    ----------
    vertices : set[int]
        Set of vertex IDs

    edges : set[tuple[int,int,int]]
        Labeled edges (a,b,c)
        where a -> b and c is the edge label

    filename : str
        Output HTML file
    """

    # Create interactive network
    net = Network(
        directed=True,
        height="750px",
        width="100%",
        notebook=True
    )

    # Add vertices
    for v in vertices:
        net.add_node(
            v,
            label=str(v)
        )

    # Add labeled edges
    for a, b, c in edges:
        net.add_edge(
            a,
            b,
            label=str(c),
            arrows="to"
        )

    # Physics makes nodes movable and auto-adjust
    net.toggle_physics(True)

    # Optional controls panel
    net.show_buttons(filter_=["physics"])

    # Save and display
    net.show(filename)

def draw_interactive_graph(vertices, edges,
                           filename="graph.html"):
    """
    Draw an interactive undirected graph.

    Parameters
    ----------
    vertices : set[int]
        Set of vertex IDs

    edges : set[tuple[int,int,int]]
        Labeled edges (a,b,c)
        where c is the edge label
    """

    # Create undirected network
    net = Network(
        directed=False,
        height="750px",
        width="100%",
        notebook=True
    )

    # Add vertices
    for v in vertices:
        net.add_node(
            v,
            label=str(v)
        )

    # Add edges
    for a, b, c in edges:
        net.add_edge(
            a,
            b,
            label=str(c)
        )

    # Enable movable physics
    net.toggle_physics(True)

    # Optional controls
    net.show_buttons(filter_=["physics"])

    net.show(filename)

############################################
# Example: this quandle has 134 elements
############################################

#gen = 3
#init = [[1,-3,1,3,2,1,2,1,2,1,2],[1,3,1,-3,1,2,1,2,1,2,1,2],[3,2,1,3]]
#sec = [[1, 1], [2, 2], [3, 3, 3], [2,1,3,1,2,-3],[-3,1,3,1,-3,1,3,2,1,2,1,2,1,2,1,2,1,2,1,2],[1,3,1,-3,1,3,1,-3,1,-3,1,3,1,-3,1,3]]

############################################
# Example: this quandle has 52 elements
############################################

gen = 4
init = [[3,1,2,3]]
sec = [[1, 1], [2, 2], [3, 3], [4,4], [3,2,1,3,1,2],[4,1,2,3,1,3,1]]

Vertices, Edges = q_graph(gen, init, sec)
print(len(Vertices))
print(Vertices)
print(Edges)

generate_graph(Vertices, Edges, os.path.join('graphs','52_element_quandle'), f'52_element_quandle', 0)
draw_labeled_digraph(Vertices, Edges)
draw_interactive_graph(Vertices, Edges)