import random
import numpy as np
import streamlit as st
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import time

class Organelle:
    def __init__(self, type):
        self.type = type  # e.g., "nucleus", "mitochondria", "chloroplast"

class Cell:
    def __init__(self, x, y, cell_type="prokaryote"):
        self.x = x
        self.y = y
        self.energy = 100
        self.cell_type = cell_type
        self.organelles = []
        self.size = 1
        self.color = "blue"
        self.division_threshold = 150
        
        if cell_type == "prokaryote":
            self.color = "lightblue"
        elif cell_type == "early_eukaryote":
            self.organelles.append(Organelle("nucleus"))
            self.color = "green"
            self.size = 2
        elif cell_type == "advanced_eukaryote":
            self.organelles.extend([Organelle("nucleus"), Organelle("mitochondria")])
            self.color = "red"
            self.size = 3
        elif cell_type == "plant_like":
            self.organelles.extend([Organelle("nucleus"), Organelle("mitochondria"), Organelle("chloroplast")])
            self.color = "darkgreen"
            self.size = 4

    def move(self, environment):
        dx = random.uniform(-1, 1)
        dy = random.uniform(-1, 1)
        self.x = max(0, min(environment.width - 1, self.x + dx))
        self.y = max(0, min(environment.height - 1, self.y + dy))
        self.energy -= 0.5 * self.size

    def feed(self, environment):
        if "chloroplast" in [org.type for org in self.organelles]:
            # Photosynthesis
            self.energy += environment.light_level * 2
        else:
            # Consume environmental nutrients
            self.energy += environment.grid[int(self.y)][int(self.x)] * 0.1
            environment.grid[int(self.y)][int(self.x)] *= 0.9

    def can_divide(self):
        return self.energy > self.division_threshold

    def divide(self):
        if self.can_divide():
            self.energy /= 2
            return Cell(self.x, self.y, self.cell_type)
        return None

    def can_fuse(self, other):
        return (self.cell_type == "prokaryote" and other.cell_type == "prokaryote" and
                random.random() < 0.001)  # 0.1% chance of fusion

    def fuse(self, other):
        new_cell = Cell(
            (self.x + other.x) / 2,
            (self.y + other.y) / 2,
            "early_eukaryote"
        )
        new_cell.energy = self.energy + other.energy
        return new_cell

class Environment:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.grid = np.random.rand(height, width) * 10  # Nutrient distribution
        self.light_level = 5  # Ambient light level
        self.cells = []
        self.time = 0
        self.population_history = {
            "prokaryote": [],
            "early_eukaryote": [],
            "advanced_eukaryote": [],
            "plant_like": []
        }

    def add_cell(self, cell):
        self.cells.append(cell)

    def update(self):
        self.time += 1
        
        # Update environment
        self.grid += np.random.rand(self.height, self.width) * 0.1
        self.light_level = 5 + np.sin(self.time / 100) * 2  # Fluctuating light levels

        new_cells = []
        cells_to_remove = []

        for cell in self.cells:
            cell.move(self)
            cell.feed(self)

            if cell.energy <= 0:
                cells_to_remove.append(cell)
            elif cell.can_divide():
                new_cell = cell.divide()
                if new_cell:
                    new_cells.append(new_cell)

        # Handle cell fusion
        for i, cell1 in enumerate(self.cells):
            for cell2 in self.cells[i+1:]:
                if cell1.can_fuse(cell2):
                    new_cell = cell1.fuse(cell2)
                    new_cells.append(new_cell)
                    cells_to_remove.extend([cell1, cell2])

        # Add new cells and remove dead/fused cells
        self.cells.extend(new_cells)
        self.cells = [cell for cell in self.cells if cell not in cells_to_remove]

        # Introduce mutations
        for cell in self.cells:
            if random.random() < 0.0001:  # 0.01% chance of mutation
                if cell.cell_type == "early_eukaryote":
                    cell.cell_type = "advanced_eukaryote"
                    cell.organelles.append(Organelle("mitochondria"))
                    cell.color = "red"
                    cell.size = 3
                elif cell.cell_type == "advanced_eukaryote" and random.random() < 0.5:
                    cell.cell_type = "plant_like"
                    cell.organelles.append(Organelle("chloroplast"))
                    cell.color = "darkgreen"
                    cell.size = 4

        # Record population counts
        for cell_type in self.population_history.keys():
            count = len([cell for cell in self.cells if cell.cell_type == cell_type])
            self.population_history[cell_type].append(count)

    def get_visualization_data(self):
        cell_data = {
            "prokaryote": {"x": [], "y": [], "size": [], "color": "lightblue"},
            "early_eukaryote": {"x": [], "y": [], "size": [], "color": "green"},
            "advanced_eukaryote": {"x": [], "y": [], "size": [], "color": "red"},
            "plant_like": {"x": [], "y": [], "size": [], "color": "darkgreen"}
        }

        for cell in self.cells:
            cell_data[cell.cell_type]["x"].append(cell.x)
            cell_data[cell.cell_type]["y"].append(cell.y)
            cell_data[cell.cell_type]["size"].append(cell.size * 3)

        return cell_data, self.population_history

def setup_figure(env):
    fig = make_subplots(rows=1, cols=2, subplot_titles=("Cell Distribution", "Population Over Time"))

    # Cell distribution
    for cell_type, data in env.get_visualization_data()[0].items():
        fig.add_trace(go.Scatter(
            x=data["x"], y=data["y"], mode='markers',
            marker=dict(color=data["color"], size=data["size"]),
            name=cell_type
        ), row=1, col=1)

    # Population over time
    for cell_type, counts in env.population_history.items():
        fig.add_trace(go.Scatter(y=counts, mode='lines', name=cell_type), row=1, col=2)

    fig.update_xaxes(title_text="X", row=1, col=1)
    fig.update_yaxes(title_text="Y", row=1, col=1)
    fig.update_xaxes(title_text="Time", row=1, col=2)
    fig.update_yaxes(title_text="Population", row=1, col=2)

    fig.update_layout(height=600, width=1200, title_text="Cell Evolution Simulation")

    return fig

# Streamlit app
st.title("Cell Evolution Simulation")

num_steps = st.slider("Number of simulation steps", 100, 1000, 500)
initial_cells = st.slider("Initial number of cells", 10, 100, 50)
update_interval = st.slider("Update interval (milliseconds)", 100, 1000, 200)

if st.button("Run Simulation"):
    env = Environment(100, 100)
    
    # Add initial cells
    for _ in range(initial_cells):
        cell = Cell(random.uniform(0, env.width), random.uniform(0, env.height))
        env.add_cell(cell)
    
    # Set up the figure
    fig = setup_figure(env)
    chart = st.plotly_chart(fig, use_container_width=True)
    
    # Run simulation
    for step in range(num_steps):
        env.update()
        
        # Update the figure data
        with fig.batch_update():
            cell_data, population_history = env.get_visualization_data()
            for i, (cell_type, data) in enumerate(cell_data.items()):
                fig.data[i].x = data["x"]
                fig.data[i].y = data["y"]
                fig.data[i].marker.size = data["size"]
            
            for i, (cell_type, counts) in enumerate(population_history.items()):
                fig.data[i+4].y = counts  # +4 because we have 4 cell types in the first subplot
        
        fig.layout.title.text = f"Cell Evolution Simulation (Time: {env.time})"
        
        # Update the chart
        chart.plotly_chart(fig, use_container_width=True)
        
        time.sleep(update_interval / 1000)  # Convert milliseconds to seconds

    st.write("Simulation complete!")