app.py

import streamlit as st
import numpy as np
import random
import matplotlib.pyplot as plt
from scipy.spatial import distance
from sklearn.cluster import KMeans
import networkx as nx
from collections import deque
from scipy.signal import convolve2d

# Constants
GRID_SIZE = 200
FOOD_SOURCES = [(20, 20), (80, 80), (150, 150), (40, 160), (180, 30)]
OBSTACLES = [(50, 50), (100, 100), (150, 50), (70, 130), (120, 80)]
PHEROMONE_DECAY_RATE = 0.02
PHEROMONE_DIFFUSION_RATE = 0.05
MAX_ANTS = 100
MUTATION_RATE = 0.01

# Pheromone Grid
pheromone_grid = np.zeros((GRID_SIZE, GRID_SIZE, 3))  # 3 channels: food, danger, exploration

# Graph representation of the environment
env_graph = nx.grid_2d_graph(GRID_SIZE, GRID_SIZE)

# Remove edges for obstacles
for obstacle in OBSTACLES:
    env_graph.remove_node(obstacle)

# Ant Class
class Ant:
    def __init__(self, position, genome):
        self.position = position
        self.genome = genome
        self.carrying_food = False
        self.energy = 100
        self.memory = deque(maxlen=20)
        self.path_home = []
        self.role = "explorer"
        self.communication_range = 10
        self.q_table = {}  # Changed to dictionary for flexible indexing

    def perceive_environment(self, pheromone_grid, ants):
        self.food_pheromone = pheromone_grid[self.position[0], self.position[1], 0]
        self.danger_pheromone = pheromone_grid[self.position[0], self.position[1], 1]
        self.exploration_pheromone = pheromone_grid[self.position[0], self.position[1], 2]
        
        # Perceive nearby ants
        self.nearby_ants = [ant for ant in ants if distance.euclidean(self.position, ant.position) <= self.communication_range]

     def act(self, pheromone_grid):
        possible_actions = self.get_possible_actions()
        
        if random.random() < self.genome['exploration_rate']:  
            action = random.choice(possible_actions)
        else:
            q_values = [self.get_q_value(action) for action in possible_actions]
            action = possible_actions[np.argmax(q_values)]

        reward = self.calculate_reward()
        self.update_q_table(action, reward)

        return action

    def get_q_value(self, action):
        return self.q_table.get((self.position, action), 0)

    def update_q_table(self, action, reward):
        current_q = self.get_q_value(action)
        max_future_q = max([self.get_q_value(future_action) for future_action in self.get_possible_actions()])
        
        new_q = (1 - self.genome['learning_rate']) * current_q + \
                self.genome['learning_rate'] * (reward + self.genome['discount_factor'] * max_future_q)
        
        self.q_table[(self.position, action)] = new_q

    def get_possible_actions(self):
        x, y = self.position
        possible_actions = []
        for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:  # right, down, left, up
            new_x, new_y = x + dx, y + dy
            if 0 <= new_x < GRID_SIZE and 0 <= new_y < GRID_SIZE and (new_x, new_y) not in OBSTACLES:
                possible_actions.append((new_x, new_y))
        return possible_actions

    def update(self, pheromone_grid, ants):
        self.perceive_environment(pheromone_grid, ants)
        action = self.act(pheromone_grid)
        self.position = action

        self.energy -= 1
        if self.energy <= 0:
            return False  # Ant dies

        if self.carrying_food:
            pheromone_grid[self.position[0], self.position[1], 0] += 5
            if self.position == (0, 0):  # Drop food at nest
                self.carrying_food = False
                self.energy = min(100, self.energy + 50)
                return True  # Food collected successfully
        
        if self.position in FOOD_SOURCES and not self.carrying_food:
            self.carrying_food = True
            pheromone_grid[self.position[0], self.position[1], 0] += 10
        
        if self.position in OBSTACLES:
            pheromone_grid[self.position[0], self.position[1], 1] += 5
        
        pheromone_grid[self.position[0], self.position[1], 2] += 1  # Exploration pheromone

        self.memory.append(self.position)
        
        # Update role based on situation
        if self.carrying_food:
            self.role = "carrier"
        elif self.food_pheromone > 5:
            self.role = "follower"
        else:
            self.role = "explorer"

        # Path planning
        if self.carrying_food and not self.path_home:
            self.path_home = nx.shortest_path(env_graph, self.position, (0, 0))

        return True  # Ant survives

# Pheromone Diffusion using Convolution
def diffuse_pheromones(pheromone_grid):
    kernel = np.array([[0.05, 0.1, 0.05],
                       [0.1,  0.4, 0.1],
                       [0.05, 0.1, 0.05]])
    for i in range(3):  # For each pheromone type
        pheromone_grid[:,:,i] = convolve2d(pheromone_grid[:,:,i], kernel, mode='same', boundary='wrap')

# Genetic Algorithm
def crossover(parent1, parent2):
    child = {}
    for key in parent1.keys():
        if random.random() < 0.5:
            child[key] = parent1[key]
        else:
            child[key] = parent2[key]
    return child

def mutate(genome):
    for key in genome.keys():
        if random.random() < MUTATION_RATE:
            genome[key] += random.uniform(-0.1, 0.1)
            genome[key] = max(0, min(1, genome[key]))
    return genome

