# interface.py
# Author: Liam Grinstead

import gradio as gr
from app import run_simulation
from registry_utils import append_to_registry
from registry_viewer import display_registry
from mutation_designer import build_mutation
from lineage_tracker import register_lineage
from lineage_visualizer import render_lineage_tree
from waveform_renderer import render_waveform
from leaderboard import generate_leaderboard
from codex.formulas import GVU_FORMULAS, rft_invariants
import stage1   # keep imports at the top

# Safety guard to ensure agent has required keys
def ensure_agent_shape(agent: dict, mutation_profile: dict) -> dict:
    if not isinstance(agent, dict):
        agent = {}
    agent.setdefault("id", mutation_profile.get("agent_id", "Agent_Unknown"))
    agent.setdefault("tier", mutation_profile.get("tier_drift", "Tier_1"))
    agent.setdefault("symbolic_operators", mutation_profile.get("symbolic_operators", ["R", "O", "T", "P"]))
    agent.setdefault("emotional_resonance", mutation_profile.get("emotional_resonance", False))

    overlay = agent.get("collapse_overlay", {})
    if not isinstance(overlay, dict):
        overlay = {}
    if mutation_profile.get("collapse_overlay"):
        ov = mutation_profile["collapse_overlay"]
        overlay.setdefault("tau_eff", ov.get("tau_eff"))
        overlay.setdefault("beta_band", ov.get("beta_band"))
        overlay.setdefault("operator_weights", ov.get("operator_weights"))
    overlay.setdefault("tau_eff", 1.8 if mutation_profile.get("collapse_torque") == "Gen6508_M5" else 1.2)
    overlay.setdefault("beta_band", 0.65 if mutation_profile.get("collapse_torque") == "Gen6508_M5" else 0.4)
    overlay.setdefault("operator_weights", {("R","O"): 0.9, ("T","P"): 0.7})
    agent["collapse_overlay"] = overlay
    return agent

# --- Simulation ---
def simulate(agent_id, collapse_torque, emotional_resonance, tier_drift):
    mutation_profile = build_mutation(agent_id, collapse_torque, tier_drift, emotional_resonance)
    agent, sha512 = run_simulation(agent_id, mutation_profile)
    agent = ensure_agent_shape(agent, mutation_profile)

    score = GVU_FORMULAS["Formula_20"].evaluate(agent)
    invariants = rft_invariants(agent) or {}
    tau = invariants.get("tau_eff", "?")
    beta = invariants.get("beta_band", "?")
    op_count = invariants.get("operator_count", 0)
    tier_level = invariants.get("tier_level", 1)

    fields = {
        "Φᵢ": f"<div><b>Φᵢ Awareness</b><br>Tier={agent.get('tier')} τ_eff={tau}</div>",
        "Kᵢⱼ": f"<div><b>Kᵢⱼ Coupling</b><br>Operators={op_count}</div>",
        "Φ_col": f"<div><b>Φ_col Collective</b><br>Score={score}</div>"
    }

    append_to_registry(agent_id, collapse_torque, tier_drift, emotional_resonance, score, sha512)

    summary = (
        f"📊 <b>Fitness (GVU):</b> {score}<br>"
        f"🧷 <b>Invariants:</b> τ_eff={tau}, β={beta}, |K|={op_count}, tier={tier_level}<br>"
        f"🔐 <b>SHA-512:</b> <code>{sha512}</code>"
    )

    wf = render_waveform(agent, score)
    return fields["Φᵢ"], fields["Kᵢⱼ"], fields["Φ_col"], wf, summary

# --- Forge ---
def forge_agent(parent_id, new_id, collapse_torque, emotional_resonance, tier_drift, max_depth):
    mutation_profile = build_mutation(new_id, collapse_torque, tier_drift, emotional_resonance)
    agent, _ = run_simulation(new_id, mutation_profile)
    agent = ensure_agent_shape(agent, mutation_profile)
    register_lineage(parent_id, new_id, {
        "tier_drift": tier_drift,
        "collapse_torque": collapse_torque,
        "symbolic_operators": agent.get("symbolic_operators", [])
    })
    return render_lineage_tree(parent_id, max_depth=max_depth)

# --- Validation Stage Wrapper ---
def run_stage(stage_name, mode, epochs, batch, lr):
    if stage_name == "Stage 1 — CIFAR-10 Baseline":
        stage1.train(
            mode=mode,
            epochs=int(epochs),
            batch=int(batch),
            lr=float(lr),
            log_path="stage1_cifar10_log.jsonl"
        )
        return f"Stage 1 complete. Telemetry saved to stage1_cifar10_log.jsonl"
    else:
        return "Stage not yet implemented."

