"""
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: velocity/acceleration 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 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)


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

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

    kind options:
      "velocity"       — derivative of a displacement source
      "acceleration"   — second derivative of a displacement source
      "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 (not serialised)
    _fn:          Optional[Callable] = field(default=None, repr=False, compare=False)

    def compile(self) -> Optional[str]:
        """
        Compile expression/script into self._fn.
        Returns None on success, or error string on failure.
        """
        try:
            if self.kind == "expression":
                # Single-line: inputs are x (list of latest values), t (time)
                code = compile(f"__result__ = {self.expression}", "<expr>", "exec")
                def _expr_fn(inputs, t, _code=code):
                    ns = {"x": inputs, "t": t, "math": math}
                    exec(_code, ns)
                    return float(ns["__result__"])
                self._fn = _expr_fn

            elif self.kind in ("function", "custom_script"):
                # User provides a def compute(x, t): ... body
                # We wrap it in a module namespace
                src = self.script
                if not src.strip().startswith("def compute"):
                    src = "def compute(x, t):\n" + "\n".join(
                        "    " + ln for ln in src.splitlines()
                    )
                ns: dict = {"math": math}
                exec(compile(src, "<script>", "exec"), ns)
                fn = ns["compute"]
                self._fn = lambda inputs, t, _f=fn: float(_f(inputs, t))

            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]] = {}

    # ── 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())

    # ── 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:
                        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
            return dc._fn(inputs, t)

        elif dc.kind == "velocity":
            # derivative of source[0]
            src = dc.sources[0] if dc.sources else None
            if src is None: return None
            state = dc.params.setdefault("_state", {})
            prev_t = state.get("t"); prev_v = state.get("v")
            state["t"] = t; state["v"] = inputs[0]
            if prev_t is None or t == prev_t: return 0.0
            return (inputs[0] - prev_v) / (t - prev_t)

        elif dc.kind == "acceleration":
            # second derivative  — derivative of velocity
            src = dc.sources[0] if dc.sources else None
            if src is None: return None
            state = dc.params.setdefault("_state", {})
            # First get velocity
            prev_t = state.get("t"); prev_v = state.get("v")
            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