import gradio as gr
import json
import os
import matplotlib.pyplot as plt
from simulation import RFTSimulation
from utils import export_json, seed_everything

# Detect Colab
IN_COLAB = 'COLAB_GPU' in os.environ or 'Google' in os.environ.get('COLAB_CLOUD_PLATFORM', '')
if IN_COLAB:
    from IPython.display import Image, display

# -----------------------------
# Core Simulation Function
# -----------------------------

def run_simulation_and_plot(num_agents, steps, mutation_rate, drift_rate, seed=None):
    if seed is not None:
        seed_everything(seed)

    simulation = RFTSimulation(num_agents, steps, mutation_rate, drift_rate, save_plots=True)
    all_agents_history, coherence_list, stability_list = simulation.run()

    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
        })

    json_filename = "final_agent_states.json"
    export_json(final_agent_states, json_filename)

    plot_filenames = [
        'phi_plot.png',
        'tau_plot.png',
        'fitness_plot.png',
        'coherence_plot.png',
        'stability_plot.png'
    ]

    return (*plot_filenames, json_filename)

# -----------------------------
# PREMIUM UI (Always Unlocked)
# -----------------------------

with gr.Blocks(title="RFT Agent Simulation Engine — Premium") as iface:

    gr.Markdown("# 🔥 RFT Agent Simulation Engine — Premium Edition")
    gr.Markdown("Full‑power simulation engine with extended limits, advanced controls, and unrestricted exports.")

    with gr.Group():
        gr.Markdown("### Premium Simulation Controls")

        num_agents = gr.Slider(1, 200, value=10, step=1, label="Number of Agents (Premium)")
        steps = gr.Slider(10, 20000, value=500, step=10, label="Steps (Premium)")
        mutation_rate = gr.Slider(0.001, 1.0, value=0.05, step=0.001, label="Mutation Rate")
        drift_rate = gr.Slider(0.001, 1.0, value=0.02, step=0.001, label="Drift Rate")
        seed = gr.Number(value=42, label="Random Seed")

    outputs = [
        gr.Image(label="Phi Plot"),
        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)")
    ]

    run_button = gr.Button("🚀 Run Premium Simulation")

    run_button.click(
        run_simulation_and_plot,
        inputs=[num_agents, steps, mutation_rate, drift_rate, seed],
        outputs=outputs
    )


# -----------------------------
# Launch Logic
# -----------------------------

if IN_COLAB:
    print("Running in Colab development mode (skipping Gradio launch).")
else:
    iface.launch()