import networkx as nx
import pyreason as pr
import gradio as gr
import pandas as pd

# ================================ CREATE GRAPH ====================================
# Create an Undirected Graph for Mutual Relationships
g = nx.Graph()

# Add workers
workers = ['John', 'Mary', 'Justin']
g.add_nodes_from(workers)

# Add equipment
equipment = ['Crane', 'Bulldozer']
g.add_nodes_from(equipment)

# Add relationships among workers (colleagues)
g.add_edge('John', 'Mary', colleagues=1)
g.add_edge('Mary', 'Justin', colleagues=1)
g.add_edge('John', 'Justin', colleagues=1)

# Define equipment usage (initially, all are using the Crane)
g.add_edge('John', 'Crane', uses=1)
g.add_edge('Mary', 'Crane', uses=1)
g.add_edge('Justin', 'Crane', uses=1)

# ================================ VERIFY GRAPH ====================================
def display_graph_edges(graph):
    print("\nGraph edges with attributes:")
    for u, v, d in graph.edges(data=True):
        print(f"  {u} -- {v}, attributes: {d}")

display_graph_edges(g)

# ================================ LOAD GRAPH INTO PYREASON ====================================
pr.load_graph(g)
pr.settings.verbose = False      # Set to False for cleaner output
pr.settings.atom_trace = False   # Disable atom tracing for clarity

# ================================ DEFINE RULES ====================================
# Define the 'certification' propagation rule with bounds
certification_rule_str = (
    'certified(x) : [1,1] <-1 certified(y) : [1,1], '
    'colleagues(x,y) : [1,1], uses(x,z) : [1,1], uses(y,z) : [1,1]'
)
certification_rule = pr.Rule(certification_rule_str, 'certification_rule')
pr.add_rule(certification_rule)

# Define the 'persistence' rule to maintain certifications over time with bounds
persistence_rule_str = 'certified(x) : [1,1] <-1 certified(x) : [1,1]'
persistence_rule = pr.Rule(persistence_rule_str, 'persistence_rule')
pr.add_rule(persistence_rule)

# ================================ ADD FACTS ====================================
# Add initial fact that Mary is certified from timestep 0 to 999 (indefinite)
initial_fact = pr.Fact(
    fact_text='certified(Mary) : true',
    name='certified_fact',
    start_time=0,
    end_time=999  # Use a large number to indicate indefinite duration
)
pr.add_fact(initial_fact)

# ================================ DEFINE DISPLAY FUNCTION ====================================
# Function to capture active 'uses' edges and certification statuses
def capture_state(graph, timestep):
    # Capture active 'uses' edges
    active_uses = [
        {'User': u, 'Equipment': v} for u, v, d in graph.edges(data=True)
        if d.get('uses', 0) == 1
    ]
    
    # Capture certification statuses
    interpretation = pr.reason(timesteps=1)
    dataframes = pr.filter_and_sort_nodes(interpretation, ['certified'])
    certification_status = {}
    
    for df in dataframes:
        for index, row in df.iterrows():
            component = row['component']
            status = row['certified']
            # Assuming the bounds [1,1] imply status is binary
            certification_status[component] = 'Yes' if status[0] >= 1 else 'No'
    
    # Convert certification_status dictionary to a list of dictionaries
    certifications = [
        {'Worker': worker, 'Certified': status}
        for worker, status in certification_status.items()
    ]
    
    # Convert lists to DataFrames
    uses_df = pd.DataFrame(active_uses) if active_uses else pd.DataFrame(columns=['User', 'Equipment'])
    cert_df = pd.DataFrame(certifications) if certifications else pd.DataFrame(columns=['Worker', 'Certified'])
    
    return uses_df, cert_df

# ================================ RUN REASONING AND CAPTURE STATES ====================================
timesteps = 4
state_data = []

for t in range(timesteps):
    print(f"\n=== TIMESTEP - {t} ===")
    
    # Dynamically update equipment usage based on timestep
    if t == 1:
        # At timestep 1, John stops using the Crane
        if g.has_edge('John', 'Crane'):
            g['John']['Crane']['uses'] = 0  # Set 'uses' to 0 instead of removing the edge
            print("Timestep 1: John stops using the Crane.")
    if t == 2:
        # At timestep 2, Justin stops using the Crane
        if g.has_edge('Justin', 'Crane'):
            g['Justin']['Crane']['uses'] = 0  # Set 'uses' to 0 instead of removing the edge
            print("Timestep 2: Justin stops using the Crane.")
    if t == 3:
        # At timestep 3, Mary stops using the Crane
        if g.has_edge('Mary', 'Crane'):
            g['Mary']['Crane']['uses'] = 0  # Set 'uses' to 0 instead of removing the edge
            print("Timestep 3: Mary stops using the Crane.")
    
    # Display current 'uses' edges for verification
    current_uses = [
        (u, v, d['uses']) for u, v, d in g.edges(data=True)
        if 'uses' in d and d['uses'] == 1
    ]
    print("\nCurrent 'uses' edges:")
    if current_uses:
        for edge in current_uses:
            print(f"  {edge[0]} uses {edge[1]}: {edge[2]}")
    else:
        print("  No active 'uses' edges.")
    
    # Reload the updated graph into PyReason
    pr.load_graph(g)
    
    # Capture the current state
    uses_df, cert_df = capture_state(g, t)
    state_data.append({
        'timestep': t,
        'uses': uses_df,
        'certifications': cert_df
    })

