path: root/core/profile.py

"""
core/profile.py

Lab profile — saves and loads the complete operator configuration:

  • Controls       — all control widgets (type, title, device, channel, params)
  • Channels       — per-channel visibility, name overrides
  • Signal pipelines — filter stacks per channel
  • Derived channels — all virtual/computed channels
  • Plot layout     — pane arrangement, traces, axes, time window

Profiles are stored as human-readable JSON files (.labdaq).
"""

from __future__ import annotations
import json
import os
from dataclasses import dataclass, field, asdict
from typing import Any, Dict, List, Optional


# ── Helpers ───────────────────────────────────────────────────────────────────

def _safe(d: dict, key: str, default=None):
    return d.get(key, default)


# ══════════════════════════════════════════════════════════════════════════════
#  Profile data model
# ══════════════════════════════════════════════════════════════════════════════

@dataclass
class ProfileChannelOverride:
    """Per-channel settings that survive device reconnects."""
    device_id:  str
    channel_id: str
    name:       str  = ""
    unit:       str  = ""
    enabled:    bool = True
    color:      str  = ""


@dataclass
class ProfilePipeline:
    """Serialised filter pipeline for one channel."""
    device_id:  str
    channel_id: str
    enabled:    bool = True
    filters:    List[Dict[str, Any]] = field(default_factory=list)
    # filters = [{"type": "low_pass", "alpha": 0.1}, ...]


@dataclass
class ProfileDerived:
    """Serialised derived/virtual channel."""
    channel_id:  str
    name:        str
    unit:        str   = ""
    color:       str   = "#f72585"
    kind:        str   = "expression"
    sources:     List  = field(default_factory=list)  # [(dev_id, ch_id), ...]
    expression:  str   = ""
    script:      str   = ""
    params:      Dict  = field(default_factory=dict)
    enabled:     bool  = True


@dataclass
class Profile:
    name:         str   = "Untitled Profile"
    version:      str   = "1.0"
    devices:      List[Dict]                   = field(default_factory=list)
    controls:     List[Dict]                   = field(default_factory=list)
    channels:     List[Dict]                   = field(default_factory=list)
    pipelines:    List[Dict]                   = field(default_factory=list)
    derived:      List[Dict]                   = field(default_factory=list)
    plot:         Optional[Dict]               = None   # LayoutConfig.to_json()
    settings:     Dict[str, Any]               = field(default_factory=dict)
    plugin_state:   Dict[str, Any]             = field(default_factory=dict)
    # plugin_state = {plugin_id: plugin.get_save_state()}
    plugins_enabled: List[str]                 = field(default_factory=list)
    # plugins_enabled = [plugin_id, ...]  — which plugins were on when saved
    script_vars:  Dict[str, Any]               = field(default_factory=dict)
    # script_vars — shared vars dict accessible as `vars` in all expressions

    def to_json(self) -> str:
        return json.dumps(asdict(self), indent=2)

    @staticmethod
    def from_json(s: str) -> "Profile":
        d = json.loads(s)
        return Profile(
            name         = d.get("name", ""),
            version      = d.get("version", "1.0"),
            devices      = d.get("devices", []),
            controls     = d.get("controls", []),
            channels     = d.get("channels", []),
            pipelines    = d.get("pipelines", []),
            derived      = d.get("derived", []),
            plot         = d.get("plot"),
            settings        = d.get("settings", {}),
            plugin_state    = d.get("plugin_state", {}),
            plugins_enabled = d.get("plugins_enabled", []),
            script_vars     = d.get("script_vars", {}),
        )

    def save(self, path: str):
        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
        with open(path, "w") as f:
            f.write(self.to_json())

    @staticmethod
    def load(path: str) -> "Profile":
        with open(path) as f:
            return Profile.from_json(f.read())


# ══════════════════════════════════════════════════════════════════════════════
#  ProfileManager — serialises / deserialises the live app state
# ══════════════════════════════════════════════════════════════════════════════

