diff --git a/scripts/GameField.gd b/scripts/GameField.gd
index 430a8fc..46b2d62 100644
--- a/scripts/GameField.gd
+++ b/scripts/GameField.gd
@@ -235,7 +235,7 @@ func _input(event: InputEvent) -> void:
 					sel_col = int(gridpos.x)
 					drag_action = DragAction.DRAW_COLOR
 					drag_color = posmod(cell_colors.get_cell(sel_row, sel_col) + (1 if event.button_index==BUTTON_LEFT else -1), N_COLORS)
-					cell_colors.set_cell(sel_row, sel_col, drag_color)
+					set_cell_color(sel_row, sel_col, drag_color)
 
 		elif event is InputEventMouseMotion and event.button_mask:
 			if drag_action == DragAction.TEST_LINE:
@@ -268,99 +268,134 @@ func _input(event: InputEvent) -> void:
 					sel_row = int(gridpos.y)
 					sel_col = int(gridpos.x)
 					set_cell_color(sel_row, sel_col, drag_color)
+	elif event is InputEventKey and event.pressed:
+		match event.scancode:
+			KEY_Z:
+				undo()
 
 var undo_stack = []
 
 class UndoAction:
 	var action_type
-	var cell_row
-	var cell_col
+	var cell_row: int
+	var cell_col: int
+	var activated: bool
 	func _init(action, row, col) -> void:
 		action_type = action
 		cell_row = row
 		cell_col = col
+		activated = false
 
 class UndoActionCorner extends UndoAction:
-	func _init(action, row, col).(action, row, col) -> void:
-		pass
+	var corner_marks
+	var corner_mark
+	var set: bool  # False = clear
+	func _init(row, col, corner_mark, set, corner_marks).(corner_mark, row, col) -> void:
+		self.corner_marks = corner_marks
+		self.corner_mark = corner_mark
+		self.set = set
 
 	func _to_string() -> String:
-		match action_type:
-			UndoActions.SET_L_RIGHT:
-				return '+L %s%d:%d'%[Common.num2alpha(cell_col), cell_row-1, cell_row]
-			UndoActions.CLEAR_L_RIGHT:
-				return '-L %s%d:%d'%[Common.num2alpha(cell_col), cell_row-1, cell_row]
-			UndoActions.SET_X_RIGHT:
-				return '+X %s%d:%d'%[Common.num2alpha(cell_col), cell_row-1, cell_row]
-			UndoActions.CLEAR_X_RIGHT:
-				return '-X %s%d:%d'%[Common.num2alpha(cell_col), cell_row-1, cell_row]
-			UndoActions.SET_L_DOWN:
-				return '+L %s%d:%s'%[Common.num2alpha(cell_col-1), cell_row, Common.num2alpha(cell_col)]
-			UndoActions.CLEAR_L_DOWN:
-				return '-L %s%d:%s'%[Common.num2alpha(cell_col-1), cell_row, Common.num2alpha(cell_col)]
-			UndoActions.SET_X_DOWN:
-				return '+X %s%d:%s'%[Common.num2alpha(cell_col-1), cell_row, Common.num2alpha(cell_col)]
-			UndoActions.CLEAR_X_DOWN:
-				return '-X %s%d:%s'%[Common.num2alpha(cell_col-1), cell_row, Common.num2alpha(cell_col)]
+		var prefix = '+' if set else '-'
+		match corner_mark:
+			CornerMark.L_RIGHT:
+				return prefix+'L %s%d:%d'%[Common.num2alpha(cell_col), cell_row-1, cell_row]
+			CornerMark.X_RIGHT:
+				return prefix+'X %s%d:%d'%[Common.num2alpha(cell_col), cell_row-1, cell_row]
+			CornerMark.L_DOWN:
+				return prefix+'L %s%d:%s'%[Common.num2alpha(cell_col-1), cell_row, Common.num2alpha(cell_col)]
+			CornerMark.X_DOWN:
+				return prefix+'X %s%d:%s'%[Common.num2alpha(cell_col-1), cell_row, Common.num2alpha(cell_col)]
 			_:
-				return 'unk %d %s%d'%[action_type, Common.num2alpha(cell_col), cell_row]
+				return 'unk %s%d %s%d'%[prefix, corner_mark, Common.num2alpha(cell_col), cell_row]
+
+	func activate() -> void:
+		activated = true
+		if set:
+			corner_marks.set_flag(cell_row, cell_col, corner_mark)
+		else:
+			corner_marks.clear_flag(cell_row, cell_col, corner_mark)
+
+	func undo() -> void:
+		activated = false
+		if set:
+			corner_marks.clear_flag(cell_row, cell_col, corner_mark)
+		else:
+			corner_marks.set_flag(cell_row, cell_col, corner_mark)
 
 class UndoActionCell extends UndoAction:
 	var set_state
-	var prev_state
-	func _init(action, row, col, set_state, prev_state).(action, row, col) -> void:
+	var old_state
+	var cell_colors
+	func _init(action, row, col, set_state, cell_colors).(action, row, col) -> void:
 		self.set_state = set_state
-		self.prev_state = prev_state
+		self.old_state = set_state
+		self.cell_colors = cell_colors
 
 	func _to_string() -> String:
 		match action_type:
 			UndoActions.SET_COLOR:
