gradio_app.py

# -*- coding: utf-8 -*-
"""
Gradio Web Interface for the Quantum Circuit Optimization Environment.

Provides a visual, interactive UI with dropdowns for task selection,
gate operations, and real-time score visualization -- no JSON typing required.
"""

import json
import logging
import os
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
_ROOT = os.path.dirname(_HERE)
for _p in (_ROOT, _HERE):
    if _p not in sys.path:
        sys.path.insert(0, _p)

import gradio as gr

from models import ActionType, GateType, QuantumAction, QuantumObservation
from server.my_env_environment import QuantumCircuitEnvironment

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# ---------------------------------------------------------------------------
# Global environment instance (per-process, single user for Spaces)
# ---------------------------------------------------------------------------
env = QuantumCircuitEnvironment(seed=42)
current_obs = None
step_history = []


# ---------------------------------------------------------------------------
# Helper: format observation for display
# ---------------------------------------------------------------------------
def format_scores(obs):
    """Build a dict of scores from obs metadata for JSON display."""
    meta = obs.metadata or {}
    return {
        "fidelity": round(meta.get("fidelity_score", obs.fidelity), 6),
        "efficiency": round(meta.get("efficiency_score", 0), 6),
        "noise": round(meta.get("noise_score", 0), 6),
        "constraints": round(meta.get("constraints_score", 0), 6),
        "aggregate": round(meta.get("aggregate_score", obs.score), 6),
    }


def format_obs_json(obs):
    """Convert observation to a clean JSON dict for display."""
    return {
        "task_id": obs.task_id,
        "num_qubits": obs.num_qubits,
        "step": obs.steps_taken,
        "max_steps": obs.max_steps,
        "done": obs.done,
        "fidelity": round(obs.fidelity, 6),
        "depth": obs.depth,
        "gate_count": obs.gate_count,
        "score": round(obs.score, 6),
        "reward": round(obs.reward, 6),
        "noise_estimate": round(obs.noise_estimate, 6),
        "circuit_gates": obs.circuit_gates,
        "valid_actions": obs.valid_actions,
        "target": obs.target_description,
        "scores_breakdown": format_scores(obs),
    }


def format_circuit_display(obs):
    """Build a human-readable circuit string."""
    if not obs.circuit_gates:
        return "(empty circuit)"
    lines = []
    for i, g in enumerate(obs.circuit_gates):
        gate = g.get("gate", "?")
        qubits = g.get("qubits", [])
        param = g.get("parameter")
        q_str = ",".join(str(q) for q in qubits)
        if param is not None:
            lines.append(f"  {i+1}. {gate}({q_str}, theta={param:.4f})")
        else:
            lines.append(f"  {i+1}. {gate}({q_str})")
    return "\n".join(lines)


def build_score_bars(obs):
    """Build HTML score bars for visual display."""
    scores = format_scores(obs)
    colors = {
        "fidelity": "#4CAF50",
        "efficiency": "#2196F3",
        "noise": "#FF9800",
        "constraints": "#9C27B0",
        "aggregate": "#F44336",
    }
    labels = {
        "fidelity": "Fidelity",
        "efficiency": "Efficiency",
        "noise": "Noise Resilience",
        "constraints": "Constraints",
        "aggregate": "AGGREGATE",
    }

    html = '<div style="font-family: monospace; padding: 10px;">'
    for key in ["fidelity", "efficiency", "noise", "constraints", "aggregate"]:
        val = scores[key]
        pct = max(0, min(100, val * 100))
        color = colors[key]
        label = labels[key]
        bold = "font-weight: bold; font-size: 16px;" if key == "aggregate" else ""
        html += f'''
        <div style="margin: 6px 0; {bold}">
            <div style="display: flex; align-items: center; gap: 8px;">
                <span style="width: 140px; text-align: right;">{label}</span>
                <div style="flex: 1; background: #333; border-radius: 4px; height: 22px; position: relative;">
                    <div style="width: {pct}%; background: {color}; height: 100%; border-radius: 4px;
                                transition: width 0.3s ease;"></div>
                </div>
                <span style="width: 60px;">{val:.4f}</span>
            </div>
        </div>'''
    html += '</div>'
    return html


