File size: 3,448 Bytes
be3d62e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
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")