# --- Interface Layout ---
with gr.Blocks(theme="soft") as demo:
    gr.Markdown("# 🧠 RFT Codex Sovereign")
    gr.Markdown("Rendered Frame Theory simulation, lineage, and GVU sealing. Author: Liam Grinstead.")

    # --- What is this Tab ---
    with gr.Tab("What is this?"):
        gr.Markdown("... your existing intro markdown ...")

    # --- Simulation Tab ---
    with gr.Tab("Simulate Agent"):
        with gr.Row():
            agent_id = gr.Dropdown(["Agent_5", "Agent_7", "Agent_1032"], label="Agent ID")
            collapse_torque = gr.Dropdown(["Gen6508_M5", "Gen26_M23"], label="Collapse Torque Overlay")
            emotional_resonance = gr.Dropdown(["Yes", "No"], label="Inject Emotional Resonance")
            tier_drift = gr.Dropdown(["Tier_1", "Tier_2", "Tier_6"], label="Tier Drift")
        simulate_btn = gr.Button("Run Simulation")

        with gr.Row():
            phi_i = gr.HTML(label="Φᵢ Awareness Field")
            k_ij = gr.HTML(label="Kᵢⱼ Correlation Kernel")
        with gr.Row():
            phi_col = gr.HTML(label="Φ_col Coherence Field")
            waveform = gr.HTML(label="Collapse Torque Waveform")

        summary = gr.HTML(label="Simulation Summary")

        simulate_btn.click(
            lambda agent_id, collapse_torque, emotional_resonance, tier_drift:
                simulate(agent_id, collapse_torque, emotional_resonance == "Yes", tier_drift),
            inputs=[agent_id, collapse_torque, emotional_resonance, tier_drift],
            outputs=[phi_i, k_ij, phi_col, waveform, summary]
        )

    # --- Registry Tab ---
    with gr.Tab("View Registry"):
        registry_output = gr.Textbox(label="Codex Registry", lines=20)
        refresh_btn = gr.Button("Refresh Registry")
        refresh_btn.click(display_registry, outputs=registry_output)

    # --- Forge Tab ---
    with gr.Tab("Codex Forge"):
        gr.Markdown("### 🧬 Evolve a New Agent from a Parent")
        parent_id = gr.Dropdown(["Agent_5", "Agent_7", "Agent_1032"], label="Parent Agent")
        new_id = gr.Textbox(label="New Agent ID")
        forge_torque = gr.Dropdown(["Gen6508_M5", "Gen26_M23"], label="Collapse Torque")
        forge_resonance = gr.Dropdown(["Yes", "No"], label="Inject Emotional Resonance")
        forge_tier = gr.Dropdown(["Tier_1", "Tier_2", "Tier_6"], label="Tier Drift")
        max_depth = gr.Slider(1, 8, value=5, step=1, label="Lineage depth")
        forge_btn = gr.Button("Forge Agent")
        lineage_svg_output = gr.HTML(label="Lineage Visualization")

        forge_btn.click(
            lambda parent_id, new_id, forge_torque, forge_resonance, forge_tier, max_depth:
                forge_agent(parent_id, new_id, forge_torque, forge_resonance == "Yes", forge_tier, max_depth),
            inputs=[parent_id, new_id, forge_torque, forge_resonance, forge_tier, max_depth],
            outputs=lineage_svg_output
        )

    # --- Leaderboard Tab ---
    with gr.Tab("Leaderboard"):
        leaderboard_output = gr.Textbox(label="Top Agents", lines=15)
        refresh_leaderboard = gr.Button("Refresh Leaderboard")
        refresh_leaderboard.click(generate_leaderboard, outputs=leaderboard_output)

    # --- Codex Reference Tab ---
    with gr.Tab("Codex Reference"):
        gr.Markdown("... your existing reference markdown ...")

    # --- Validation Stages Tab ---
    with gr.Tab("Validation Stages"):
        stage = gr.Dropdown(
            ["Stage 1 — CIFAR-10 Baseline"],  # extend with Stage 2–12 later
            label="Select Stage"
        )
        mode = gr.Dropdown(["RFT", "BASE"], label="Mode")
        epochs = gr.Number(label="Epochs", value=5)
        batch = gr.Number(label="Batch Size", value=256)
        lr = gr.Number(label="Learning Rate", value=5e-4)
        val_output = gr.Textbox(label="Validation Output")
        run_button = gr.Button("Run Stage")
        run_button.click(
            fn=run_stage,
            inputs=[stage, mode, epochs, batch, lr],
            outputs=val_output
        )

if __name__ == "__main__":
    demo.launch()