app.py

# app.py
import streamlit as st
import numpy as np
import random
import copy
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go # Using graph_objects for more control over the grid
# import matplotlib.pyplot as plt # No longer needed for grid
# import matplotlib.colors # No longer needed for grid
import time # For the delay in continuous run

# --- Simulation Core Classes ---

class Cell:
    """Base class for all cells."""
    def __init__(self, x, y):
        self.x = x
        self.y = y # Represents ROW index in grid (origin top-left)

class CancerCell(Cell):
    """Represents a cancer cell."""
    CELL_TYPE = 1 # Grid representation
    LABEL = 'Cancer'
    COLOR = 'red'
    def __init__(self, x, y, growth_prob, metastasis_prob, resistance, mutation_rate):
        super().__init__(x, y)
        self.growth_prob = growth_prob
        self.metastasis_prob = metastasis_prob
        self.resistance = resistance # 0.0 (susceptible) to 1.0 (fully resistant)
        self.mutation_rate = mutation_rate
        self.is_alive = True

    def attempt_division(self, sim):
        """Attempt to divide into an adjacent empty cell."""
        if random.random() < self.growth_prob:
            neighbors = sim.get_neighbors(self.x, self.y)
            empty_neighbors = [(nx, ny) for nx, ny, cell_type in neighbors if cell_type == 0]
            if empty_neighbors:
                nx, ny = random.choice(empty_neighbors)
                # Create a new cell with possibly mutated properties
                new_cell = copy.deepcopy(self)
                new_cell.x, new_cell.y = nx, ny
                new_cell.mutate() # Mutate the offspring
                return new_cell
        return None

    def attempt_metastasis(self, sim):
        """Attempt to move to a random empty cell on the grid."""
        if random.random() < self.metastasis_prob:
            empty_cells = np.argwhere(sim.grid == 0)
            if len(empty_cells) > 0:
                new_y, new_x = random.choice(empty_cells) # Note: numpy argwhere returns (row, col) -> (y, x)
                # Return the new position for the simulation to handle the move
                return (int(new_x), int(new_y)) # Convert numpy types to int
        return None

    def mutate(self):
        """Potentially mutate resistance and growth probability."""
        if random.random() < self.mutation_rate:
            # Mutate resistance slightly
            self.resistance += random.uniform(-0.1, 0.1)
            self.resistance = max(0.0, min(1.0, self.resistance)) # Keep within [0, 1]
        if random.random() < self.mutation_rate:
             # Mutate growth prob slightly
            self.growth_prob += random.uniform(-0.05, 0.05)
            self.growth_prob = max(0.0, min(1.0, self.growth_prob)) # Keep within [0, 1]

    def check_drug_effect(self, drug_effect_base, drug_resistance_interaction):
        """Check if the drug kills this cell."""
        effective_drug = drug_effect_base * max(0, (1.0 - self.resistance * drug_resistance_interaction))
        if random.random() < effective_drug:
            self.is_alive = False
            return True # Killed by drug
        return False