# ---------------------------------------------------------------------------
# Gradio callbacks
# ---------------------------------------------------------------------------

def do_reset(task_name):
    """Reset environment with selected task."""
    global current_obs, step_history
    task_map = {"Bell State (Easy)": "easy", "GHZ State (Medium)": "medium", "Unitary Approx (Hard)": "hard"}
    task_id = task_map.get(task_name, "easy")
    current_obs = env.reset(task_id=task_id)
    step_history = []

    obs_json = format_obs_json(current_obs)
    circuit_text = format_circuit_display(current_obs)
    score_html = build_score_bars(current_obs)
    status = f"Environment reset to: {task_name}\nTarget: {current_obs.target_description}\nQubits: {current_obs.num_qubits} | Max steps: {current_obs.max_steps}"

    return (
        status,
        circuit_text,
        score_html,
        json.dumps(obs_json, indent=2),
        "",  # clear history
    )


def do_step(action_type, gate_type, qubit_0, qubit_1, parameter):
    """Execute one step in the environment."""
    global current_obs, step_history

    if current_obs is None:
        return (
            "ERROR: Reset the environment first!",
            "(no circuit)", "", "{}", "",
        )
    if current_obs.done:
        return (
            "Episode is DONE. Click Reset to start a new one.",
            format_circuit_display(current_obs),
            build_score_bars(current_obs),
            json.dumps(format_obs_json(current_obs), indent=2),
            "\n".join(step_history),
        )

    # Build action
    try:
        at = ActionType(action_type)

        gate = None
        qubits = []
        param = None

        if at == ActionType.ADD:
            gate = GateType(gate_type) if gate_type else None
            if gate in (GateType.CNOT,):
                qubits = [int(qubit_0), int(qubit_1)]
            elif gate == GateType.SWAP:
                qubits = [int(qubit_0), int(qubit_1)]
            else:
                qubits = [int(qubit_0)]
            if gate in (GateType.RX, GateType.RZ) and parameter is not None:
                param = float(parameter)

        elif at == ActionType.SWAP:
            qubits = [int(qubit_0), int(qubit_1)]

        elif at == ActionType.REMOVE:
            pass  # no args needed

        elif at == ActionType.PARAM:
            if parameter is not None:
                param = float(parameter)

        action = QuantumAction(
            action_type=at,
            gate=gate,
            qubits=qubits,
            parameter=param,
        )
        current_obs = env.step(action)

    except Exception as e:
        return (
            f"ACTION ERROR: {e}",
            format_circuit_display(current_obs) if current_obs else "(no circuit)",
            build_score_bars(current_obs) if current_obs else "",
            json.dumps(format_obs_json(current_obs), indent=2) if current_obs else "{}",
            "\n".join(step_history),
        )

    # Record step
    obs_json = format_obs_json(current_obs)
    step_entry = (
        f"Step {current_obs.steps_taken}: {action_type}"
        + (f" {gate_type}({','.join(str(q) for q in qubits)})" if gate else "")
        + f" -> fid={current_obs.fidelity:.4f} score={current_obs.score:.4f}"
        + f" reward={current_obs.reward:+.4f}"
        + (" [DONE]" if current_obs.done else "")
    )
    step_history.append(step_entry)

    circuit_text = format_circuit_display(current_obs)
    score_html = build_score_bars(current_obs)

    done_msg = ""
    if current_obs.done:
        done_msg = f"\n\nEPISODE COMPLETE! Final score: {current_obs.score:.4f}"

    status = (
        f"Step {current_obs.steps_taken}/{current_obs.max_steps}"
        + f" | Fidelity: {current_obs.fidelity:.4f}"
        + f" | Score: {current_obs.score:.4f}"
        + f" | Reward: {current_obs.reward:+.4f}"
        + done_msg
    )

    return (
        status,
        circuit_text,
        score_html,
        json.dumps(obs_json, indent=2),
        "\n".join(step_history),
    )