# Simulation Loop
def simulate(ants):
    global pheromone_grid
    food_collected = 0
    for ant in ants:
        if ant.update(pheromone_grid, ants):
            if ant.position == (0, 0) and not ant.carrying_food:
                food_collected += 1

    pheromone_grid *= (1 - PHEROMONE_DECAY_RATE)
    diffuse_pheromones(pheromone_grid)

    # Genetic Algorithm
    if len(ants) > MAX_ANTS:
        ants.sort(key=lambda x: x.energy, reverse=True)
        survivors = ants[:MAX_ANTS//2]
        new_ants = []
        while len(new_ants) < MAX_ANTS//2:
            parent1, parent2 = random.sample(survivors, 2)
            child_genome = crossover(parent1.genome, parent2.genome)
            child_genome = mutate(child_genome)
            new_ant = Ant((random.randint(0, GRID_SIZE-1), random.randint(0, GRID_SIZE-1)), child_genome)
            new_ants.append(new_ant)
        ants = survivors + new_ants
    
    return ants, food_collected

# Clustering for strategic analysis
def analyze_ant_clusters(ants):
    positions = np.array([ant.position for ant in ants])
    kmeans = KMeans(n_clusters=3)
    kmeans.fit(positions)
    return kmeans.cluster_centers_

# Visualization Functions
def plot_environment(pheromone_grid, ants, cluster_centers):
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(np.sum(pheromone_grid, axis=2), cmap='viridis', alpha=0.7)
    
    for ant in ants:
        color = 'blue' if ant.role == 'explorer' else 'red' if ant.role == 'carrier' else 'green'
        ax.plot(ant.position[1], ant.position[0], 'o', color=color, markersize=4)
    
    for food_x, food_y in FOOD_SOURCES:
        ax.plot(food_y, food_x, 'go', markersize=10)
    
    for obstacle_x, obstacle_y in OBSTACLES:
        ax.plot(obstacle_y, obstacle_x, 'ro', markersize=10)
    
    for center in cluster_centers:
        ax.plot(center[1], center[0], 'mo', markersize=15, alpha=0.7)
    
    ax.set_xlim([0, GRID_SIZE])
    ax.set_ylim([GRID_SIZE, 0])
    return fig

# Streamlit App
st.title("Advanced Ant Hivemind Simulation")

# Sidebar controls
st.sidebar.header("Simulation Parameters")
num_ants = st.sidebar.slider("Number of Ants", 10, MAX_ANTS, 50)
exploration_rate = st.sidebar.slider("Exploration Rate", 0.0, 1.0, 0.2)
learning_rate = st.sidebar.slider("Learning Rate", 0.0, 1.0, 0.1)
discount_factor = st.sidebar.slider("Discount Factor", 0.0, 1.0, 0.9)

# Initialize ants
ants = [Ant((random.randint(0, GRID_SIZE-1), random.randint(0, GRID_SIZE-1)), 
            {'exploration_rate': exploration_rate, 
             'learning_rate': learning_rate, 
             'discount_factor': discount_factor}) 
        for _ in range(num_ants)]

# Simulation control
start_simulation = st.sidebar.button("Start Simulation")
stop_simulation = st.sidebar.button("Stop Simulation")
reset_simulation = st.sidebar.button("Reset Simulation")

# Initialize variables
total_food_collected = 0
iterations = 0

# Main simulation loop
if start_simulation:
    cluster_centers = np.array([[0, 0], [0, 0], [0, 0]])
    
    progress_bar = st.progress(0)
    stats_placeholder = st.empty()
    plot_placeholder = st.empty()
    
    while not stop_simulation:
        ants, food_collected = simulate(ants)
        total_food_collected += food_collected
        iterations += 1
        
        if iterations % 10 == 0:
            cluster_centers = analyze_ant_clusters(ants)
        
        if iterations % 5 == 0:
            progress_bar.progress(min(iterations / 1000, 1.0))
            stats_placeholder.write(f"Iterations: {iterations}, Total Food Collected: {total_food_collected}")
            fig = plot_environment(pheromone_grid, ants, cluster_centers)
            plot_placeholder.pyplot(fig)
            plt.close(fig)

if reset_simulation:
    pheromone_grid = np.zeros((GRID_SIZE, GRID_SIZE, 3))
    ants = [Ant((random.randint(0, GRID_SIZE-1), random.randint(0, GRID_SIZE-1)), 
                {'exploration_rate': exploration_rate, 
                 'learning_rate': learning_rate, 
                 'discount_factor': discount_factor}) 
            for _ in range(num_ants)]
    total_food_collected = 0
    iterations = 0

# Display final statistics only if simulation has run
if iterations > 0:
    st.write("## Final Statistics")
    st.write(f"Total Food Collected: {total_food_collected}")
    st.write(f"Average Food per Iteration: {total_food_collected / iterations}")

    # Display heatmap of pheromone concentration
    st.write("## Pheromone Concentration Heatmap")
    fig, ax = plt.subplots(figsize=(10, 10))
    heatmap = ax.imshow(np.sum(pheromone_grid, axis=2), cmap='hot', interpolation='nearest')
    plt.colorbar(heatmap)
    st.pyplot(fig)

    # Display ant role distribution
    roles = [ant.role for ant in ants]
    role_counts = {role: roles.count(role) for role in set(roles)}
    st.write("## Ant Role Distribution")
    st.bar_chart(role_counts)

    # Display network graph of ant communication
    st.write("## Ant Communication Network")
    G = nx.Graph()
    for ant in ants:
        G.add_node(ant.position)
        for nearby_ant in ant.nearby_ants:
            G.add_edge(ant.position, nearby_ant.position)

    fig, ax = plt.subplots(figsize=(10, 10))
    pos = nx.spring_layout(G)
    nx.draw(G, pos, with_labels=False, node_size=30, node_color='skyblue', edge_color='gray', ax=ax)
    st.pyplot(fig)