import numpy as np
from agent import RFTAgent
import visualization # Added import for visualization

class RFTSimulation:
    def __init__(self, num_agents, steps, mutation_rate, drift_rate, save_plots=False): # Added save_plots
        self.agents = []
        self.steps = steps
        self.mutation_rate = mutation_rate
        self.drift_rate = drift_rate
        self.save_plots = save_plots # Stored save_plots
        self.coherence_list = [] # Initialized coherence_list
        self.stability_list = [] # Initialized stability_list

        for _ in range(num_agents):
            phi_init = np.random.uniform(-1, 1)
            tier = np.random.randint(1, 5)
            agent = RFTAgent(phi_init, tier, mutation_rate, drift_rate)
            self.agents.append(agent)

    def run(self):
        for _ in range(self.steps):
            for agent in self.agents:
                agent.step()
            # Append current coherence and stability after all agents have stepped
            self.coherence_list.append(self.compute_coherence())
            self.stability_list.append(self.compute_stability())

        # Conditional plot saving
        if self.save_plots:
            all_agents_history = self.get_history()
            visualization.plot_phi(all_agents_history, filename='phi_plot.png')
            visualization.plot_tau(all_agents_history, filename='tau_plot.png')
            visualization.plot_fitness(all_agents_history, filename='fitness_plot.png')
            visualization.plot_coherence(self.coherence_list, filename='coherence_plot.png')
            visualization.plot_stability(self.stability_list, filename='stability_plot.png')

        return (self.get_history(), self.coherence_list, self.stability_list) # Modified return value

    def get_history(self):
        return [agent.history for agent in self.agents]

    def compute_coherence(self):
        # Average phi across agents at the current step
        if not self.agents:
            return 0.0
        return np.mean([agent.phi for agent in self.agents])

    def compute_stability(self):
        # Variance of phi across agents at the current step
        if not self.agents:
            return 0.0
        return np.var([agent.phi for agent in self.agents])

print("simulation.py updated successfully.")