class ImmuneCell(Cell):
    """Represents an immune cell."""
    CELL_TYPE = 2 # Grid representation
    LABEL = 'Immune'
    COLOR = 'blue'
    def __init__(self, x, y, base_kill_prob, movement_prob, lifespan, activation_boost):
        super().__init__(x, y)
        self.base_kill_prob = base_kill_prob
        self.movement_prob = movement_prob
        self.lifespan = lifespan
        self.activation_boost = activation_boost # Added boost when activated by drug
        self.steps_alive = 0
        self.is_activated = False # Can be temporarily boosted by drug
        self.is_alive = True

    def attempt_move(self, sim):
        """Attempt to move to a random adjacent cell (can be empty or occupied)."""
        if random.random() < self.movement_prob:
            neighbors = sim.get_neighbors(self.x, self.y, radius=1, include_self=False)
            potential_moves = [(nx, ny) for nx, ny, _ in neighbors]
            if potential_moves:
                 # Prioritize moving towards cancer cells slightly
                cancer_neighbors = [(nx, ny) for nx, ny, cell_type in neighbors if cell_type == CancerCell.CELL_TYPE]
                if cancer_neighbors and random.random() < 0.5: # 50% chance to prioritize cancer
                     nx, ny = random.choice(cancer_neighbors)
                else:
                    nx, ny = random.choice(potential_moves)
                 # Return the new position for the simulation to handle the move
                return (nx, ny)
        return None

    def attempt_kill(self, sim):
        """Attempt to kill adjacent cancer cells."""
        killed_coords = []
        neighbors = sim.get_neighbors(self.x, self.y)
        cancer_neighbors = [(nx, ny) for nx, ny, cell_type in neighbors if cell_type == CancerCell.CELL_TYPE]

        current_kill_prob = self.base_kill_prob
        if self.is_activated:
            current_kill_prob += self.activation_boost
            current_kill_prob = min(1.0, current_kill_prob) # Cap at 1.0

        for nx, ny in cancer_neighbors:
            if random.random() < current_kill_prob:
                target_cell = sim.get_cell_at(nx, ny)
                if target_cell and isinstance(target_cell, CancerCell) and target_cell.is_alive:
                    target_cell.is_alive = False
                    killed_coords.append((nx, ny))
        return killed_coords # Return coords of killed cells

    def step(self):
        """Increment age and check lifespan."""
        self.steps_alive += 1
        if self.steps_alive >= self.lifespan:
            self.is_alive = False
        self.is_activated = False # Reset activation each step unless reactivated

    def activate_by_drug(self, drug_immune_boost_prob):
         """Potentially activate based on drug presence."""
         if random.random() < drug_immune_boost_prob:
             self.is_activated = True