class ProfileManager:
    """
    Converts between a Profile dict and the live application objects.

    Usage:
        # Save
        profile = ProfileManager.capture(registry, processor, chart_cfg,
                                         control_specs, settings)
        profile.save("my_lab.labdaq")

        # Load
        profile = Profile.load("my_lab.labdaq")
        ProfileManager.apply(profile, registry, processor, ...)
    """

    # Plugin device factories registered at runtime (device_type → class)
    _extra_factories: dict = {}

    @staticmethod
    def register_device_factory(device_type: str, factory):
        """Register a plugin device class so it can be restored from profiles."""
        ProfileManager._extra_factories[device_type] = factory

    @staticmethod
    def unregister_device_factory(device_type: str):
        ProfileManager._extra_factories.pop(device_type, None)

    # ── Capture ──────────────────────────────────────────────────────────────

    @staticmethod
    def capture(
        registry,
        processor,
        plot_cfg,
        control_specs: List,
        settings: dict,
        profile_name: str = "Profile",
        plugin_manager=None,
    ) -> Profile:
        p = Profile(name=profile_name)

        # Devices
        for dev in registry.all_instances():
            if hasattr(dev, "get_save_config"):
                p.devices.append(dev.get_save_config())

        # Controls
        p.controls = [spec.to_dict() for spec in control_specs]

        # Channel overrides
        for dev in registry.all_instances():
            for ch in dev.info.channels:
                p.channels.append({
                    "device_id":  dev.info.device_id,
                    "channel_id": ch.channel_id,
                    "name":       ch.name,
                    "unit":       ch.unit,
                    "enabled":    ch.enabled,
                    "color":      ch.color,
                })

        # Signal pipelines
        for key, pipeline in processor._pipelines.items():
            p.pipelines.append({
                "device_id":  pipeline.device_id,
                "channel_id": pipeline.channel_id,
                "enabled":    pipeline.enabled,
                "filters":    [f.to_dict() for f in pipeline.filters],
            })

        # Derived channels
        for dc in processor.get_derived():
            p.derived.append({
                "channel_id": dc.channel_id,
                "name":       dc.name,
                "unit":       dc.unit,
                "color":      dc.color,
                "kind":       dc.kind,
                "sources":      [list(s) for s in dc.sources],
                "source_names": list(dc.source_names),
                "expression":   dc.expression,
                "script":       dc.script,
                "params":       {k: v for k, v in dc.params.items()
                                 if not k.startswith("_")},
                "enabled":      dc.enabled,
            })

        # Plot layout
        if plot_cfg is not None:
            try:
                from ui.windows.plot_window import LayoutConfig
                if hasattr(plot_cfg, "to_json"):
                    p.plot = json.loads(plot_cfg.to_json())
            except Exception:
                pass

        p.settings = dict(settings)

        # Shared script variables
        with processor._lock:
            p.script_vars = dict(processor._script_vars)

        # Plugin state + enabled list
        if plugin_manager is not None:
            p.plugins_enabled = plugin_manager.get_enabled_ids()
            for plugin in plugin_manager.get_loaded():
                try:
                    state = plugin.get_save_state()
                    if state:
                        p.plugin_state[plugin.plugin_id] = state
                except Exception:
                    pass

        return p

    # ── Apply ─────────────────────────────────────────────────────────────────

    @staticmethod
    def apply(
        profile: Profile,
        registry,
        processor,
        control_panel,
        settings_ref: dict,
        engine=None,
        plugin_manager=None,
    ):
        """Apply a loaded profile to the live application state."""
        from ui.control_editor import ControlSpec
        from core.signal_processor import (
            ChannelPipeline, DerivedChannel,
            filter_from_dict,
        )

        # ── Devices ──────────────────────────────────────────────────────
        _DEVICE_FACTORIES = dict(ProfileManager._extra_factories)
        try:
            from devices.arduino_device import ArduinoDevice
            _DEVICE_FACTORIES["arduino"] = ArduinoDevice
        except Exception:
            pass
        try:
            from devices.nidaqmx_device import NidaqmxDevice
            _DEVICE_FACTORIES["nidaqmx"] = NidaqmxDevice
        except Exception:
            pass
        try:
            from devices.serial_device import SerialDevice
            _DEVICE_FACTORIES["serial"] = SerialDevice
        except Exception:
            pass
        # Backward compatibility with old profiles
        try:
            from devices.analog_input import AnalogInputDevice
            _DEVICE_FACTORIES["analog_input"] = AnalogInputDevice
        except Exception:
            pass
        try:
            from devices.digital_io import DigitalIODevice
            _DEVICE_FACTORIES["digital_io"] = DigitalIODevice
        except Exception:
            pass

        # Clear all existing devices — profile defines the complete device set
        for dev in list(registry.all_instances()):
            try:
                dev.disconnect()
            except Exception:
                pass
            registry.remove_instance(dev.info.device_id)
            if engine is not None:
                engine.remove_device(dev.info.device_id)

        for dev_cfg in profile.devices:
            dev_type = dev_cfg.get("device_type")
            dev_id   = dev_cfg.get("device_id")
            if not dev_id:
                continue
            factory = _DEVICE_FACTORIES.get(dev_type)
            if not factory:
                print(f"[Profile] Unknown device type: {dev_type}")
                continue
            try:
                # "name" is a display label, not a constructor arg — every device
                # factory builds its own default name internally, so apply it
                # after construction instead of passing it through.
                custom_name = dev_cfg.get("name")
                kwargs = {k: v for k, v in dev_cfg.items()
                          if k not in ("device_type", "name")}
                dev = factory(**kwargs)
                if custom_name:
                    dev.info.name = custom_name
                registry.add_instance(dev)
                dev.connect()
                if engine is not None:
                    engine.add_device(dev)
            except Exception as e:
                print(f"[Profile] Could not restore device {dev_id}: {e}")

        # ── Channel overrides ────────────────────────────────────────────
        for ch_data in profile.channels:
            dev = registry.get_instance(ch_data["device_id"])
            if not dev:
                continue
            ch = dev.get_channel(ch_data["channel_id"])
            if not ch:
                continue
            ch.name    = ch_data.get("name", ch.name)
            ch.unit    = ch_data.get("unit", ch.unit)
            ch.enabled = ch_data.get("enabled", ch.enabled)
            if ch_data.get("color"):
                ch.color = ch_data["color"]

        # ── Signal pipelines ─────────────────────────────────────────────
        for pd in profile.pipelines:
            filters = [filter_from_dict(fd) for fd in pd.get("filters", [])]
            pipeline = ChannelPipeline(
                device_id=pd["device_id"],
                channel_id=pd["channel_id"],
                filters=filters,
                enabled=pd.get("enabled", True),
            )
            processor.set_pipeline(pipeline)

        # ── Derived channels ─────────────────────────────────────────────
        for dd in profile.derived:
            sources = [tuple(s) for s in dd.get("sources", [])]
            dc = DerivedChannel(
                channel_id=dd["channel_id"],
                name=dd.get("name", dd["channel_id"]),
                unit=dd.get("unit", ""),
                color=dd.get("color", "#f72585"),
                kind=dd.get("kind", "expression"),
                sources=sources,
                source_names=dd.get("source_names", []),
                expression=dd.get("expression", ""),
                script=dd.get("script", ""),
                params=dd.get("params", {}),
                enabled=dd.get("enabled", True),
            )
            processor.add_derived(dc)

        # ── Controls ─────────────────────────────────────────────────────
        specs = [ControlSpec.from_dict(d) for d in profile.controls]
        control_panel.load_specs(specs)

        # ── Script vars ──────────────────────────────────────────────────
        if profile.script_vars:
            with processor._lock:
                processor._script_vars.update(profile.script_vars)

        # ── Settings ─────────────────────────────────────────────────────
        settings_ref.update(profile.settings)

        # ── Plot layout ───────────────────────────────────────────────────
        plot_cfg = None
        if profile.plot:
            try:
                from ui.windows.plot_window import LayoutConfig
                plot_cfg = LayoutConfig.from_json(json.dumps(profile.plot))
            except Exception as e:
                print(f"[Profile] Could not restore plot layout: {e}")

        # Plugin state
        if plugin_manager is not None and profile.plugin_state:
            for plugin_id, state in profile.plugin_state.items():
                plugin = plugin_manager.get_plugin(plugin_id)
                if plugin is not None:
                    try:
                        plugin.apply_save_state(state)
                    except Exception as e:
                        print(f"[Profile] Plugin '{plugin_id}' state restore error: {e}")

        return plot_cfg   # caller applies this to the chart