import gradio as io
import networkx as nx
import matplotlib.pyplot as plt
import random

# --- VENUE SETUP ---
class VenueLayout:
    def __init__(self):
        self.graph = nx.DiGraph()
        connections = [
            ('Main_Gate_A', 'Walkway_North', 2, 50),
            ('Main_Gate_B', 'Walkway_South', 2, 20),  # Narrow gate
            ('Walkway_North', 'Food_Court_1', 3, 40),
            ('Walkway_South', 'Food_Court_2', 3, 40),
            ('Food_Court_1', 'Arena_Zone_X', 4, 60),
            ('Food_Court_2', 'Arena_Zone_Y', 4, 60),
            ('Walkway_North', 'Arena_Zone_X', 5, 80),
            ('Walkway_South', 'Arena_Zone_Y', 5, 80),
            ('Walkway_North', 'Walkway_South', 3, 20),
            ('Walkway_South', 'Walkway_North', 3, 20),
            ('Arena_Zone_X', 'Main_Exit_1', 3, 50),
            ('Arena_Zone_Y', 'Main_Exit_2', 3, 50),
        ]
        for src, dest, weight, capacity in connections:
            self.graph.add_edge(src, dest, weight=weight, capacity=capacity, current_load=0)

    def get_dynamic_weight(self, src, dest):
        edge = self.graph[src][dest]
        load = edge['current_load']
        capacity = edge['capacity']
        base_time = edge['weight']
        if load >= capacity:
            return base_time * 10.0
        elif load >= capacity * 0.75:
            return base_time * 3.5
        return base_time

# --- AGENT LOGIC ---
class PersonAgent:
    def __init__(self, agent_id, origin, destination):
        self.agent_id = agent_id
        self.origin = origin
        self.destination = destination
        self.current_node = origin
        self.route = []
        self.route_index = 0
        self.completed = False

    def calculate_route(self, venue, use_rerouting):
        try:
            weight_param = venue.get_dynamic_weight if use_rerouting else 'weight'
            self.route = nx.shortest_path(venue.graph, source=self.current_node, target=self.destination, weight=weight_param)
            self.route_index = 0
        except nx.NetworkNoPath:
            pass

    def step(self):
        if self.current_node == self.destination:
            self.completed = True
            return self.current_node, None
        if self.route_index < len(self.route) - 1:
            from_node = self.route[self.route_index]
            to_node = self.route[self.route_index + 1]
            self.current_node = to_node
            self.route_index += 1
            return from_node, to_node
        self.completed = True
        return self.current_node, None

# --- SIMULATION PIPELINE FOR GRADIO ---
def run_ui_simulation(crowd_size, use_ai_rerouting):
    venue = VenueLayout()
    agents = []
    origins = ['Main_Gate_A', 'Main_Gate_B']
    destinations = ['Arena_Zone_X', 'Arena_Zone_Y', 'Main_Exit_1', 'Main_Exit_2']
    
    random.seed(42)
    for i in range(int(crowd_size)):
        agents.append(PersonAgent(i, random.choice(origins), random.choice(destinations)))

    # Run for 5 timeline steps to accumulate traffic loads
    for _ in range(5):
        for u, v in venue.graph.edges():
            venue.graph[u][v]['current_load'] = 0
        active_agents = [a for a in agents if not a.completed]
        if not active_agents:
            break
        for agent in active_agents:
            agent.calculate_route(venue, use_ai_rerouting)
        for agent in active_agents:
            u, v = agent.step()
            if v is not None:
                venue.graph[u][v]['current_load'] += 1

    # Generate Google-maps style Traffic Report Text
    report = "📋 SYSTEM LIVE REPORT:\n"
    bottlenecks_found = False
    edge_colors = []
    
    for u, v, data in venue.graph.edges(data=True):
        load = data['current_load']
        cap = data['capacity']
        ratio = load / cap if cap > 0 else 0
        
        if ratio >= 1.0:
            report += f"🔴 CRITICAL BOTTLENECK: {u} -> {v} ({load}/{cap} people)\n"
            edge_colors.append('red')
            bottlenecks_found = True
        elif ratio >= 0.75:
            report += f"🟡 WARNING CONGESTION: {u} -> {v} ({load}/{cap} people)\n"
            edge_colors.append('orange')
            bottlenecks_found = True
        else:
            edge_colors.append('green')

    if not bottlenecks_found:
        report += "🟢 All routes operating smoothly. Crowd distributed successfully."

    # Create Visual Map Plot
    fig, ax = plt.subplots(figsize=(10, 6))
    pos = nx.spring_layout(venue.graph, seed=42)
    nx.draw_networkx_nodes(venue.graph, pos, node_size=700, node_color='skyblue', ax=ax)
    nx.draw_networkx_labels(venue.graph, pos, font_size=8, font_weight='bold', ax=ax)
    nx.draw_networkx_edges(venue.graph, pos, edge_color=edge_colors, width=3, arrowsize=15, ax=ax)
    plt.title("Venue Crowd Traffic Density Layout Map")
    plt.axis('off')
    
    return fig, report

# --- GRADIO INTERFACE CONFIGURATION ---
interface = io.Interface(
    fn=run_ui_simulation,
    inputs=[
        io.Slider(minimum=10, maximum=500, value=150, label="Expected Crowd Size"),
        io.Checkbox(value=True, label="Enable Google Maps Style AI Rerouting")
    ],
    outputs=[
        io.Plot(label="Live Congestion Map Layout"),
        io.Textbox(label="Analytics Report Console", lines=6)
    ],
    title="🏢 Real-Time AI Crowd Flow Optimiser",
    description="Simulate venue patterns and clear path networks automatically using dynamic weight calculations."
)

if __name__ == "__main__":
    interface.launch()