class Simulation:
    """Manages the simulation grid and cells."""
    def __init__(self, params):
        self.params = params
        self.grid_size = params['grid_size']
        self.grid = np.zeros((self.grid_size, self.grid_size), dtype=int)
        self.cells = {} # Using dict {(x, y): cell_obj} for faster lookup
        self.history = [] # To store cell counts over time
        self.current_step = 0
        self._initialize_cells()
        self._record_history() # Record initial state

    def _initialize_cells(self):
        """Place initial cells on the grid."""
        center_x, center_y = self.grid_size // 2, self.grid_size // 2

        # Initial Cancer Cells (cluster in the center)
        radius = max(1, int(np.sqrt(self.params['initial_cancer_cells']) / 2))
        count = 0
        placed_coords = set() # Keep track of where we placed cells initially
        for r in range(radius + 2): # Search slightly larger radius if needed
             for x in range(center_x - r, center_x + r + 1):
                 for y in range(center_y - r, center_y + r + 1):
                     if count >= self.params['initial_cancer_cells']: break
                     if 0 <= x < self.grid_size and 0 <= y < self.grid_size:
                         coords = (x,y)
                         if self.grid[y, x] == 0 and coords not in placed_coords: # Ensure cell is placed in empty spot
                            cell = CancerCell(x, y,
                                              self.params['cancer_growth_prob'],
                                              self.params['cancer_metastasis_prob'],
                                              self.params['cancer_initial_resistance'],
                                              self.params['cancer_mutation_rate'])
                            self.grid[y, x] = CancerCell.CELL_TYPE
                            self.cells[coords] = cell
                            placed_coords.add(coords)
                            count += 1
                 if count >= self.params['initial_cancer_cells']: break
             if count >= self.params['initial_cancer_cells']: break


        # Initial Immune Cells (randomly distributed)
        immune_count = 0
        attempts = 0 # Prevent infinite loop if grid is too full
        max_attempts = self.grid_size * self.grid_size * 2
        while immune_count < self.params['initial_immune_cells'] and attempts < max_attempts:
            x, y = random.randint(0, self.grid_size - 1), random.randint(0, self.grid_size - 1)
            coords = (x,y)
            if self.grid[y, x] == 0 and coords not in placed_coords: # Place only in empty spots
                cell = ImmuneCell(x, y,
                                  self.params['immune_base_kill_prob'],
                                  self.params['immune_movement_prob'],
                                  self.params['immune_lifespan'],
                                  self.params['drug_immune_activation_boost'])
                self.grid[y, x] = ImmuneCell.CELL_TYPE
                self.cells[coords] = cell
                placed_coords.add(coords)
                immune_count += 1
            attempts += 1
        if attempts >= max_attempts and immune_count < self.params['initial_immune_cells']:
             st.warning(f"Could only place {immune_count}/{self.params['initial_immune_cells']} immune cells due to space constraints.")


    def get_neighbors(self, x, y, radius=1, include_self=False):
        """Get neighbors within a radius, handling grid boundaries."""
        neighbors = []
        for dx in range(-radius, radius + 1):
            for dy in range(-radius, radius + 1):
                if not include_self and dx == 0 and dy == 0:
                    continue
                nx, ny = x + dx, y + dy
                # Check boundaries (no wrap-around)
                if 0 <= nx < self.grid_size and 0 <= ny < self.grid_size:
                    neighbors.append((nx, ny, self.grid[ny, nx]))
        return neighbors

    def get_cell_at(self, x, y):
        """Retrieve cell object at given coordinates."""
        return self.cells.get((x, y), None)

    def _apply_drug_effects(self):
        """Apply drug effects: killing cancer cells and activating immune cells."""
        killed_by_drug = []
        immune_cells_to_activate = []

        # Iterate through a copy of keys because dict size might change
        for coords, cell in list(self.cells.items()):
            if isinstance(cell, CancerCell) and cell.is_alive:
                if cell.check_drug_effect(self.params['drug_effect_base'], self.params['drug_resistance_interaction']):
                     killed_by_drug.append(coords)
                     # Check for nearby immune cells to activate
                     neighbors = self.get_neighbors(cell.x, cell.y, radius=self.params['drug_immune_activation_radius'])
                     for nx, ny, cell_type in neighbors:
                         if cell_type == ImmuneCell.CELL_TYPE:
                             immune_cell = self.get_cell_at(nx, ny)
                             if immune_cell and immune_cell.is_alive:
                                 immune_cells_to_activate.append(immune_cell)

        # Apply activation (using a set to avoid duplicate activations if multiple cancer cells are nearby)
        for immune_cell in set(immune_cells_to_activate):
             immune_cell.activate_by_drug(self.params['drug_immune_boost_prob']) # Pass prob here

        return killed_by_drug

    def _immune_cell_actions(self):
        """Handle immune cell movement and killing actions."""
        immune_moves = {} # {old_coords: new_coords}
        killed_by_immune = []

        # Iterate through a copy of keys
        immune_cells_list = [cell for cell in self.cells.values() if isinstance(cell, ImmuneCell) and cell.is_alive]
        random.shuffle(immune_cells_list) # Randomize order of action

        for cell in immune_cells_list:
             if not cell.is_alive: continue

             # 1. Aging
             cell.step()
             if not cell.is_alive: continue # Died of old age

             # 2. Attempt Kill
             killed_coords = cell.attempt_kill(self)
             killed_by_immune.extend(killed_coords)

             # 3. Attempt Move (only if it didn't die)
             new_pos = cell.attempt_move(self)
             if new_pos:
                nx, ny = new_pos
                # Check if target is empty OR occupied by a cancer cell immune cell can kill/displace
                # Prevent moving onto another immune cell's intended spot in this step
                # Simplification: allow overlap for now, let update handle conflict
                if new_pos not in immune_moves.values(): # Avoid multiple immune cells targetting same spot
                   immune_moves[(cell.x, cell.y)] = new_pos


        return immune_moves, killed_by_immune

    def _cancer_cell_actions(self):
        """Handle cancer cell division, metastasis, and mutation."""
        new_cancer_cells = []
        cancer_moves = {} # {old_coords: new_coords} for metastasis

        # Iterate through a copy of keys
        cancer_cells_list = [cell for cell in self.cells.values() if isinstance(cell, CancerCell) and cell.is_alive]
        random.shuffle(cancer_cells_list) # Randomize order

        for cell in cancer_cells_list:
            if not cell.is_alive: continue # Could have been killed earlier in the step

            # 1. Mutation (apply first, affects division/metastasis below)
            # Moved mutation inside attempt_division/metastasis for offspring only in prev ver, let's keep it there.
            # Parent mutation should also occur
            cell.mutate() # Parent mutates regardless of division/metastasis

            # 2. Attempt Division
            offspring = cell.attempt_division(self) # Offspring inherits parent's (possibly mutated) state and then mutates itself
            if offspring:
                 # Check if the target spot is still empty (could have been taken by metastasis/immune move)
                 # This check will happen more definitively in _update_grid_and_cells
                 new_cancer_cells.append(offspring)

            # 3. Attempt Metastasis (only if division didn't occur?) Let's allow both for now.
            new_pos = cell.attempt_metastasis(self)
            if new_pos:
                # Check if the target spot is still empty AND not targeted by another metastasis
                # Defer final check to _update_grid_and_cells
                if new_pos not in cancer_moves.values(): # Avoid multiple metastases targeting same spot
                   cancer_moves[(cell.x, cell.y)] = new_pos

        return new_cancer_cells, cancer_moves


    def _update_grid_and_cells(self, killed_coords_list, immune_moves, cancer_moves, new_cancer_cells):
        """Update the grid and cell dictionary based on actions."""

        # 1. Process deaths first
        all_killed_coords = set()
        for coords_list in killed_coords_list:
            all_killed_coords.update(coords_list)

        # Also add immune cells that died of old age
        for coords, cell in list(self.cells.items()):
             if isinstance(cell, ImmuneCell) and not cell.is_alive:
                 all_killed_coords.add(coords)

        for x, y in all_killed_coords:
            if (x, y) in self.cells:
                self.grid[y, x] = 0
                if (x, y) in self.cells: # Check again, might be double-killed?
                    del self.cells[(x, y)]

        # --- Resolve Move Conflicts & Update ---
        # Priority: Immune > Cancer Metastasis. If conflict, immune wins, cancer move fails.
        # If immune A wants to move to B, and immune C wants to move to B, one fails randomly.
        # If cancer A wants to move to B, and cancer C wants to move to B, one fails randomly.
        # If immune A wants to move to B, and cancer C wants to move to B, immune wins.

        occupied_targets = set()
        resolved_immune_moves = {} # {old_coords: new_coords}
        resolved_cancer_moves = {} # {old_coords: new_coords}

        # Shuffle move order for fairness in conflicts
        immune_move_items = list(immune_moves.items())
        random.shuffle(immune_move_items)
        cancer_move_items = list(cancer_moves.items())
        random.shuffle(cancer_move_items)

        # Process immune moves first
        for old_coords, new_coords in immune_move_items:
            if old_coords not in self.cells: continue # Cell died before moving
            if new_coords not in occupied_targets:
                target_cell = self.get_cell_at(new_coords[0], new_coords[1])
                # Allow move if target is empty, or is a cancer cell (implicit displacement/kill)
                if target_cell is None or isinstance(target_cell, CancerCell):
                     resolved_immune_moves[old_coords] = new_coords
                     occupied_targets.add(new_coords)
                # else: blocked by another immune cell already there or moving there

        # Process cancer metastasis moves
        for old_coords, new_coords in cancer_move_items:
            if old_coords not in self.cells: continue # Cell died before moving
            if new_coords not in occupied_targets:
                 target_cell = self.get_cell_at(new_coords[0], new_coords[1])
                 # Allow move ONLY if target is empty
                 if target_cell is None:
                     resolved_cancer_moves[old_coords] = new_coords
                     occupied_targets.add(new_coords)
                 # else: blocked by existing cell or an immune cell moving there


        # Apply moves: remove old, add new
        moved_cells_buffer = {} # Store {new_coords: cell} before adding back

        for old_coords, new_coords in resolved_immune_moves.items():
            if old_coords in self.cells: # Check if cell still exists
                cell = self.cells.pop(old_coords)
                self.grid[old_coords[1], old_coords[0]] = 0
                cell.x, cell.y = new_coords
                moved_cells_buffer[new_coords] = cell

        for old_coords, new_coords in resolved_cancer_moves.items():
             if old_coords in self.cells: # Check if cell still exists
                cell = self.cells.pop(old_coords)
                self.grid[old_coords[1], old_coords[0]] = 0
                cell.x, cell.y = new_coords
                moved_cells_buffer[new_coords] = cell

        # Add moved cells back, handling displacement
        for new_coords, cell in moved_cells_buffer.items():
            # If an immune cell lands on a cancer cell's spot, the cancer cell should be gone.
            if isinstance(cell, ImmuneCell) and new_coords in self.cells and isinstance(self.cells[new_coords], CancerCell):
                 del self.cells[new_coords] # Remove the displaced cancer cell

            # Place the moved cell
            self.cells[new_coords] = cell
            self.grid[new_coords[1], new_coords[0]] = cell.CELL_TYPE


        # 3. Process births (add new cancer cells)
        # Shuffle order for fairness if multiple births target same location (unlikely but possible)
        random.shuffle(new_cancer_cells)
        added_cells_count = 0
        for cell in new_cancer_cells:
            coords = (cell.x, cell.y)
            # Final check if the spot is truly empty *after* deaths and moves
            if self.grid[cell.y, cell.x] == 0 and coords not in self.cells:
                self.grid[cell.y, cell.x] = CancerCell.CELL_TYPE
                self.cells[coords] = cell
                added_cells_count += 1
            # else: Birth failed due to space conflict

    def step(self):
        """Perform one step of the simulation."""
        # Check if simulation should stop before proceeding
        cancer_count = sum(1 for cell in self.cells.values() if isinstance(cell, CancerCell))
        if cancer_count == 0 and self.current_step > 0: # Check after at least one step
             st.session_state.final_message = "Cancer eliminated!"
             return False # Stop simulation
        if self.current_step >= self.params['max_steps']:
             st.session_state.final_message = "Maximum steps reached."
             return False # Stop simulation
        if not self.cells: # Stop if absolutely no cells left for some reason
             st.session_state.final_message = "No cells remaining."
             return False

        # --- Action Phase ---
        # 1. Drug effects (kill cancer, activate immune)
        killed_by_drug = self._apply_drug_effects()

        # 2. Immune cell actions (move, kill, age)
        immune_moves, killed_by_immune = self._immune_cell_actions()

        # 3. Cancer cell actions (divide, metastasize) - Mutation happens within these
        new_cancer_cells, cancer_moves = self._cancer_cell_actions()

        # --- Update Phase ---
        # Consolidate killed cells list
        all_killed_this_step = [killed_by_drug, killed_by_immune]

        # Apply all changes to grid and cell list
        self._update_grid_and_cells(all_killed_this_step, immune_moves, cancer_moves, new_cancer_cells)

        # --- End Step ---
        self.current_step += 1
        self._record_history()


        return True # Continue simulation


    def _record_history(self):
        """Record the number of each cell type."""
        cancer_count = sum(1 for cell in self.cells.values() if isinstance(cell, CancerCell))
        immune_count = sum(1 for cell in self.cells.values() if isinstance(cell, ImmuneCell))
        avg_resistance = 0
        if cancer_count > 0:
            avg_resistance = sum(cell.resistance for cell in self.cells.values() if isinstance(cell, CancerCell)) / cancer_count

        self.history.append({
            'Step': self.current_step,
            'Cancer Cells': cancer_count,
            'Immune Cells': immune_count,
            'Average Resistance': avg_resistance
        })

    def get_history_df(self):
        """Return the recorded history as a Pandas DataFrame."""
        return pd.DataFrame(self.history)

    def get_plotly_grid_data(self):
        """Prepare data for Plotly scatter plot grid visualization."""
        if not self.cells:
            return pd.DataFrame(columns=['x', 'y_plotly', 'Type', 'Color', 'Resistance', 'Info'])

        cell_data = []
        for coords, cell in self.cells.items():
            # Plotly scatter typically has origin at bottom-left.
            # Our grid (y, x) has origin top-left. Transform y for plotting.
            plotly_y = self.grid_size - 1 - cell.y
            info_str = f"Type: {cell.LABEL}<br>Pos: ({cell.x}, {cell.y})"
            resistance = None
            if isinstance(cell, CancerCell):
                resistance = round(cell.resistance, 2)
                info_str += f"<br>Resistance: {resistance}"
            elif isinstance(cell, ImmuneCell):
                 info_str += f"<br>Steps Alive: {cell.steps_alive}"
                 if cell.is_activated:
                      info_str += "<br>Status: Activated"


            cell_data.append({
                'x': cell.x,
                'y_plotly': plotly_y,
                'Type': cell.LABEL,
                'Color': cell.COLOR,
                'Resistance': resistance, # Store for potential coloring/hover later
                'Info': info_str
            })
        return pd.DataFrame(cell_data)