-				return 'C%d->%d %s%d'%[set_state, prev_state, Common.num2alpha(cell_col), cell_row]
+				return 'C%d<-%d %s%d'%[set_state, old_state, Common.num2alpha(cell_col), cell_row]
 			_:
 				return 'unk %d %s%d'%[action_type, Common.num2alpha(cell_col), cell_row]
 
-func set_cell_color(row, col, new_color):
-	var old_color = cell_colors.get_cell(row, col)
-	if old_color != new_color:
-		cell_colors.set_cell(row, col, new_color)
-		undo_stack.append(UndoActionCell.new(UndoActions.SET_COLOR, row, col, new_color, old_color))
+	func activate() -> void:
+		assert(not activated)
+		activated = true
+		match action_type:
+			UndoActions.SET_COLOR:
+				old_state = cell_colors.get_cell(cell_row, cell_col)
+				cell_colors.set_cell(cell_row, cell_col, set_state)
 
+	func undo() -> void:
+		assert(activated)
+		activated = false
+		match action_type:
+			UndoActions.SET_COLOR:
+				cell_colors.set_cell(cell_row, cell_col, old_state)
+
+func do_action(action):
+	action.activate()
+	undo_stack.append(action)
+
+func undo():
+	if !undo_stack.empty():
+		var action = undo_stack.pop_back()
+		action.undo()
+
+func set_cell_color(row, col, new_color):
+	if new_color != cell_colors.get_cell(row, col):
+		do_action(UndoActionCell.new(UndoActions.SET_COLOR, row, col, new_color, cell_colors))
+
+# TODO: Refactor all this nonsense a bit more
 func set_line_right(row, col):
 	if !corner_marks.get_flag(row, col, CornerMark.L_RIGHT) and !corner_marks.get_flag(row, col, CornerMark.X_RIGHT):
-		corner_marks.set_flag(row, col, CornerMark.L_RIGHT)
-		undo_stack.append(UndoActionCorner.new(UndoActions.SET_L_RIGHT, row, col))
+		do_action(UndoActionCorner.new(row, col, CornerMark.L_RIGHT, true, corner_marks))
 
 func set_line_down(row, col):
 	if !corner_marks.get_flag(row, col, CornerMark.L_DOWN) and !corner_marks.get_flag(row, col, CornerMark.X_DOWN):
-		corner_marks.set_flag(row, col, CornerMark.L_DOWN)
-		undo_stack.append(UndoActionCorner.new(UndoActions.SET_L_DOWN, row, col))
+		do_action(UndoActionCorner.new(row, col, CornerMark.L_DOWN, true, corner_marks))
 
 func clear_line_right(row, col):
 	if corner_marks.get_flag(row, col, CornerMark.L_RIGHT):
-		corner_marks.clear_flag(row, col, CornerMark.L_RIGHT)
-		undo_stack.append(UndoActionCorner.new(UndoActions.CLEAR_L_RIGHT, row, col))
+		do_action(UndoActionCorner.new(row, col, CornerMark.L_RIGHT, false, corner_marks))
 
 func clear_line_down(row, col):
 	if corner_marks.get_flag(row, col, CornerMark.L_DOWN):
-		corner_marks.clear_flag(row, col, CornerMark.L_DOWN)
-		undo_stack.append(UndoActionCorner.new(UndoActions.CLEAR_L_DOWN, row, col))
+		do_action(UndoActionCorner.new(row, col, CornerMark.L_DOWN, false, corner_marks))
 
 func set_x_right(row, col):
 	if !corner_marks.get_flag(row, col, CornerMark.L_RIGHT) and !corner_marks.get_flag(row, col, CornerMark.X_RIGHT):
-		corner_marks.set_flag(row, col, CornerMark.X_RIGHT)
-		undo_stack.append(UndoActionCorner.new(UndoActions.SET_X_RIGHT, row, col))
+		do_action(UndoActionCorner.new(row, col, CornerMark.X_RIGHT, true, corner_marks))
 
 func set_x_down(row, col):
 	if !corner_marks.get_flag(row, col, CornerMark.L_DOWN) and !corner_marks.get_flag(row, col, CornerMark.X_DOWN):
-		corner_marks.set_flag(row, col, CornerMark.X_DOWN)
-		undo_stack.append(UndoActionCorner.new(UndoActions.SET_X_DOWN, row, col))
+		do_action(UndoActionCorner.new(row, col, CornerMark.X_DOWN, true, corner_marks))
 
 func clear_x_right(row, col):
 	if corner_marks.get_flag(row, col, CornerMark.X_RIGHT):
-		corner_marks.clear_flag(row, col, CornerMark.X_RIGHT)
-		undo_stack.append(UndoActionCorner.new(UndoActions.CLEAR_X_RIGHT, row, col))
+		do_action(UndoActionCorner.new(row, col, CornerMark.X_RIGHT, false, corner_marks))
 
 func clear_x_down(row, col):
 	if corner_marks.get_flag(row, col, CornerMark.X_DOWN):
-		corner_marks.clear_flag(row, col, CornerMark.X_DOWN)
-		undo_stack.append(UndoActionCorner.new(UndoActions.CLEAR_X_DOWN, row, col))
+		do_action(UndoActionCorner.new(row, col, CornerMark.X_DOWN, false, corner_marks))