import numpy as np
import cv2

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

class Turtle:
    def __init__(self, canvas_size=(2000, 2000)):
        self.x = canvas_size[0] // 2
        self.y = canvas_size[1] // 2
        self.heading = 0
        self.canvas = np.ones((canvas_size[1], canvas_size[0], 3), dtype=np.uint8) * 255
        self.is_down = True

    def forward(self, dist):
        dist = dist * SCALE
        x0, y0 = self.x, self.y
        x1 = int(x0 + dist * np.cos(self.heading))
        y1 = int(y0 - dist * np.sin(self.heading))
        if self.is_down:
            cv2.line(self.canvas, (x0, y0), (x1, y1), (0, 0, 0), 3)
        self.x, self.y = x1, 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):
        if path:
            cv2.imwrite(path, self.canvas)
        return self.canvas

    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("test2.png")

    def plot2():
        for j in range(2):
            forward(2)
            left(0.0)
            for i in range(4):
                forward(2)
                left(90)
            forward(0)
            left(180.0)
            forward(2)
            left(180.0)
        FINAL_IMAGE = turtle.save("")

    example_plot()
    # plot2()