# --- Streamlit App ---

st.set_page_config(layout="wide")
st.title("Cancer Simulation: Tumor Growth, Immune Response & Drug Treatment")

# --- Instructions ---
st.markdown("""
Welcome to the Cancer Simulation!
*   Use the **sidebar** on the left to set the initial parameters for the simulation.
*   Click **Start / Restart Simulation** to initialize the grid with the chosen parameters.
*   **Run N Step(s):** Executes a fixed number of simulation steps.
*   **Run Continuously:** Automatically runs the simulation step-by-step with a short delay (approx. 100ms) between steps. The grid and plots will update dynamically.
*   **Stop:** Pauses the simulation (either manual steps or continuous run).
*   The **Simulation Grid** visualizes the cells (Red=Cancer, Blue=Immune). Hover over cells for details.
*   The **Plots** below the grid show population dynamics and average cancer cell resistance over time.
""")
st.divider()


# --- Parameters Sidebar ---
with st.sidebar:
    st.header("Simulation Parameters")

    st.subheader("Grid & General")
    grid_size = st.slider("Grid Size (N x N)", 20, 100, 50, key="grid_size_slider")
    max_steps = st.number_input("Max Simulation Steps", 50, 1000, 200, key="max_steps_input")

    st.subheader("Initial Cells")
    initial_cancer_cells = st.slider("Initial Cancer Cells", 1, max(1,grid_size*grid_size//4), 10, key="init_cancer_slider") # Limit initial cells
    initial_immune_cells = st.slider("Initial Immune Cells", 0, max(1,grid_size*grid_size//2), 50, key="init_immune_slider")

    st.subheader("Cancer Cell Properties")
    cancer_growth_prob = st.slider("Growth Probability", 0.0, 1.0, 0.2, 0.01, key="cancer_growth_slider")
    cancer_metastasis_prob = st.slider("Metastasis Probability", 0.0, 0.1, 0.005, 0.001, format="%.3f", key="cancer_meta_slider")
    cancer_initial_resistance = st.slider("Initial Drug Resistance", 0.0, 1.0, 0.1, 0.01, key="cancer_res_slider")
    cancer_mutation_rate = st.slider("Mutation Rate", 0.0, 0.1, 0.01, 0.001, format="%.3f", key="cancer_mut_slider")

    st.subheader("Immune Cell Properties")
    immune_base_kill_prob = st.slider("Base Kill Probability", 0.0, 1.0, 0.3, 0.01, key="immune_kill_slider")
    immune_movement_prob = st.slider("Movement Probability", 0.0, 1.0, 0.8, 0.01, key="immune_move_slider")
    immune_lifespan = st.number_input("Lifespan (steps)", 10, 500, 100, key="immune_life_input")

    st.subheader("Drug Properties")
    drug_effect_base = st.slider("Base Drug Effect (Kill Prob)", 0.0, 1.0, 0.4, 0.01, key="drug_effect_slider")
    drug_resistance_interaction = st.slider("Resistance Interaction Factor", 0.0, 2.0, 1.0, 0.05, help="How much resistance reduces drug effect (1.0=linear)", key="drug_resint_slider")
    drug_immune_activation_boost = st.slider("Immune Activation Boost", 0.0, 1.0, 0.3, 0.01, help="Added kill prob when activated", key="drug_immune_boost_slider")
    drug_immune_boost_prob = st.slider("Immune Activation Probability", 0.0, 1.0, 0.7, 0.01, help="Prob. an immune cell near dying cancer gets activated", key="drug_immune_prob_slider")
    drug_immune_activation_radius = st.slider("Immune Activation Radius", 0, 5, 1, help="Radius around dying cancer cell to activate immune cells", key="drug_immune_rad_slider")


# Store parameters in a dictionary
simulation_params = {
    'grid_size': grid_size,
    'max_steps': max_steps,
    'initial_cancer_cells': initial_cancer_cells,
    'initial_immune_cells': initial_immune_cells,
    'cancer_growth_prob': cancer_growth_prob,
    'cancer_metastasis_prob': cancer_metastasis_prob,
    'cancer_initial_resistance': cancer_initial_resistance,
    'cancer_mutation_rate': cancer_mutation_rate,
    'immune_base_kill_prob': immune_base_kill_prob,
    'immune_movement_prob': immune_movement_prob,
    'immune_lifespan': immune_lifespan,
    'drug_effect_base': drug_effect_base,
    'drug_resistance_interaction': drug_resistance_interaction,
    'drug_immune_activation_boost': drug_immune_activation_boost,
    'drug_immune_boost_prob': drug_immune_boost_prob,
    'drug_immune_activation_radius': drug_immune_activation_radius,
}

# --- Simulation Control and State ---

# Initialize simulation state
if 'simulation' not in st.session_state:
    st.session_state.simulation = None
    st.session_state.running = False # Overall simulation active (not paused/stopped)
    st.session_state.continuously_running = False # Auto-step mode active
    st.session_state.history_df = pd.DataFrame()
    st.session_state.final_message = "" # To display end condition

col1, col2, col3, col4 = st.columns(4)

with col1:
    if st.button("Start / Restart Simulation", key="start_button"):
        st.session_state.simulation = Simulation(copy.deepcopy(simulation_params)) # Use deepcopy for params
        st.session_state.running = True
        st.session_state.continuously_running = False # Stop continuous if restarting
        st.session_state.history_df = st.session_state.simulation.get_history_df()
        st.session_state.final_message = ""
        st.success("Simulation Initialized.")
        st.rerun() # Rerun to update displays immediately

with col2:
    steps_to_run = st.number_input("Run Steps", min_value=1, max_value=max_steps, value=10, key="steps_input_manual")
    run_button = st.button(f"Run {steps_to_run} Step(s)", disabled=(st.session_state.simulation is None or not st.session_state.running or st.session_state.continuously_running), key="run_steps_button")

    if run_button:
        sim = st.session_state.simulation
        if sim:
            progress_bar = st.progress(0)
            steps_taken = 0
            for i in range(steps_to_run):
                if not st.session_state.running: break # Check if stopped externally
                keep_running = sim.step()
                steps_taken += 1
                # Need to update history inside loop if we want live plot updates during manual steps, but simpler to update after
                if not keep_running:
                    st.session_state.running = False # Simulation ended naturally
                    break
                progress_bar.progress((i + 1) / steps_to_run)

            progress_bar.empty()
            st.session_state.history_df = sim.get_history_df() # Update history after batch run
            st.info(f"Ran {steps_taken} steps. Current step: {sim.current_step}")
            if not st.session_state.running and st.session_state.final_message:
                 st.success(st.session_state.final_message) # Show end reason
            st.rerun() # Update displays

with col3:
     run_cont_button = st.button("Run Continuously", disabled=(st.session_state.simulation is None or not st.session_state.running or st.session_state.continuously_running), key="run_cont_button")
     if run_cont_button:
         st.session_state.continuously_running = True
         st.info("Running continuously...")
         st.rerun() # Start the continuous loop

with col4:
     stop_button = st.button("Stop", disabled=(st.session_state.simulation is None or (not st.session_state.running) or (not st.session_state.continuously_running and not run_button) ), key="stop_button") # Enable if running or continuously running
     if stop_button:
         st.session_state.running = False # Stop the simulation process
         st.session_state.continuously_running = False # Turn off continuous mode
         st.warning("Simulation stopped by user.")
         st.rerun()


# --- Dynamic Update Logic for Continuous Run ---
if st.session_state.get('simulation') and st.session_state.get('continuously_running') and st.session_state.get('running'):
    sim = st.session_state.simulation
    keep_running = sim.step()
    st.session_state.history_df = sim.get_history_df() # Update history

    if not keep_running:
        st.session_state.running = False # Simulation ended naturally
        st.session_state.continuously_running = False # Stop continuous mode
        if st.session_state.final_message:
            st.success(st.session_state.final_message) # Show end reason

    # Schedule the next rerun with a delay
    time.sleep(0.1) # 100 ms delay
    st.rerun()

# --- Visualization ---
# Use placeholders to potentially update plots faster
grid_placeholder = st.empty()
charts_placeholder = st.container() # Use a container for the two charts

if st.session_state.simulation:
    sim = st.session_state.simulation

    # --- Plotly Grid Visualization ---
    with grid_placeholder.container(): # Draw in the placeholder
        st.subheader(f"Simulation Grid (Step: {sim.current_step})")
        df_grid = sim.get_plotly_grid_data()

        fig_grid = go.Figure()

        if not df_grid.empty:
             # Add scatter trace for cells
             fig_grid.add_trace(go.Scatter(
                 x=df_grid['x'],
                 y=df_grid['y_plotly'],
                 mode='markers',
                 marker=dict(
                     color=df_grid['Color'],
                     size=max(5, 400 / sim.grid_size), # Adjust marker size based on grid size
                     symbol='square'
                 ),
                 text=df_grid['Info'], # Text appearing on hover
                 hoverinfo='text',
                 showlegend=False
             ))

        # Configure layout
        fig_grid.update_layout(
            xaxis=dict(
                range=[-0.5, sim.grid_size - 0.5],
                showgrid=True,
                gridcolor='lightgrey',
                zeroline=False,
                showticklabels=False,
                fixedrange=True # Prevent zoom/pan
            ),
            yaxis=dict(
                range=[-0.5, sim.grid_size - 0.5],
                showgrid=True,
                gridcolor='lightgrey',
                zeroline=False,
                showticklabels=False,
                scaleanchor="x", # Ensure square cells
                scaleratio=1,
                fixedrange=True # Prevent zoom/pan
            ),
            width=min(600, 800), # Adjust size as needed
            height=min(600, 800),
            margin=dict(l=10, r=10, t=40, b=10),
            paper_bgcolor='white',
            plot_bgcolor='white',
            # Add manual legend items if needed, or rely on text/color
             legend=dict(
                itemsizing='constant',
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            )
        )
        # Add dummy traces for legend (if desired)
        fig_grid.add_trace(go.Scatter(x=[None], y=[None], mode='markers', marker=dict(color=CancerCell.COLOR, size=10, symbol='square'), name='Cancer Cell'))
        fig_grid.add_trace(go.Scatter(x=[None], y=[None], mode='markers', marker=dict(color=ImmuneCell.COLOR, size=10, symbol='square'), name='Immune Cell'))

        st.plotly_chart(fig_grid, use_container_width=True) # Make it responsive


    # --- Time Series Plots ---
    with charts_placeholder: # Draw in the placeholder
        st.divider()
        col_chart1, col_chart2 = st.columns(2)

        if not st.session_state.history_df.empty:
            df_history = st.session_state.history_df

            with col_chart1:
                st.subheader("Cell Counts Over Time")
                df_melt = df_history.melt(id_vars=['Step'],
                                          value_vars=['Cancer Cells', 'Immune Cells'],
                                          var_name='Cell Type', value_name='Count')

                fig_line = px.line(df_melt, x='Step', y='Count', color='Cell Type',
                                   title="Population Dynamics", markers=False, # Use markers=False for potentially smoother continuous updates
                                   color_discrete_map={'Cancer Cells': CancerCell.COLOR, 'Immune Cells': ImmuneCell.COLOR})
                fig_line.update_layout(legend_title_text='Cell Type')
                st.plotly_chart(fig_line, use_container_width=True)

            with col_chart2:
                st.subheader("Average Cancer Cell Drug Resistance")
                fig_res = px.line(df_history, x='Step', y='Average Resistance',
                                  title="Average Resistance", markers=False) # Use markers=False
                fig_res.update_yaxes(range=[0, 1.05]) # Resistance is between 0 and 1, add buffer
                st.plotly_chart(fig_res, use_container_width=True)

        elif st.session_state.simulation: # If sim exists but no history yet (step 0)
             st.info("Run the simulation to see the plots.")


else:
    st.info("Click 'Start / Restart Simulation' to begin.")

# Add some explanations at the bottom as well if desired
# st.markdown(""" --- Explanation ... """)