"""
core/signal_processor.py

Signal processing pipeline engine.

Provides two capabilities:

1. FILTERS — applied to a raw channel buffer before plotting:
     LowPass, HighPass, MovingAverage, Median, Derivative, Integral, Scale+Offset

2. DERIVED CHANNELS — virtual channels computed from one or more physical
   channels. Each derived channel runs a user-defined function every time
   new data arrives. Built-ins: derivative/second_derivative from displacement,
   power from voltage+current, RMS, etc.  Custom: arbitrary Python snippet.

Architecture
------------
   SignalProcessor sits between AcquisitionEngine and StripChartWidget.
   engine.new_data  →  SignalProcessor.process(dev, ch, t, val)
                    →  emits processed_data(virtual_or_real_id, ch_id, t, val)

The processor maintains its own ring buffers for derived channels so the
strip chart can query history just like physical channels.
"""

from __future__ import annotations

import math
import re
import threading
import traceback
from collections import deque
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Tuple, Any

from PyQt6.QtCore import QObject, pyqtSignal

# ── Constants ──────────────────────────────────────────────────────────────

MAX_BUF = 20_000


# ══════════════════════════════════════════════════════════════════════════════
#  Filter definitions
# ══════════════════════════════════════════════════════════════════════════════

class FilterBase:
    """All filters implement __call__(value: float) -> float."""
    name:    str = "identity"
    params:  dict = {}

    def __call__(self, value: float) -> float:
        return value

    def reset(self): pass

    def to_dict(self) -> dict:
        return {"type": self.name, **self.params}


class MovingAverageFilter(FilterBase):
    name = "moving_average"
    def __init__(self, window: int = 10):
        self.params = {"window": window}
        self._buf   = deque(maxlen=window)

    def __call__(self, v: float) -> float:
        self._buf.append(v)
        return sum(self._buf) / len(self._buf)

    def reset(self): self._buf.clear()


