add support for x^2+y^2=0,4 and alike

Signed-off-by: salmonik <raisersoft@proton.me>

1 files changed, 61 insertions, 0 deletions

diff --git a/cogs/plot.py b/cogs/plot.py
index 65dcae4..4fcde44 100644
--- a/cogs/plot.py
+++ b/cogs/plot.py
@@ -39,9 +39,22 @@ def parse_expression(expr: str) -> str:
     expr = expr.replace("^", "**")
     expr = re.sub(r"(\d)([x])", r"\1*\2", expr)
     expr = re.sub(r"([x])(\d)", r"\1*\2", expr)
+    expr = re.sub(r"(\d)([y])", r"\1*\2", expr)
+    expr = re.sub(r"([y])(\d)", r"\1*\2", expr)
     return expr
 
 
+def is_implicit(expr: str) -> bool:
+    return "=" in expr and "y" in expr and "x" in expr
+
+
+def parse_implicit(expr: str) -> tuple[str, str]:
+    lhs, rhs = expr.split("=", 1)
+    lhs = parse_expression(lhs.strip())
+    rhs = parse_expression(rhs.strip())
+    return lhs, rhs
+
+
 class Plot(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
@@ -53,6 +66,11 @@ class Plot(commands.Cog):
             raise RuntimeError("matplotlib not available")
 
         plt = _plt_module
+        
+        if is_implicit(func_expr):
+            self._generate_implicit_plot(func_expr, output_path)
+            return
+
         parsed_expr = parse_expression(func_expr)
 
         x = [i * 0.1 for i in range(-100, 101)]
@@ -96,6 +114,49 @@ class Plot(commands.Cog):
         plt.savefig(output_path, dpi=100, bbox_inches="tight")
         plt.close()
 
+    def _generate_implicit_plot(self, func_expr: str, output_path: str):
+        global _plt_module
+        plt = _plt_module
+        
+        lhs, rhs = parse_implicit(func_expr)
+        
+        x_range = 20
+        y_range = 20
+        
+        x_vals = []
+        y_vals = []
+        
+        for xi in range(-x_range * 10, x_range * 10 + 1):
+            x = xi / 10.0
+            for yi in range(-y_range * 10, y_range * 10 + 1):
+                y = yi / 10.0
+                try:
+                    left_val = eval(lhs, {"x": x, "y": y, "math": math, "sqrt": math.sqrt, "pi": math.pi, "e": math.e})
+                    right_val = eval(rhs, {"x": x, "y": y, "math": math, "sqrt": math.sqrt, "pi": math.pi, "e": math.e})
+                    diff = abs(left_val - right_val)
+                    if math.isfinite(diff) and diff < 0.3:
+                        x_vals.append(x)
+                        y_vals.append(y)
+                except Exception:
+                    pass
+        
+        fig, ax = plt.subplots(figsize=(10, 10))
+        ax.scatter(x_vals, y_vals, s=1, color="#00ff00", alpha=0.5)
+        ax.set_title(f"{func_expr}", fontsize=14)
+        ax.set_xlabel("x", fontsize=12)
+        ax.set_ylabel("y", fontsize=12)
+        ax.grid(True, alpha=0.3)
+        ax.axhline(y=0, color="white", linewidth=0.5)
+        ax.axvline(x=0, color="white", linewidth=0.5)
+        ax.set_xlim(-x_range, x_range)
+        ax.set_ylim(-y_range, y_range)
+        ax.set_aspect("equal")
+        ax.set_facecolor("#1a1a1a")
+        fig.patch.set_facecolor("#1a1a1a")
+        
+        plt.savefig(output_path, dpi=100, bbox_inches="tight")
+        plt.close()
+
     @commands.slash_command(name="funkcja", description="Generuje wykres funkcji matematycznej")
     async def plot(self, ctx, funkcja: str = discord.Option(str, "Funkcja np. x^3, sin(x), x^2 + 2*x - 1", required=True)):
         if not _ensure_plot():