em/state.py

"""
EM Embedded - State Management Module

Contains deferred state/controller proxy classes and state initialization.
These allow using state.update(), @state.change(), and ctrl.xxx at module
load time, then applying them when set_server() is called.
"""

from .globals import GRID_SIZES

__all__ = [
    "state", "ctrl", "set_server", "init_state",
    "enable_point_picking_on_plotter",
    "_apply_workflow_highlights", "_determine_workflow_step",
    "is_statevector_estimator_selected",
    "is_ibm_qpu_selected",
]


class _DeferredStateProxy:
    """
    A proxy that collects state defaults and change decorators at module load time,
    then applies them to the real server.state when bind() is called.
    """

    def __init__(self):
        self._state = None
        self._defaults = {}
        self._pending_changes = []  # list of (keys, func)

    def bind(self, real_state):
        """Bind to the real state object and apply all pending operations."""
        self._state = real_state
        # Apply all collected defaults
        if self._defaults:
            real_state.update(self._defaults)
            self._defaults.clear()
        # Apply all pending @state.change decorators
        for keys, func in self._pending_changes:
            real_state.change(*keys)(func)
        self._pending_changes.clear()

    @property
    def bound(self):
        return self._state is not None

    def update(self, d):
        """Collect defaults; apply immediately if bound."""
        if self._state is not None:
            self._state.update(d)
        else:
            self._defaults.update(d)

    def change(self, *keys):
        """
        Decorator factory that mimics @state.change("key1", "key2").
        If already bound, apply immediately. Otherwise, queue for later.
        """
        def decorator(func):
            if self._state is not None:
                # Already bound - register directly
                self._state.change(*keys)(func)
            else:
                # Queue for later
                self._pending_changes.append((keys, func))
            return func
        return decorator

    def __getattr__(self, name):
        if name.startswith("_"):
            raise AttributeError(name)
        if self._state is None:
            raise AttributeError(f"State not bound yet; cannot access '{name}'")
        return getattr(self._state, name)

    def __setattr__(self, name, value):
        if name.startswith("_"):
            object.__setattr__(self, name, value)
        elif self._state is not None:
            setattr(self._state, name, value)
        else:
            # Store as a default
            self._defaults[name] = value


class _DeferredControllerProxy:
    """
    A proxy that collects controller attribute assignments at module load time,
    then applies them when bind() is called.
    """

    def __init__(self):
        self._ctrl = None
        self._pending = {}

    def bind(self, real_ctrl):
        """Bind to the real controller and apply pending attributes."""
        self._ctrl = real_ctrl
        for name, value in self._pending.items():
            setattr(real_ctrl, name, value)
        self._pending.clear()

    @property
    def bound(self):
        return self._ctrl is not None

    def __getattr__(self, name):
        if name.startswith("_"):
            raise AttributeError(name)
        if self._ctrl is None:
            raise AttributeError(f"Controller not bound yet; cannot access '{name}'")
        return getattr(self._ctrl, name)

    def __setattr__(self, name, value):
        if name.startswith("_"):
            object.__setattr__(self, name, value)
        elif self._ctrl is not None:
            setattr(self._ctrl, name, value)
        else:
            self._pending[name] = value


# Module-level proxies (will be bound when set_server is called)
_server = None
state = _DeferredStateProxy()
ctrl = _DeferredControllerProxy()


def _noop(*_, **__):
    """No-op placeholder for controller methods."""
    pass


# Pre-register controller methods as no-ops
ctrl.qpu_ts_update = _noop
ctrl.qpu_ts_other_update = _noop
ctrl.view_update = _noop
ctrl.sim_ts_update = _noop
ctrl.geometry_preview_update = _noop
ctrl.excitation_preview_update = _noop
ctrl.qubit_plot_update = _noop


def set_server(server):
    """Bind the embedded EM module to the shared Trame server."""
    global _server
    _server = server
    state.bind(server.state)
    ctrl.bind(server.controller)


def get_server():
    """Get the bound server instance."""
    return _server


# --- Workflow Highlighting ---
_WORKFLOW_CARD_KEYS = [
    "overview_card_style",
    "geometry_card_style",
    "excitation_card_style",
    "meshing_card_style",
    "backend_card_style",
    "output_card_style",
]


def _workflow_highlight_style(active: bool) -> str:
    base = "font-size: 0.8rem; border: 1px solid transparent; transition: box-shadow 0.2s ease;"
    return f"{base} box-shadow: 0 0 0 2px #6200ea;" if active else base


def _determine_workflow_step() -> int:
    """Determine the current workflow step based on state."""
    if not state.bound:
        return 0
    if state.problem_selection is None:
        return 0
    if state.geometry_selection is None:
        return 1
    if state.dist_type is None:
        return 2
    if state.nx is None:
        return 3
    if state.backend_type is None:
        return 4
    return 5


def _apply_workflow_highlights(step_index: int):
    """Apply highlight styles to workflow cards."""
    for i, key in enumerate(_WORKFLOW_CARD_KEYS):
        setattr(state, key, _workflow_highlight_style(i == step_index))