# ================================ GRADIO INTERFACE ====================================
def visualize_timestep(timestep):
    if timestep < 0 or timestep >= timesteps:
        return "<p style='color:red;'>Invalid timestep selected.</p>", "<p style='color:red;'>Invalid timestep selected.</p>"
    state = state_data[timestep]
    uses_df = state['uses']
    cert_df = state['certifications']
    
    # Convert DataFrames to HTML tables for better display
    uses_html = uses_df.to_html(index=False) if not uses_df.empty else "<p>No active 'uses' edges.</p>"
    cert_html = cert_df.to_html(index=False) if not cert_df.empty else "<p>No certifications found.</p>"
    
    return uses_html, cert_html

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 Neuro-Symbolic AI: Certification Propagation Visualization")
    gr.Markdown(
        """
        Welcome to the **Certification Propagation Visualization** tool powered by **Neuro-Symbolic AI** and **PyReason**.
        
        ## **Understanding the Technology**
        
        ### **Neuro-Symbolic AI**
        Neuro-Symbolic AI is an advanced approach that combines the pattern recognition capabilities of neural networks with the logical reasoning strengths of symbolic AI. This fusion enables systems to learn from data and perform complex reasoning tasks, offering both flexibility and interpretability.
        
        ### **PyReason**
        PyReason is a symbolic reasoning engine that facilitates logical inferences on graph-based knowledge representations. It allows the definition of rules and facts to simulate reasoning processes, such as how certifications propagate among workers based on their relationships and equipment usage.
        
        ## **How It Works**
        
        1. **Graph Representation:** Workers and equipment are represented as nodes in a graph, with edges denoting relationships like colleagueship and equipment usage.
        2. **Defining Rules:** Logical rules dictate how certifications spread based on these relationships and equipment usage.
        3. **Reasoning Process:** PyReason applies these rules iteratively to update certification statuses over time.
        4. **Interactive Visualization:** The Gradio interface allows you to observe these dynamics interactively across different timesteps.
        
        ## **Instructions**
        
        - **Select Timestep:** Use the slider to choose a timestep (0 to 3).
        - **Visualize:** Click the "Visualize" button or simply move the slider to update the displays.
        - **Interpret Results:**
            - **Active Equipment Usages:** See which workers are currently using which equipment.
            - **Certification Statuses:** View the certification status (`Yes` or `No`) of each worker.
        """
    )
    
    with gr.Row():
        with gr.Column():
            timestep_slider = gr.Slider(
                minimum=0,
                maximum=timesteps-1,
                step=1,
                label="🔄 Select Timestep",
                value=0
            )
            submit_button = gr.Button("🔍 Visualize")
        with gr.Column():
            uses_output = gr.HTML(label="💼 Active Equipment Usages")
            cert_output = gr.HTML(label="🏅 Certification Statuses")
    
    submit_button.click(
        visualize_timestep,
        inputs=timestep_slider,
        outputs=[uses_output, cert_output]
    )
    
    # Optionally, synchronize slider changes with visualization
    timestep_slider.change(
        visualize_timestep,
        inputs=timestep_slider,
        outputs=[uses_output, cert_output]
    )
    
    gr.Markdown(
        """
        ---
        
        ### **Behind the Scenes**
        
        - **Initialization:** The system starts by defining workers, equipment, and their initial relationships and equipment usage.
        - **Rule Application:** At each timestep, rules defined in PyReason determine how certifications propagate based on current equipment usage and colleague relationships.
        - **State Capture:** The active equipment usages and certification statuses are captured and displayed in tables for each timestep.
        
        ### **Benefits of This Approach**
        
        - **Interpretability:** The logical rules provide clear insights into how decisions (certifications) are made.
        - **Flexibility:** Easily adjust rules or relationships to simulate different scenarios.
        - **Scalability:** Can be extended to larger networks with more complex relationships and rules.
        """
    )

# Launch the Gradio app
demo.launch()