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

# 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 = np.zeros((GRID_SIZE, GRID_SIZE, 4))

    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:
            action = np.argmax(self.q_table[self.position[0], self.position[1], possible_actions])

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

        return action

    def calculate_reward(self):
        if self.carrying_food:
            return 10
        elif self.position in FOOD_SOURCES:
            return 20
        elif self.position in OBSTACLES:
            return -10
        else:
            return -1 + self.food_pheromone - self.danger_pheromone + 0.5 * self.exploration_pheromone

    def update_q_table(self, action, reward):
        self.q_table[self.position[0], self.position[1], action] = (
            (1 - self.genome['learning_rate']) * self.q_table[self.position[0], self.position[1], action] +
            self.genome['learning_rate'] * (reward + self.genome['discount_factor'] * np.max(self.q_table[self.position[0], self.position[1]]))
        )

    def get_possible_actions(self):
        return list(env_graph.neighbors(self.position))

    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] = np.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")

# Main simulation loop
if start_simulation:
    total_food_collected = 0
    iterations = 0
    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)]

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

# 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)