import streamlit as st
from graphviz import Digraph
import time
import random

# Define the emoji to use for the swim lanes
SWIM_LANES = {
    "Data Pipelines": "🔍",
    "Build and Train Models": "🧪",
    "Deploy and Predict": "🚀"
}

# Define the graph structure
graph = Digraph()
graph.attr(rankdir="TB") # Top to Bottom or LR Left to Right
graph.attr(fontsize="20")
graph.attr(compound="true")
graph.attr(nodesep="0.5")

# Define the nodes
graph.node("📊 Data Collection")
graph.node("🧹 Data Cleaning")
graph.node("🔧 Data Transformation")
graph.node("🔎 Feature Engineering")
graph.node("⚙️ Model Selection")
graph.node("🎓 Model Training")
graph.node("🚢 Model Deployment")
graph.node("📡 Model Serving")
graph.node("🔮 Predictions")
graph.node("👍 Feedback Collection")
graph.node("🤔 Feedback Processing")
graph.node("✍️ Model Updating")

# Add the edges
graph.edge("📊 Data Collection", "🧹 Data Cleaning")
graph.edge("🧹 Data Cleaning", "🔧 Data Transformation")
graph.edge("🔧 Data Transformation", "🔎 Feature Engineering")
graph.edge("🔎 Feature Engineering", "⚙️ Model Selection")
graph.edge("⚙️ Model Selection", "🎓 Model Training")
graph.edge("🎓 Model Training", "🚢 Model Deployment")
graph.edge("🚢 Model Deployment", "📡 Model Serving")
graph.edge("📡 Model Serving", "🔮 Predictions")
graph.edge("🔮 Predictions", "👍 Feedback Collection")
graph.edge("👍 Feedback Collection", "🤔 Feedback Processing")
graph.edge("🤔 Feedback Processing", "✍️ Model Updating")
graph.edge("✍️ Model Updating", "🎓 Model Training")

# Add the swim lanes
with graph.subgraph(name="cluster_0") as c:
    c.attr(rank="1")
    c.attr(label=SWIM_LANES["Data Pipelines"])
    c.edge("📊 Data Collection", "🧹 Data Cleaning", style="invis")
    c.edge("🧹 Data Cleaning", "🔧 Data Transformation", style="invis")

with graph.subgraph(name="cluster_1") as c:
    c.attr(rank="2")
    c.attr(label=SWIM_LANES["Build and Train Models"])
    c.edge("🔎 Feature Engineering", "⚙️ Model Selection", style="invis")
    c.edge("⚙️ Model Selection", "🎓 Model Training", style="invis")

with graph.subgraph(name="cluster_2") as c:
    c.attr(rank="3")
    c.attr(label=SWIM_LANES["Deploy and Predict"])
    c.edge("🚢 Model Deployment", "📡 Model Serving", style="invis")
    c.edge("📡 Model Serving", "🔮 Predictions", style="invis")

with graph.subgraph(name="cluster_3") as c:
    c.attr(rank="4")
    c.attr(label="Reinforcement Learning Human Feedback")
    c.edge("🔮 Predictions", "👍 Feedback Collection", style="invis")
    c.edge("👍 Feedback Collection", "🤔 Feedback Processing", style="invis")
    c.edge("🤔 Feedback Processing", "✍️ Model Updating", style="invis")

def render_graph():
    st.graphviz_chart(graph.source)

def update_graph():
    for i in range(60):
        # Update the graph with new inputs randomly
        graph.node("📊 Data Collection", label=f"📊 Data Collection\nData {random.randint(0,100)}")
        graph.node("🧹 Data Cleaning", label=f"🧹 Data Cleaning\nCleaned Data {random.randint(0,100)}")
        graph.node("🔧 Data Transformation", label=f"🔧 Data Transformation\nTransformed Data {random.randint(0,100)}")
        graph.node("🔎 Feature Engineering", label=f"🔎 Feature Engineering\nFeatures {random.randint(0,100)}")
        graph.node("⚙️ Model Selection", label=f"⚙️ Model Selection\nSelected Model {random.randint(0,100)}")
        graph.node("🎓 Model Training", label=f"🎓 Model Training\nTrained Model {random.randint(0,100)}")
        graph.node("🚢 Model Deployment", label=f"🚢 Model Deployment\nDeployed Model {random.randint(0,100)}")
        graph.node("📡 Model Serving", label=f"📡 Model Serving\nServed Model {random.randint(0,100)}")
        graph.node("🔮 Predictions", label=f"🔮 Predictions\nPredicted Results {random.randint(0,100)}")
        graph.node("👍 Feedback Collection", label=f"👍 Feedback Collection\nFeedback {random.randint(0,100)}")
        graph.node("🤔 Feedback Processing", label=f"🤔 Feedback Processing\nProcessed Feedback {random.randint(0,100)}")
        graph.node("✍️ Model Updating", label=f"✍️ Model Updating\nUpdated Model {random.randint(0,100)}")
        
        # Render the updated graph
        render_graph()
        
        # Wait for 1 second
        time.sleep(1)
        
# Render the initial graph
render_graph()

# Update the graph every second for 60 seconds
update_graph()