import matplotlib.pyplot as plt
import numpy as np

def plot_phi(all_agents_history, filename=None):
    plt.figure(figsize=(10, 6))
    for i, agent_history in enumerate(all_agents_history):
        plt.plot(agent_history['phi'], label=f'Agent {i+1} Phi')
    plt.xlabel('Step')
    plt.ylabel('Phi')
    plt.title('Agent Phi Over Time')
    plt.legend()
    plt.grid(True)
    if filename is not None:
        plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.close()

def plot_tau(all_agents_history, filename=None):
    plt.figure(figsize=(10, 6))
    for i, agent_history in enumerate(all_agents_history):
        plt.plot(agent_history['tau_eff'], label=f'Agent {i+1} Tau Eff')
    plt.xlabel('Step')
    plt.ylabel('Tau Effective')
    plt.title('Agent Tau Effective Over Time')
    plt.legend()
    plt.grid(True)
    if filename is not None:
        plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.close()

def plot_fitness(all_agents_history, filename=None):
    plt.figure(figsize=(10, 6))
    for i, agent_history in enumerate(all_agents_history):
        plt.plot(agent_history['fitness'], label=f'Agent {i+1} Fitness')
    plt.xlabel('Step')
    plt.ylabel('Fitness')
    plt.title('Agent Fitness Over Time')
    plt.legend()
    plt.grid(True)
    if filename is not None:
        plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.close()

def plot_coherence(coherence_list, filename=None):
    plt.figure(figsize=(10, 6))
    plt.plot(coherence_list)
    plt.xlabel('Step')
    plt.ylabel('Average Phi')
    plt.title('System Coherence Over Time (Average Phi)')
    plt.grid(True)
    if filename is not None:
        plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.close()

def plot_stability(stability_list, filename=None):
    plt.figure(figsize=(10, 6))
    plt.plot(stability_list)
    plt.xlabel('Step')
    plt.ylabel('Variance of Phi')
    plt.title('System Stability Over Time (Variance of Phi)')
    plt.grid(True)
    if filename is not None:
        plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.close()

print("visualization.py updated successfully.")