import gradio as gr
import json
import os
import matplotlib.pyplot as plt # Kept for potential future gr.Plot use or if figures are handled directly
from simulation import RFTSimulation
from utils import export_json, seed_everything
from IPython.display import Image, display # For Colab dev mode inline image display

# Helper function to run simulation and generate plots/export data
def run_simulation_and_plot(num_agents, steps, mutation_rate, drift_rate, seed=None):
    if seed is not None:
        seed_everything(seed)

    # Instantiate simulation with save_plots=True, so it saves plots during its run() method
    simulation = RFTSimulation(num_agents, steps, mutation_rate, drift_rate, save_plots=True)

    # Run the simulation. The internal run method will save plots if save_plots=True.
    # It returns all_agents_history, coherence_list, stability_list.
    all_agents_history, coherence_list, stability_list = simulation.run()

    # Prepare final agent states for JSON export
    final_agent_states = []
    for i, agent in enumerate(simulation.agents):
        final_agent_states.append({
            'id': i + 1,
            'phi': agent.phi,
            'tau_eff': agent.tau_eff,
            'tier': agent.tier,
            'fitness': agent.fitness
        })

    # Export JSON
    json_filename = "final_agent_states.json"
    export_json(final_agent_states, json_filename)

    # Return filenames of the generated plots and the JSON for Gradio or dev mode display
    plot_filenames = [
        'phi_plot.png',
        'tau_plot.png',
        'fitness_plot.png',
        'coherence_plot.png',
        'stability_plot.png'
    ]

    return (*plot_filenames, json_filename)

# Gradio Interface setup
iface = gr.Interface(
    fn=run_simulation_and_plot,
    inputs=[
        gr.Slider(1, 10, value=3, step=1, label="Number of Agents"),
        gr.Slider(10, 500, value=100, step=10, label="Steps"),
        gr.Slider(0.01, 0.5, value=0.05, step=0.01, label="Mutation Rate"),
        gr.Slider(0.01, 0.5, value=0.02, step=0.01, label="Drift Rate"),
        gr.Number(value=42, label="Random Seed (optional, leave empty for random)") # Clarified label
    ],
    outputs=[
        gr.Image(label="Phi Plot"), # Changed to gr.Image to display saved PNGs
        gr.Image(label="Tau Effective Plot"),
        gr.Image(label="Fitness Plot"),
        gr.Image(label="Coherence Plot"),
        gr.Image(label="Stability Plot"),
        gr.File(label="Exported Final Agent States (JSON)")
    ],
    title="RFT Agent Simulation Engine MVP",
    description="Simulate RFT Agents and visualize their dynamics."
)

# Colab-friendly "dev mode" or Gradio launch logic
# Check if running in Colab environment (simple heuristic)
IN_COLAB = 'COLAB_GPU' in os.environ or 'Google' in os.environ.get('COLAB_CLOUD_PLATFORM', '')

if IN_COLAB:
    print("Running in Colab development mode (skipping Gradio launch).")
    # Run the test simulation logic as per requirements
    num_agents = 3
    steps = 100
    mutation_rate = 0.05
    drift_rate = 0.02
    seed = 42

    print("\n--- Running Test Simulation in Dev Mode ---")
    # Call the core simulation function. It will save plots and JSON.
    plot_files_and_json_file = run_simulation_and_plot(num_agents, steps, mutation_rate, drift_rate, seed)
    plot_files = plot_files_and_json_file[:-1]
    json_file = plot_files_and_json_file[-1]

    print("\nSimulation complete. Displaying plots and exported JSON path:")
    for plot_file in plot_files:
        print(f"  - {os.path.abspath(plot_file)}")
        if os.path.exists(plot_file):
            display(Image(filename=plot_file)) # Display inline in Colab
        else:
            print(f"Warning: Plot file {plot_file} not found for inline display.")

    print(f"  - {os.path.abspath(json_file)}")
    print(f"Final agent states exported to {json_file}")
    print("\n--- End of Test Simulation in Dev Mode ---")

else:
    print("Running Gradio interface. Call iface.launch() to start it.")
    # For local development, uncomment this to launch Gradio:
    # iface.launch()

print("app.py updated successfully.")