import re
from PIL import Image
from matplotlib import pyplot as plt
import matplotlib as mpl
import matplotlib.style as mplstyle
import numpy as np

mpl.use('Agg')
mpl.rcParams["path.simplify_threshold"] = 0.0
mpl.rcParams['agg.path.chunksize'] = 10000
mplstyle.use('fast')

HALF_INF = 63
INF = 126
EPS_DIST = 1/20
EPS_ANGLE = 2.86

class Turtle:
    def __init__(self, ax=None):
        self.x = 0
        self.y = 0
        self.heading = 0
        if ax is None:
            self.fig, self.ax = plt.subplots(1, 1, figsize=(20, 20))
        else:
            self.ax = ax
        self.ax.set_xlim(-50, 50)
        self.ax.set_ylim(-50, 50)
        self.ax.set_xticklabels([])
        self.ax.set_yticklabels([])
        self.ax.set_xticks([])
        self.ax.set_yticks([])
        self.ax.spines['top'].set_visible(False)
        self.ax.spines['bottom'].set_visible(False)
        self.ax.spines['left'].set_visible(False)
        self.ax.spines['right'].set_visible(False)
        self.is_down = True

    def forward(self, dist):
        x0, y0 = self.x, self.y
        x1 = x0 + dist * np.cos(self.heading)
        y1 = y0 + dist * np.sin(self.heading)
        if self.is_down:
            self.ax.plot([x0, x1], [y0, y1], color='black', linewidth=3)
        self.x = x1
        self.y = y1

    def left(self, angle):
        self.heading += angle * np.pi / 180

    def right(self, angle):
        self.heading -= angle * np.pi / 180

    def penup(self):
        self.is_down = False

    def pendown(self):
        self.is_down = True

    def save(self, path):
        self.fig.canvas.draw()
        pil_img = Image.frombytes('RGB', self.fig.canvas.get_width_height(), self.fig.canvas.tostring_rgb())
        if path:
            pil_img.save(path)
        return pil_img

    class _TurtleState:
        def __init__(self, turtle):
            self.turtle = turtle
            self.position = None
            self.heading = None
            self.pen_status = None

        def __enter__(self):
            self.position = (self.turtle.x, self.turtle.y)
            self.heading = self.turtle.heading
            self.pen_status = self.turtle.is_down
            return self

        def __exit__(self, exc_type, exc_val, exc_tb):
            self.turtle.penup()
            self.turtle.x, self.turtle.y = self.position
            self.turtle.heading = self.heading
            if self.pen_status:
                self.turtle.pendown()

if __name__ == "__main__":
    turtle = Turtle()

    def forward(dist):
        turtle.forward(dist)

    def left(angle):
        turtle.left(angle)

    def right(angle):
        turtle.right(angle)

    def penup():
        turtle.penup()

    def pendown():
        turtle.pendown()

    def save(path):
        turtle.save(path)

    def fork_state():
        """
        Clone the current state of the turtle.

        Usage:
        with clone_state():
            forward(100)
            left(90)
            forward(100)
        """
        return turtle._TurtleState(turtle)

    # Example usage
    def example_plot():
        forward(5)

        with fork_state():
            forward(10)
            left(90)
            forward(10)
            with fork_state():
                right(90)
                forward(20)
                left(90)
                forward(10)
            left(90)
            forward(10)

        right(90)
        forward(50)
        save("test.png")