"""Math expression transform.

Allows users to create derived channels from arbitrary mathematical
expressions involving other channels. Uses a safe evaluation approach
with numpy functions.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any

import numpy as np

from impakt.channel.code import ChannelCode
from impakt.channel.model import Channel


# Allowed numpy functions in expressions
_SAFE_NUMPY = {
    "abs": np.abs,
    "sqrt": np.sqrt,
    "square": np.square,
    "sin": np.sin,
    "cos": np.cos,
    "tan": np.tan,
    "arcsin": np.arcsin,
    "arccos": np.arccos,
    "arctan": np.arctan,
    "arctan2": np.arctan2,
    "exp": np.exp,
    "log": np.log,
    "log10": np.log10,
    "log2": np.log2,
    "max": np.maximum,
    "min": np.minimum,
    "clip": np.clip,
    "sign": np.sign,
    "pi": np.pi,
    "e": np.e,
    "where": np.where,
    "zeros_like": np.zeros_like,
    "ones_like": np.ones_like,
    "cumsum": np.cumsum,
    "diff": np.diff,
    "gradient": np.gradient,
    "interp": np.interp,
    "mean": np.mean,
    "std": np.std,
    "trapz": np.trapezoid if hasattr(np, "trapezoid") else np.trapz,
}


def _validate_expression(expression: str) -> None:
    """Basic validation of a math expression.

    Rejects potentially dangerous constructs.
    """
    # Reject obvious dangers
    forbidden = [
        "import",
        "__",
        "exec",
        "eval",
        "open",
        "file",
        "os.",
        "sys.",
        "subprocess",
        "compile",
        "globals",
        "locals",
        "getattr",
        "setattr",
        "delattr",
        "breakpoint",
    ]
    lower = expression.lower()
    for word in forbidden:
        if word in lower:
            raise ValueError(f"Forbidden construct in expression: {word}")


@dataclass(frozen=True)
class MathExpr:
    """Math expression transform.

    Creates a new channel by evaluating a mathematical expression
    involving named channel variables.

    Example:
        MathExpr(
            expression="0.6 * chest_x + 0.4 * chest_z",
            channel_names=["chest_x", "chest_z"],
        )
    """

    expression: str
    channel_names: tuple[str, ...] = ()

    @property
    def name(self) -> str:
        return "math_expr"

    @property
    def params(self) -> dict[str, Any]:
        return {
            "expression": self.expression,
            "channel_names": list(self.channel_names),
        }

    def apply(self, channel: Channel) -> Channel:
        """Not meaningful for single-channel — use evaluate() instead."""
        raise NotImplementedError(
            "MathExpr operates on multiple channels. Use evaluate() with a channel dict."
        )

    def evaluate(
        self,
        channels: dict[str, Channel],
        result_name: str = "derived",
        result_unit: str = "",
    ) -> Channel:
        """Evaluate the expression with the given channels.

        Args:
            channels: Dict mapping variable names to Channel objects.
            result_name: Name for the output channel.
            result_unit: Unit for the output channel.

        Returns:
            New Channel with the computed data.
        """
        _validate_expression(self.expression)

        # Build namespace
        namespace: dict[str, Any] = dict(_SAFE_NUMPY)
        ref_channel: Channel | None = None

        for var_name, ch in channels.items():
            namespace[var_name] = ch.data
            if ref_channel is None:
                ref_channel = ch

        if ref_channel is None:
            raise ValueError("At least one channel is required")

        # Evaluate
        try:
            result = eval(self.expression, {"__builtins__": {}}, namespace)  # noqa: S307
        except Exception as e:
            raise ValueError(f"Error evaluating expression '{self.expression}': {e}") from e

        if not isinstance(result, np.ndarray):
            result = np.full_like(ref_channel.data, float(result))

        return Channel(
            name=result_name,
            code=ChannelCode.parse(result_name),
            data=result.astype(np.float64),
            time=ref_channel.time,
            unit=result_unit,
            sample_rate=ref_channel.sample_rate,
            metadata={},
            source_test_id=ref_channel.source_test_id,
            transform_history=(f"math_expr({self.expression})",),
        )


def math_expr(
    expression: str,
    channels: dict[str, Channel],
    name: str = "derived",
    unit: str = "",
) -> Channel:
    """Convenience function to evaluate a math expression.

    Args:
        expression: Math expression using channel variable names.
        channels: Dict mapping variable names used in expression to Channels.
        name: Name for the resulting channel.
        unit: Unit for the resulting channel.

    Returns:
        New Channel with computed data.

    Example:
        result = math_expr(
            "sqrt(ax**2 + az**2)",
            {"ax": head_accel_x, "az": head_accel_z},
            name="Head Accel XZ",
            unit="g",
        )
    """
    return MathExpr(
        expression=expression,
        channel_names=tuple(channels.keys()),
    ).evaluate(channels, result_name=name, result_unit=unit)