# --- State Defaults ---
def _init_state_defaults():
    """Initialize all EM state defaults."""
    state.update({
        "problem_selection": None,
        "geometry_options": ["Square Domain", "Square Metallic Body"],
        "dist_type": None, 
        "impulse_x": 0.5, 
        "impulse_y": 0.5,
        "peak_pair": "(0.5, 0.5)",
        "mu_x": 0.5, 
        "mu_y": 0.5, 
        "sigma_x": 0.25, 
        "sigma_y": 0.15,
        "mu_pair": "(0.5, 0.5)",
        "sigma_pair": "(0.25, 0.15)",
        "nx": None, 
        "T": 1.0, 
        "time_val": 0.0,
        "L": 1.0,
        "output_type": "Surface Plot",
        "surface_field": "Ez",
        "timeseries_field": "Ez",
        "timeseries_points": "(0.5, 0.5)",
        "timeseries_gridpoints": "",
        "timeseries_point_info": "",
        "error_message": "",
        "is_running": False,
        "simulation_has_run": False,
        "geometry_selection": None,
        "show_upload_dialog": False,
        "uploaded_file_info": None,
        "show_upload_status": False,
        "upload_status_message": "",
        "coeff_permittivity": 1.0,
        "coeff_permeability": 1.0,
        "run_button_text": "RUN!",
        "backend_type": None,
        "selected_simulator": "IBM Qiskit simulator",
        "selected_qpu": "IBM QPU",
        "stop_button_disabled": True,
        "export_format": "vtk",
        "nx_slider_index": None,
    "grid_size_labels": [int(val) for val in GRID_SIZES],
        "show_export_status": False,
        "export_status_message": "",
        "logo_src": None,
        # Geometry-hole controls
        "hole_size_edge": 0.2,
        "hole_center_x": 0.5,
        "hole_center_y": 0.5,
        "hole_center_pair": "(0.5, 0.5)",
        "hole_error_message": "",
        "excitation_error_message": "",
        "excitation_info_message": "",
        "excitation_config_open": False,
        "dt_user": 0.1,
        "temporal_warning": "",
        # QPU monitor controls
        "qpu_field_components": "Ez",
        "qpu_monitor_gridpoints": "",
        "qpu_monitor_samples": "(0.5, 0.5)",
        "qpu_monitor_sample_info": "",
        "console_output": "Console initialized.\n",
        "pyvista_view_style": "aspect-ratio: 1 / 1; width: 100%;",
        # QPU Plotly controls
        "qpu_ts_fig": None,
        "qpu_ts_selected_time": None,
        "qpu_ts_ready": False,
        "qpu_plot_style": "display: none; width: 900px; height: 660px; margin: 0 auto;",
        "qpu_ts_other_ready": False,
        "qpu_other_plot_style": "display: none; width: 900px; height: 660px; margin: 0 auto;",
        "qpu_monitor_configs": [],
        "qpu_plot_filter": "All",
        "qpu_plot_field_options": ["All"],
        "qpu_plot_position_filter": "All positions",
        "qpu_plot_position_options": ["All positions"],
        # IBM QPU specific options (single field only - no "All")
        "ibm_qpu_field_options": ["Ez", "Hx", "Hy"],
        # Additional QPU monitor slots
        "qpu_monitor_count": 0,
        "qpu_field_components_2": "Ez",
        "qpu_monitor_gridpoints_2": "",
        "qpu_monitor_samples_2": "",
        "qpu_monitor_sample_info_2": "",
        "qpu_field_components_3": "Ez",
        "qpu_monitor_gridpoints_3": "",
        "qpu_monitor_samples_3": "",
        "qpu_monitor_sample_info_3": "",
        "qpu_field_components_4": "Ez",
        "qpu_monitor_gridpoints_4": "",
        "qpu_monitor_samples_4": "",
        "qpu_monitor_sample_info_4": "",
        "qpu_field_components_5": "Ez",
        "qpu_monitor_gridpoints_5": "",
        "qpu_monitor_samples_5": "",
        "qpu_monitor_sample_info_5": "",
        # Status
        "status_visible": False,
        "status_message": "Ready",
        "status_type": "info",
        "simulation_progress": 0,
        "show_progress": False,
        "console_logs": "Console initialized...\n",
        # --- EM Job Upload State ---
        "em_job_upload_error": "",
        "em_job_upload_success": "",
        "em_job_is_processing": False,
        "em_job_platform": "IBM",
        "em_job_id": "",
        "em_job_field_type": "Ez",
        "em_job_monitor_point": "(0, 0)",
        "em_job_total_time": 1.0,
        "em_job_snapshot_dt": 0.1,
        "em_job_nx": 4,
    })

    # Ensure hole snap state exists
    state.hole_snap = True


def init_state(force: bool = False):
    """Initialize EM state. Called after set_server()."""
    if not state.bound:
        return
    
    # Apply workflow highlights
    _apply_workflow_highlights(0)
    
    # Load logo
    from .utils import load_logo_data_uri
    state.logo_src = load_logo_data_uri()
    
    # NOTE: Point picking is enabled later, after the UI is built,
    # by calling enable_point_picking_on_plotter() or in redraw_surface_plot()
    
    if force:
        from .simulation import reset_to_defaults
        reset_to_defaults()
    
    # Initialize preview
    try:
        from .excitation import update_initial_state_preview
        update_initial_state_preview()
    except Exception:
        pass


def enable_point_picking_on_plotter():
    """Enable point picking on the plotter. Call AFTER build_ui()."""
    from .globals import plotter
    from .simulation import update_value_display
    try:
        plotter.disable_picking()
    except Exception:
        pass
    try:
        plotter.enable_point_picking(callback=update_value_display, show_message=False)
    except Exception:
        pass


def is_statevector_estimator_selected() -> bool:
    """Return True when the simulator dropdown targets the Statevector Estimator."""
    try:
        if state.backend_type != "Simulator":
            return False
        return (state.selected_simulator or "").strip().lower() == "statevector estimator"
    except AttributeError:
        return False


def is_ibm_qpu_selected() -> bool:
    """Return True when the QPU dropdown targets the IBM QPU."""
    try:
        if state.backend_type != "QPU":
            return False
        return (state.selected_qpu or "").strip().lower() == "ibm qpu"
    except AttributeError:
        return False


# Initialize state defaults at module load time
_init_state_defaults()