# ---------------------------------------------------------------------------
# Build Gradio UI
# ---------------------------------------------------------------------------

def create_gradio_app():
    """Create the Gradio Blocks interface."""

    with gr.Blocks(
        title="Quantum Circuit Optimizer",
        theme=gr.themes.Base(
            primary_hue="indigo",
            secondary_hue="blue",
            neutral_hue="slate",
            font=gr.themes.GoogleFont("Inter"),
        ),
        css="""
        .score-panel { border: 1px solid #444; border-radius: 8px; padding: 12px; background: #1a1a2e; }
        .circuit-box { font-family: 'Courier New', monospace; font-size: 14px; }
        .header-row { text-align: center; margin-bottom: 16px; }
        """,
    ) as demo:

        gr.Markdown(
            """
            # Quantum Circuit Optimizer
            ### Noise-aware, hardware-constrained RL environment for quantum circuit design

            **How to use:** Select a task, click Reset, then build your circuit step-by-step using the controls below.
            Try to maximize the aggregate score by matching the target quantum state!
            """,
        )

        with gr.Row():
            # --- Left column: Controls ---
            with gr.Column(scale=1):
                gr.Markdown("### 1. Select Task & Reset")
                task_dropdown = gr.Dropdown(
                    choices=["Bell State (Easy)", "GHZ State (Medium)", "Unitary Approx (Hard)"],
                    value="Bell State (Easy)",
                    label="Task",
                )
                reset_btn = gr.Button("Reset Environment", variant="primary", size="lg")

                gr.Markdown("### 2. Build Circuit")
                action_dropdown = gr.Dropdown(
                    choices=["ADD", "REMOVE", "SWAP", "PARAM", "STOP"],
                    value="ADD",
                    label="Action Type",
                )
                gate_dropdown = gr.Dropdown(
                    choices=["H", "X", "CNOT", "RX", "RZ"],
                    value="H",
                    label="Gate (for ADD)",
                )

                with gr.Row():
                    qubit_0 = gr.Number(value=0, label="Qubit 0", precision=0)
                    qubit_1 = gr.Number(value=1, label="Qubit 1 (CNOT/SWAP)", precision=0)

                parameter = gr.Number(value=None, label="Parameter (RX/RZ angle)", precision=4)
                step_btn = gr.Button("Execute Step", variant="secondary", size="lg")

            # --- Right column: Display ---
            with gr.Column(scale=2):
                status_box = gr.Textbox(label="Status", lines=3, interactive=False)

                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### Circuit")
                        circuit_display = gr.Textbox(
                            label="Current Circuit",
                            lines=8,
                            interactive=False,
                            elem_classes=["circuit-box"],
                        )
                    with gr.Column():
                        gr.Markdown("### Scores")
                        score_bars = gr.HTML(elem_classes=["score-panel"])

                with gr.Accordion("Full JSON Response", open=False):
                    json_display = gr.Code(language="json", label="Observation JSON", lines=20)

                with gr.Accordion("Step History", open=True):
                    history_display = gr.Textbox(label="History", lines=8, interactive=False)

        # --- Wire events ---
        reset_outputs = [status_box, circuit_display, score_bars, json_display, history_display]
        reset_btn.click(fn=do_reset, inputs=[task_dropdown], outputs=reset_outputs)

        step_outputs = [status_box, circuit_display, score_bars, json_display, history_display]
        step_btn.click(
            fn=do_step,
            inputs=[action_dropdown, gate_dropdown, qubit_0, qubit_1, parameter],
            outputs=step_outputs,
        )

    return demo


# ---------------------------------------------------------------------------
# Standalone entry point
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    demo = create_gradio_app()
    demo.launch(server_name="0.0.0.0", server_port=7860)