controlai_agent/tools/plotting.py

"""General mathematical function and signal plotting tools."""

from __future__ import annotations

import hashlib
import re
import time
from pathlib import Path
from typing import Any

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from controlai_agent.registry import registry

PLOTS_DIR = Path("outputs/plots")
PLOTS_DIR.mkdir(parents=True, exist_ok=True)


SAFE_MATH_ENV: dict[str, Any] = {
    "sin": np.sin,
    "cos": np.cos,
    "tan": np.tan,
    "sinc": np.sinc,
    "arcsin": np.arcsin,
    "asin": np.arcsin,
    "arccos": np.arccos,
    "acos": np.arccos,
    "arctan": np.arctan,
    "atan": np.arctan,
    "arctan2": np.arctan2,
    "atan2": np.arctan2,
    "exp": np.exp,
    "log": np.log,
    "log10": np.log10,
    "sqrt": np.sqrt,
    "abs": np.abs,
    "pi": np.pi,
    "e": np.e,
    "sinh": np.sinh,
    "cosh": np.cosh,
    "tanh": np.tanh,
    "heaviside": lambda x: np.heaviside(x, 1.0),
    "step": lambda x: (x >= 0).astype(float),
    "sign": np.sign,
    "rad2deg": np.rad2deg,
    "deg2rad": np.deg2rad,
}


def _expr_to_mathtext(expr: str) -> str:
    """Best-effort conversion of a Python-style expression into valid matplotlib mathtext.

    Model-generated expressions use Python syntax (`**` for power, bare `*` for
    multiplication) which mathtext does not understand -- it renders the raw
    asterisks literally (e.g. "t * *2 - 3 * t") instead of a clean formula.
    """
    text = expr.replace("**", "^")
    text = re.sub(r"\^\(([^()]+)\)", r"^{\1}", text)
    text = re.sub(r"\^([a-zA-Z0-9_.]{2,})", r"^{\1}", text)
    text = text.replace("*", r" \cdot ")
    return text