class MedianFilter(FilterBase):
    name = "median"
    def __init__(self, window: int = 5):
        self.params = {"window": window}
        self._buf   = deque(maxlen=window)

    def __call__(self, v: float) -> float:
        self._buf.append(v)
        s = sorted(self._buf)
        n = len(s)
        return s[n // 2] if n % 2 else (s[n//2 - 1] + s[n//2]) / 2

    def reset(self): self._buf.clear()


class LowPassFilter(FilterBase):
    """Exponential moving average (single-pole IIR low-pass)."""
    name = "low_pass"
    def __init__(self, alpha: float = 0.1):
        """alpha=0.0 → no change, 1.0 → unfiltered."""
        self.params = {"alpha": alpha}
        self._prev  = None

    def __call__(self, v: float) -> float:
        if self._prev is None:
            self._prev = v
        self._prev = self._prev + self.params["alpha"] * (v - self._prev)
        return self._prev

    def reset(self): self._prev = None


class HighPassFilter(FilterBase):
    """Simple single-pole IIR high-pass (compliment of low-pass)."""
    name = "high_pass"
    def __init__(self, alpha: float = 0.9):
        self.params  = {"alpha": alpha}
        self._prev_v = None
        self._prev_y = 0.0

    def __call__(self, v: float) -> float:
        if self._prev_v is None:
            self._prev_v = v
        y = self.params["alpha"] * (self._prev_y + v - self._prev_v)
        self._prev_y = y
        self._prev_v = v
        return y

    def reset(self): self._prev_v = None; self._prev_y = 0.0


class ScaleOffsetFilter(FilterBase):
    """y = scale * x + offset  (unit conversion, calibration)."""
    name = "scale_offset"
    def __init__(self, scale: float = 1.0, offset: float = 0.0):
        self.params = {"scale": scale, "offset": offset}

    def __call__(self, v: float) -> float:
        return self.params["scale"] * v + self.params["offset"]


class DerivativeFilter(FilterBase):
    """Numerical first derivative  dy/dt."""
    name = "derivative"
    def __init__(self): self.params = {}; self._prev_v = None; self._prev_t = None

    def process_with_t(self, v: float, t: float) -> float:
        if self._prev_t is None or t == self._prev_t:
            self._prev_v = v; self._prev_t = t; return 0.0
        dy = (v - self._prev_v) / (t - self._prev_t)
        self._prev_v = v; self._prev_t = t
        return dy

    def __call__(self, v: float) -> float:
        return 0.0   # use process_with_t for real output

    def reset(self): self._prev_v = None; self._prev_t = None


class IntegralFilter(FilterBase):
    """Numerical integration (trapezoidal rule)."""
    name = "integral"
    def __init__(self): self.params = {}; self._sum = 0.0; self._prev_v = None; self._prev_t = None

    def process_with_t(self, v: float, t: float) -> float:
        if self._prev_t is not None and t != self._prev_t:
            self._sum += 0.5 * (v + self._prev_v) * (t - self._prev_t)
        self._prev_v = v; self._prev_t = t
        return self._sum

    def __call__(self, v: float) -> float:
        return self._sum

    def reset(self): self._sum = 0.0; self._prev_v = None; self._prev_t = None


FILTER_CLASSES = {
    "moving_average": MovingAverageFilter,
    "median":         MedianFilter,
    "low_pass":       LowPassFilter,
    "high_pass":      HighPassFilter,
    "scale_offset":   ScaleOffsetFilter,
    "derivative":     DerivativeFilter,
    "integral":       IntegralFilter,
}


def filter_from_dict(d: dict) -> FilterBase:
    cls = FILTER_CLASSES.get(d.get("type", ""))
    if cls is None:
        return FilterBase()
    params = {k: v for k, v in d.items() if k != "type"}
    return cls(**params)


def register_filter_class(name: str, cls) -> None:
    """Register a plugin-provided filter class by type name."""
    FILTER_CLASSES[name] = cls


def unregister_filter_class(name: str) -> None:
    """Remove a previously registered plugin filter class."""
    FILTER_CLASSES.pop(name, None)


# ══════════════════════════════════════════════════════════════════════════════
#  Helpers
# ══════════════════════════════════════════════════════════════════════════════

def _make_src_names(sources: List[Tuple[str, str]]) -> List[str]:
    """
    Build a list of Python-safe variable names from (device_id, channel_id) pairs.
    Uses channel_id alone when unambiguous; prefixes device_id when two sources
    share the same channel_id from different devices.
    """
    seen: Dict[str, str] = {}  # ch_id -> dev_id of first occurrence
    names = []
    for dev_id, ch_id in sources:
        if ch_id in seen and seen[ch_id] != dev_id:
            raw = f"{dev_id}_{ch_id}"
        else:
            seen[ch_id] = dev_id
            raw = ch_id
        safe = re.sub(r"\W", "_", raw)
        if safe and safe[0].isdigit():
            safe = "_" + safe
        names.append(safe or "_ch")
    return names


# ══════════════════════════════════════════════════════════════════════════════
#  Derived channel definitions
# ══════════════════════════════════════════════════════════════════════════════

@dataclass
class DerivedChannel:
    """
    A virtual channel computed from one or more physical channels.

    kind options:
      "derivative"        — derivative of a source (dy/dt)
      "second_derivative" — second derivative of a source (d²y/dt²)
      "power"          — voltage_source * current_source
      "rms"            — rolling RMS of a source (window samples)
      "expression"     — arbitrary Python expression string
      "function"       — multi-line Python function body (def compute(...))
      "custom_script"  — full Python script, must define compute(inputs, t)
    """
    channel_id:   str                           # virtual ID, e.g. "vel_0"
    name:         str                           # display name
    unit:         str          = ""
    color:        str          = "#f72585"
    kind:         str          = "expression"   # see above
    # Source channel references  [("dev_id", "ch_id"), ...]
    sources:      List[Tuple[str, str]] = field(default_factory=list)
    # For built-in kinds
    params:       Dict[str, Any] = field(default_factory=dict)
    # For expression / function / custom_script
    expression:   str          = ""             # single-line: "x[0] * 2"
    script:       str          = ""             # multi-line function body
    enabled:      bool         = True
    # Runtime: compiled callable and per-channel script state (not serialised)
    _fn:          Optional[Callable] = field(default=None, repr=False, compare=False)
    _exec_state:  dict               = field(default_factory=dict, repr=False, compare=False)

    def compile(self) -> Optional[str]:
        """
        Compile expression/script into self._fn.
        Returns None on success, or error string on failure.

        Namespace available in all expression/script kinds:
          x        — list of source values in declaration order
          <name>   — each source channel_id as a named variable (x[0] == first source name)
          t        — elapsed time in seconds
          math     — Python math module
          vars     — shared SignalProcessor._script_vars dict (read/write)
          state    — per-channel persistent dict (survives between evaluations)
        """
        try:
            src_names = _make_src_names(self.sources)

            if self.kind == "expression":
                code = compile(f"__result__ = {self.expression}", "<expr>", "exec")
                _state = self._exec_state

                def _expr_fn(inputs, t, sv, _code=code, _names=src_names, _st=_state):
                    ns = {"x": inputs, "t": t, "math": math, "vars": sv, "state": _st}
                    for i, name in enumerate(_names):
                        if i < len(inputs):
                            ns[name] = inputs[i]
                    exec(_code, ns)
                    return float(ns["__result__"])

                self._fn = _expr_fn

            elif self.kind in ("function", "custom_script"):
                src = self.script
                if not src.strip().startswith("def compute"):
                    src = "def compute(x, t):\n" + "\n".join(
                        "    " + ln for ln in src.splitlines()
                    )
                _exec_ns: dict = {"math": math}
                exec(compile(src, "<script>", "exec"), _exec_ns)
                _fn_ref = _exec_ns["compute"]
                _state  = self._exec_state

                def _script_fn(inputs, t, sv,
                               _f=_fn_ref, _ns=_exec_ns, _st=_state, _names=src_names):
                    # refresh mutable globals before each call so scripts see live values
                    _ns["vars"]  = sv
                    _ns["state"] = _st
                    for i, name in enumerate(_names):
                        if i < len(inputs):
                            _ns[name] = inputs[i]
                    return float(_f(inputs, t))

                self._fn = _script_fn

            else:
                # Built-in kinds handled in SignalProcessor._compute_derived
                self._fn = None

            return None
        except Exception as e:
            self._fn = None
            return str(e)


# ══════════════════════════════════════════════════════════════════════════════
#  ChannelPipeline — per-channel filter stack
# ══════════════════════════════════════════════════════════════════════════════

@dataclass
class ChannelPipeline:
    device_id:  str
    channel_id: str
    filters:    List[FilterBase] = field(default_factory=list)
    enabled:    bool = True

    def process(self, value: float, timestamp: float) -> float:
        if not self.enabled:
            return value
        v = value
        for f in self.filters:
            if isinstance(f, (DerivativeFilter, IntegralFilter)):
                v = f.process_with_t(v, timestamp)
            else:
                v = f(v)
        return v

    def reset(self):
        for f in self.filters:
            f.reset()


# ══════════════════════════════════════════════════════════════════════════════
#  SignalProcessor
# ══════════════════════════════════════════════════════════════════════════════

class SignalProcessor(QObject):
    """
    Sits between AcquisitionEngine and StripChartWidget.

    Applies filter pipelines to raw channel data, then evaluates all
    derived channels and emits processed_data for everything.

    Connect:  engine.new_data  →  processor.on_raw_data
    Connect:  processor.processed_data  →  chart.on_new_data
    """

    processed_data = pyqtSignal(str, str, float, float)
    # device_id, channel_id, timestamp, value
    # For derived channels: device_id = "derived", channel_id = derived.channel_id

    derived_error  = pyqtSignal(str, str)   # channel_id, error_message

    def __init__(self):
        super().__init__()
        self._pipelines: Dict[Tuple[str, str], ChannelPipeline] = {}
        self._derived:   List[DerivedChannel] = []
        self._lock       = threading.Lock()

        # Latest raw values cache for derived evaluation
        # (dev_id, ch_id) -> (timestamp, value)
        self._latest: Dict[Tuple[str, str], Tuple[float, float]] = {}

        # Ring buffers for derived channels (so strip chart can query history)
        self._derived_bufs: Dict[str, Tuple[deque, deque]] = {}

        # Shared variable store — accessible as `vars` in all expressions/scripts
        self._script_vars: Dict[str, Any] = {}

    # ── Pipeline management ───────────────────────────────────────────────

    def set_pipeline(self, pipeline: ChannelPipeline):
        with self._lock:
            self._pipelines[(pipeline.device_id, pipeline.channel_id)] = pipeline

    def remove_pipeline(self, device_id: str, channel_id: str):
        with self._lock:
            self._pipelines.pop((device_id, channel_id), None)

    def get_pipeline(self, device_id: str, channel_id: str) -> Optional[ChannelPipeline]:
        return self._pipelines.get((device_id, channel_id))

    # ── Derived channel management ────────────────────────────────────────

    def add_derived(self, dc: DerivedChannel) -> Optional[str]:
        """Add a derived channel. Returns compile error string or None."""
        err = dc.compile()
        if err:
            return err
        with self._lock:
            self._derived = [d for d in self._derived if d.channel_id != dc.channel_id]
            self._derived.append(dc)
            self._derived_bufs[dc.channel_id] = (deque(maxlen=MAX_BUF), deque(maxlen=MAX_BUF))
        return None

    def remove_derived(self, channel_id: str):
        with self._lock:
            self._derived = [d for d in self._derived if d.channel_id != channel_id]
            self._derived_bufs.pop(channel_id, None)

    def get_derived(self) -> List[DerivedChannel]:
        return list(self._derived)

    def get_derived_buffer(self, channel_id: str):
        """Returns (times_deque, values_deque) or None."""
        return self._derived_bufs.get(channel_id)

    def all_virtual_channel_ids(self) -> List[str]:
        return list(self._derived_bufs.keys())

    def clear_history(self):
        """Clear derived channel buffers and latest-value cache."""
        with self._lock:
            for t_buf, v_buf in self._derived_bufs.values():
                t_buf.clear()
                v_buf.clear()
            self._latest.clear()

    # ── Shared script variables ───────────────────────────────────────────

    def set_var(self, name: str, value: Any) -> None:
        """Set a named variable accessible as `vars[name]` in all scripts."""
        with self._lock:
            self._script_vars[name] = value

    def get_var(self, name: str, default: Any = 0.0) -> Any:
        with self._lock:
            return self._script_vars.get(name, default)

    # ── Main data path ────────────────────────────────────────────────────

    def on_raw_data(self, device_id: str, channel_id: str,
                    timestamp: float, value: float):
        """Slot: receive raw data, apply filters, emit processed, update derived."""
        key = (device_id, channel_id)

        # Apply filter pipeline
        with self._lock:
            pipeline = self._pipelines.get(key)
        processed = pipeline.process(value, timestamp) if pipeline else value

        # Cache latest processed value
        with self._lock:
            self._latest[key] = (timestamp, processed)

        # Emit processed physical channel
        self.processed_data.emit(device_id, channel_id, timestamp, processed)

        # Evaluate all derived channels whose sources include this channel
        self._evaluate_derived(timestamp)

    def _evaluate_derived(self, timestamp: float):
        with self._lock:
            derived = list(self._derived)
            latest  = dict(self._latest)

        for dc in derived:
            if not dc.enabled:
                continue
            # Check all sources have recent data
            inputs = []
            for src in dc.sources:
                entry = latest.get(src)
                if entry is None:
                    break
                inputs.append(entry[1])   # value only
            else:
                # All sources present
                try:
                    result = self._compute_derived(dc, inputs, timestamp, latest)
                    if result is not None:
                        # Write back into local snapshot so downstream derived
                        # channels that depend on this one can see it this pass
                        latest[("derived", dc.channel_id)] = (timestamp, result)
                        with self._lock:
                            self._latest[("derived", dc.channel_id)] = (timestamp, result)
                        bufs = self._derived_bufs.get(dc.channel_id)
                        if bufs:
                            bufs[0].append(timestamp)
                            bufs[1].append(result)
                        self.processed_data.emit("derived", dc.channel_id,
                                                  timestamp, result)
                except Exception as e:
                    self.derived_error.emit(dc.channel_id,
                                            traceback.format_exc(limit=3))

    def _compute_derived(self, dc: DerivedChannel, inputs: List[float],
                          t: float, latest: dict) -> Optional[float]:
        if dc.kind in ("expression", "function", "custom_script"):
            if dc._fn is None:
                return None
            with self._lock:
                sv = self._script_vars
            return dc._fn(inputs, t, sv)

        elif dc.kind == "derivative":
            if not dc.sources: return None
            state = dc.params.setdefault("_state", {})
            prev_v = state.get("v")
            state["v"] = inputs[0]
            if dc.params.get("deriv_wrt") == "channel":
                if len(inputs) < 2: return None
                prev_x = state.get("x")
                state["x"] = inputs[1]
                if prev_x is None or inputs[1] == prev_x: return 0.0
                return (inputs[0] - prev_v) / (inputs[1] - prev_x)
            else:
                prev_t = state.get("t")
                state["t"] = t
                if prev_t is None or t == prev_t: return 0.0
                return (inputs[0] - prev_v) / (t - prev_t)

        elif dc.kind == "second_derivative":
            if not dc.sources: return None
            state = dc.params.setdefault("_state", {})
            prev_v = state.get("v")
            if dc.params.get("deriv_wrt") == "channel":
                if len(inputs) < 2: return None
                prev_x = state.get("x")
                cur_vel = 0.0
                if prev_x is not None and inputs[1] != prev_x:
                    cur_vel = (inputs[0] - prev_v) / (inputs[1] - prev_x)
                prev_vel = state.get("vel", 0.0)
                state["x"] = inputs[1]; state["v"] = inputs[0]; state["vel"] = cur_vel
                if prev_x is None or inputs[1] == prev_x: return 0.0
                return (cur_vel - prev_vel) / (inputs[1] - prev_x)
            else:
                prev_t = state.get("t")
                cur_vel = 0.0
                if prev_t is not None and t != prev_t:
                    cur_vel = (inputs[0] - prev_v) / (t - prev_t)
                prev_vel = state.get("vel", 0.0)
                state["t"] = t; state["v"] = inputs[0]; state["vel"] = cur_vel
                if prev_t is None or t == prev_t: return 0.0
                return (cur_vel - prev_vel) / (t - prev_t)

        elif dc.kind == "power":
            # V * I
            if len(inputs) < 2: return None
            return inputs[0] * inputs[1]

        elif dc.kind == "rms":
            window = dc.params.get("window", 20)
            buf = dc.params.setdefault("_buf", deque(maxlen=window))
            buf.append(inputs[0])
            return math.sqrt(sum(x*x for x in buf) / len(buf))

        elif dc.kind == "difference":
            if len(inputs) < 2: return None
            return inputs[0] - inputs[1]

        elif dc.kind == "sum":
            return sum(inputs)

        return None