tools/plotter.py

from __future__ import annotations

import io
import tempfile

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


def plot_function(
    expression: str,
    x_range: tuple[float, float] = (-10, 10),
    title: str = "",
    xlabel: str = "x",
    ylabel: str = "y",
) -> str:
    """Plot a mathematical function and return the path to the saved image.

    Args:
        expression: A numpy-compatible expression string using 'x' as variable.
            Examples: "np.sin(x)", "x**2 - 3*x + 2", "np.exp(-x**2)"
        x_range: Tuple of (x_min, x_max).
        title: Plot title.
        xlabel: X-axis label.
        ylabel: Y-axis label.

    Returns:
        Path to the saved PNG image.
    """
    x = np.linspace(x_range[0], x_range[1], 500)
    try:
        y = eval(expression, {"x": x, "np": np, "pi": np.pi, "e": np.e})
    except Exception as e:
        return _error_plot(f"Cannot evaluate: {expression}\n{e}")

    fig, ax = plt.subplots(figsize=(8, 5))
    ax.plot(x, y, "b-", linewidth=2)
    ax.set_title(title or f"$y = {expression}$")
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.grid(True, alpha=0.3)
    ax.axhline(y=0, color="k", linewidth=0.5)
    ax.axvline(x=0, color="k", linewidth=0.5)

    return _save_fig(fig)


def plot_multiple(
    expressions: list[dict],
    x_range: tuple[float, float] = (-10, 10),
    title: str = "",
) -> str:
    """Plot multiple functions on the same axes.

    Args:
        expressions: List of {"expr": str, "label": str} dicts.
        x_range: Tuple of (x_min, x_max).
        title: Plot title.

    Returns:
        Path to the saved PNG image.
    """
    x = np.linspace(x_range[0], x_range[1], 500)
    fig, ax = plt.subplots(figsize=(8, 5))

    for item in expressions:
        expr = item["expr"]
        label = item.get("label", expr)
        try:
            y = eval(expr, {"x": x, "np": np, "pi": np.pi, "e": np.e})
            ax.plot(x, y, linewidth=2, label=label)
        except Exception:
            continue

    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    ax.axhline(y=0, color="k", linewidth=0.5)
    ax.axvline(x=0, color="k", linewidth=0.5)
    ax.legend()

    return _save_fig(fig)


def _save_fig(fig) -> str:
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    fig.savefig(tmp.name, dpi=100, bbox_inches="tight")
    plt.close(fig)
    return tmp.name


def _error_plot(message: str) -> str:
    fig, ax = plt.subplots(figsize=(8, 5))
    ax.text(0.5, 0.5, message, ha="center", va="center", fontsize=12, color="red")
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis("off")
    return _save_fig(fig)