@registry.register(
    name="plot_math_expression",
    description="Plot any mathematical function, curve, or signal f(t) or f(x) (e.g. 'sin(t)', 'exp(-t)*cos(2*t)', 'sin(x)', 't**2 - 3*t') and save the plot figure.",
    parameters_schema={
        "type": "object",
        "properties": {
            "expression": {
                "type": "string",
                "description": "Mathematical expression in terms of variable 't' or 'x', e.g. 'sin(t)', 'sin(x)', 'exp(-0.5*t)*sin(3*t)'.",
            },
            "t_start": {
                "type": "number",
                "description": "Start of range (default -10.0 for x, or 0.0 for t).",
            },
            "t_end": {
                "type": "number",
                "description": "End of range (default 10.0).",
            },
            "title": {
                "type": "string",
                "description": "Optional title for the plot.",
            },
            "xlabel": {
                "type": "string",
                "description": "Optional label for horizontal axis (default 't' or 'x').",
            },
            "ylabel": {
                "type": "string",
                "description": "Optional label for vertical axis (default 'f(t)' or 'y').",
            },
        },
        "required": ["expression"],
    },
)
def plot_math_expression(
    expression: str,
    t_start: float | None = None,
    t_end: float | None = None,
    title: str | None = None,
    xlabel: str | None = None,
    ylabel: str | None = None,
) -> dict[str, Any]:
    # Normalize expression (replace ^ with **)
    clean_expr = expression.replace("^", "**")
    
    # Detect independent variable (t or x)
    var_name = "x" if "x" in clean_expr and "t" not in clean_expr else "t"
    
    # Default range
    if t_start is None:
        t_start = -2 * np.pi if var_name == "x" else 0.0
    if t_end is None:
        t_end = 2 * np.pi if var_name == "x" else 10.0

    t_vals = np.linspace(t_start, t_end, 600)

    # Evaluate safely
    eval_env = dict(SAFE_MATH_ENV)
    for v in ["x", "t", "alpha", "theta", "phi", "omega", "u", "y", "s", "rad"]:
        eval_env[v] = t_vals

    try:
        y_vals = eval(clean_expr, {"__builtins__": {}}, eval_env)
        # Handle scalar constant output
        if np.isscalar(y_vals):
            y_vals = np.full_like(t_vals, float(y_vals))
        else:
            y_vals = np.asarray(y_vals)
    except Exception as exc:
        return {
            "status": "error",
            "error": f"Failed to evaluate expression '{expression}': {exc}",
        }

    # A complex result means the expression was not a real-valued signal --
    # most often a Laplace/transfer-function expression containing `1j` that
    # was mistakenly passed to a time-domain plotter. Silently casting it to
    # float discards the imaginary part and yields a meaningless curve, so
    # reject it and tell the model what to do instead.
    if np.iscomplexobj(y_vals):
        return {
            "status": "error",
            "error": (
                f"Expression '{expression}' evaluates to complex values, so it is not a real "
                "time-domain signal that can be plotted. Do not pass transfer functions or "
                "expressions containing the imaginary unit here -- to simulate a system's response "
                "use simulate_step_response (transfer function) or simulate_state_feedback_response "
                "(state-space with optional gain K)."
            ),
        }

    y_vals = np.asarray(y_vals, dtype=float)
    if not np.any(np.isfinite(y_vals)):
        return {
            "status": "error",
            "error": f"Expression '{expression}' produced no finite values over the range [{t_start}, {t_end}].",
        }

    # Plot styling
    fig, ax = plt.subplots(figsize=(8, 4.2), dpi=120)
    fig.patch.set_facecolor("#151b23")
    ax.set_facecolor("#0f1217")

    line_color = "#58a6ff"
    mathtext_expr = _expr_to_mathtext(expression)
    (line,) = ax.plot(t_vals, y_vals, color=line_color, linewidth=2, label=f"${mathtext_expr}$")
    ax.axhline(0, color="#484f58", linestyle="--", linewidth=0.8, alpha=0.7)
    ax.axvline(0, color="#484f58", linestyle="--", linewidth=0.8, alpha=0.7)

    plot_title = title or f"Plot of $f({var_name}) = {mathtext_expr}$"
    plot_xlabel = xlabel or var_name
    plot_ylabel = ylabel or f"f({var_name})"

    title_obj = ax.set_title(plot_title, color="#f0f6fc", fontsize=12, pad=10, fontweight="bold")
    ax.set_xlabel(plot_xlabel, color="#8b949e", fontsize=10)
    ax.set_ylabel(plot_ylabel, color="#8b949e", fontsize=10)
    ax.tick_params(colors="#8b949e", labelsize=9)
    ax.grid(True, linestyle=":", alpha=0.3, color="#8b949e")

    for spine in ax.spines.values():
        spine.set_color("#30363d")

    ax.legend(loc="best", facecolor="#151b23", edgecolor="#30363d", labelcolor="#f0f6fc", fontsize=9)
    plt.tight_layout()

    # Save unique plot
    hash_str = hashlib.md5(f"{clean_expr}_{t_start}_{t_end}_{time.time()}".encode()).hexdigest()[:8]
    out_file = PLOTS_DIR / f"plot_{hash_str}.png"
    try:
        plt.savefig(str(out_file), facecolor=fig.get_facecolor(), edgecolor="none", dpi=120)
    except Exception:
        # Mathtext couldn't parse the expression (rare, malformed LaTeX-ish input)
        # -- fall back to plain, non-math text labels rather than losing the plot.
        line.set_label(expression)
        title_obj.set_text(title or f"Plot of f({var_name}) = {expression}")
        ax.legend(loc="best", facecolor="#151b23", edgecolor="#30363d", labelcolor="#f0f6fc", fontsize=9)
        plt.savefig(str(out_file), facecolor=fig.get_facecolor(), edgecolor="none", dpi=120)
    plt.close(fig)

    return {
        "status": "success",
        "plot_path": str(out_file),
        "expression": expression,
        "range": [float(t_start), float(t_end)],
    }