diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4ee72fb..ddfacf5 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -14,6 +14,8 @@ add_executable(masssprings
     src/renderer.cpp
     src/mass_springs.cpp
     src/klotski.cpp
+    src/state.cpp
+    src/input.cpp
 )
 
 target_include_directories(masssprings PUBLIC ${RAYLIB_CPP_INCLUDE_DIR})
diff --git a/include/input.hpp b/include/input.hpp
new file mode 100644
index 0000000..15055b3
--- /dev/null
+++ b/include/input.hpp
@@ -0,0 +1,43 @@
+#ifndef __INPUT_HPP_
+#define __INPUT_HPP_
+
+#include "renderer.hpp"
+#include "state.hpp"
+
+class InputHandler {
+public:
+  StateManager &state;
+  Renderer &renderer;
+
+  int hov_x;
+  int hov_y;
+  int sel_x;
+  int sel_y;
+
+  bool has_block_add_xy = false;
+  int block_add_x = -1;
+  int block_add_y = -1;
+
+public:
+  InputHandler(StateManager &state, Renderer &renderer)
+      : state(state), renderer(renderer), hov_x(-1), hov_y(-1), sel_x(-1),
+        sel_y(-1), has_block_add_xy(false), block_add_x(-1), block_add_y(-1) {}
+
+  InputHandler(const InputHandler &copy) = delete;
+  InputHandler &operator=(const InputHandler &copy) = delete;
+  InputHandler(InputHandler &&move) = delete;
+  InputHandler &operator=(InputHandler &&move) = delete;
+
+  ~InputHandler() {}
+
+public:
+  auto HandleMouseHover() -> void;
+
+  auto HandleMouse() -> void;
+
+  auto HandleKeys() -> void;
+
+  auto HandleInput() -> void;
+};
+
+#endif
diff --git a/include/mass_springs.hpp b/include/mass_springs.hpp
index adf1803..964adab 100644
--- a/include/mass_springs.hpp
+++ b/include/mass_springs.hpp
@@ -1,6 +1,8 @@
 #ifndef __MASS_SPRINGS_HPP_
 #define __MASS_SPRINGS_HPP_
 
+#include "config.hpp"
+#include "klotski.hpp"
 #include <raylib.h>
 #include <raymath.h>
 #include <string>
@@ -97,7 +99,7 @@ public:
   std::unordered_map<std::string, Spring> springs;
 
 public:
-  MassSpringSystem() : last_build(1000) {};
+  MassSpringSystem() : last_build(REPULSION_GRID_REFRESH) {};
 
   MassSpringSystem(const MassSpringSystem &copy) = delete;
   MassSpringSystem &operator=(const MassSpringSystem &copy) = delete;
@@ -110,14 +112,13 @@ private:
   auto BuildGrid() -> void;
 
 public:
-  auto AddMass(float mass, Vector3 position, bool fixed,
-               const std::string &state) -> void;
+  auto AddMass(float mass, Vector3 position, bool fixed, const State &state)
+      -> void;
 
-  auto GetMass(const std::string &state) -> Mass &;
+  auto GetMass(const State &state) -> Mass &;
 
-  auto AddSpring(const std::string &massA, const std::string &massB,
-                 float spring_constant, float dampening_constant,
-                 float rest_length) -> void;
+  auto AddSpring(const State &massA, const State &massB, float spring_constant,
+                 float dampening_constant, float rest_length) -> void;
 
   auto Clear() -> void;
 
@@ -130,6 +131,8 @@ public:
   auto EulerUpdate(float delta_time) -> void;
 
   auto VerletUpdate(float delta_time) -> void;
+
+  auto InvalidateGrid() -> void;
 };
 
 #endif
diff --git a/include/states.hpp b/include/presets.hpp
similarity index 97%
rename from include/states.hpp
rename to include/presets.hpp
index 9c74c0f..f805563 100644
--- a/include/states.hpp
+++ b/include/presets.hpp
@@ -1,5 +1,5 @@
-#ifndef __STATES_HPP_
-#define __STATES_HPP_
+#ifndef __PRESETS_HPP_
+#define __PRESETS_HPP_
 
 #include <functional>
 #include <vector>
@@ -200,13 +200,13 @@ inline auto state_new_century_wc(const State &state) -> bool {
   return state.state == "F4x5:21......1121..12bb..12........12111111..";
 }
 
-std::vector<StateGenerator> generators{
+static std::vector<StateGenerator> generators{
     state_simple_1r,  state_simple_2r,  state_simple_3r,  state_complex_1r,
     state_complex_2r, state_complex_3r, state_complex_4f, state_complex_5r,
     state_complex_6r, state_klotski,    state_century,    state_super_century,
     state_new_century};
 
-std::vector<WinCondition> win_conditions{
+static std::vector<WinCondition> win_conditions{
     state_simple_1r_wc,  state_simple_2r_wc,  state_simple_3r_wc,
     state_complex_1r_wc, state_complex_2r_wc, state_complex_3r_wc,
     state_complex_4f_wc, state_complex_5r_wc, state_complex_6r_wc,
diff --git a/include/renderer.hpp b/include/renderer.hpp
index aee53c5..40826ef 100644
--- a/include/renderer.hpp
+++ b/include/renderer.hpp
@@ -48,7 +48,6 @@ private:
   RenderTexture render_target;
   RenderTexture klotski_target;
   RenderTexture menu_target;
-  std::unordered_set<State> winning_states;
 
 public:
   bool mark_solutions;
@@ -77,26 +76,22 @@ public:
   }
 
 public:
-  auto UpdateWinningStates(const MassSpringSystem &masssprings,
-                           const WinCondition win_condition) -> void;
-
-  auto AddWinningState(const State &state, const WinCondition win_condition)
-      -> void;
-
   auto UpdateCamera(const MassSpringSystem &masssprings, const State &current)
       -> void;
 
   auto UpdateTextureSizes() -> void;
 
   auto DrawMassSprings(const MassSpringSystem &masssprings,
-                       const State &current) -> void;
+                       const State &current,
+                       const std::unordered_set<State> &winning_states) -> void;
 
   auto DrawKlotski(const State &state, int hov_x, int hov_y, int sel_x,
                    int sel_y, int block_add_x, int block_add_y,
                    const WinCondition win_condition) -> void;
 
   auto DrawMenu(const MassSpringSystem &masssprings, int current_preset,
-                const State &current_state) -> void;
+                const State &current_state,
+                const std::unordered_set<State> &winning_states) -> void;
 
   auto DrawTextures() -> void;
 };
diff --git a/include/state.hpp b/include/state.hpp
new file mode 100644
index 0000000..fdd1358
--- /dev/null
+++ b/include/state.hpp
@@ -0,0 +1,58 @@
+#ifndef __STATE_HPP_
+#define __STATE_HPP_
+
+#include "config.hpp"
+#include "klotski.hpp"
+#include "mass_springs.hpp"
+#include "presets.hpp"
+
+#include <raymath.h>
+
+class StateManager {
+public:
+  MassSpringSystem &mass_springs;
+
+  int current_preset;
+  State current_state;
+  State previous_state;
+
+  bool edited = false;
+
+  std::unordered_set<State> winning_states;
+
+public:
+  StateManager(MassSpringSystem &mass_springs)
+      : mass_springs(mass_springs), current_preset(0),
+        current_state(generators[current_preset]()),
+        previous_state(current_state), edited(false) {
+    mass_springs.AddMass(MASS, Vector3Zero(), false, current_state);
+  }
+
+  StateManager(const StateManager &copy) = delete;
+  StateManager &operator=(const StateManager &copy) = delete;
+  StateManager(StateManager &&move) = delete;
+  StateManager &operator=(StateManager &&move) = delete;
+
+  ~StateManager() {}
+
+public:
+  auto LoadPreset(int preset) -> void;
+
+  auto ResetState() -> void;
+
+  auto PreviousPreset() -> void;
+
+  auto NextPreset() -> void;
+
+  auto FillGraph() -> void;
+
+  auto UpdateGraph() -> void;
+
+  auto ClearGraph() -> void;
+
+  auto FindWinningStates() -> void;
+
+  auto CurrentWinCondition() -> WinCondition;
+};
+
+#endif
diff --git a/src/input.cpp b/src/input.cpp
new file mode 100644
index 0000000..3c3db2e
--- /dev/null
+++ b/src/input.cpp
@@ -0,0 +1,169 @@
+#include <algorithm>
+#include <raylib.h>
+
+#include "config.hpp"
+#include "input.hpp"
+
+auto InputHandler::HandleMouseHover() -> void {
+  const int board_width = GetScreenWidth() / 2.0 - 2 * BOARD_PADDING;
+  const int board_height = GetScreenHeight() - MENU_HEIGHT - 2 * BOARD_PADDING;
+  int block_size = std::min(board_width / state.current_state.width,
+                            board_height / state.current_state.height) -
+                   2 * BLOCK_PADDING;
+  int x_offset = (board_width - (block_size + 2 * BLOCK_PADDING) *
+                                    state.current_state.width) /
+                 2.0;
+  int y_offset = (board_height - (block_size + 2 * BLOCK_PADDING) *
+                                     state.current_state.height) /
+                 2.0;
+
+  Vector2 m = GetMousePosition();
+  if (m.x < x_offset) {
+    hov_x = 100;
+  } else {
+    hov_x = (m.x - x_offset) / (block_size + 2 * BLOCK_PADDING);
+  }
+  if (m.y - MENU_HEIGHT < y_offset) {
+    hov_y = 100;
+  } else {
+    hov_y = (m.y - MENU_HEIGHT - y_offset) / (block_size + 2 * BLOCK_PADDING);
+  }
+}
+
+auto InputHandler::HandleMouse() -> void {
+  if (IsMouseButtonPressed(MOUSE_BUTTON_LEFT)) {
+    // If we clicked a block...
+    if (state.current_state.GetBlock(hov_x, hov_y).IsValid()) {
+      sel_x = hov_x;
+      sel_y = hov_y;
+    }
+    // If we clicked empty space...
+    else {
+      // Select a position
+      if (!has_block_add_xy) {
+        if (hov_x >= 0 && hov_x < state.current_state.width && hov_y >= 0 &&
+            hov_y < state.current_state.height) {
+          block_add_x = hov_x;
+          block_add_y = hov_y;
+          has_block_add_xy = true;
+        }
+      }
+      // If we have already selected a position
+      else {
+        int block_add_width = hov_x - block_add_x + 1;
+        int block_add_height = hov_y - block_add_y + 1;
+        if (block_add_width <= 0 || block_add_height <= 0) {
+          block_add_x = -1;
+          block_add_y = -1;
+          has_block_add_xy = false;
+        } else if (block_add_x >= 0 &&
+                   block_add_x + block_add_width <= state.current_state.width &&
+                   block_add_y >= 0 &&
+                   block_add_y + block_add_height <=
+                       state.current_state.height) {
+          bool success = state.current_state.AddBlock(
+              Block(block_add_x, block_add_y, block_add_width, block_add_height,
+                    false));
+
+          if (success) {
+            block_add_x = -1;
+            block_add_y = -1;
+            has_block_add_xy = false;
+            state.ClearGraph();
+            state.edited = true;
+          }
+        }
+      }
+    }
+  } else if (IsMouseButtonPressed(MOUSE_BUTTON_RIGHT)) {
+    if (state.current_state.RemoveBlock(hov_x, hov_y)) {
+      state.ClearGraph();
+      state.edited = true;
+    } else if (has_block_add_xy) {
+      block_add_x = -1;
+      block_add_y = -1;
+      has_block_add_xy = false;
+    }
+  }
+}
+
+auto InputHandler::HandleKeys() -> void {
+  if (IsKeyPressed(KEY_W)) {
+    if (state.current_state.MoveBlockAt(sel_x, sel_y, Direction::NOR)) {
+      sel_y--;
+    }
+  } else if (IsKeyPressed(KEY_A)) {
+    if (state.current_state.MoveBlockAt(sel_x, sel_y, Direction::WES)) {
+      sel_x--;
+    }
+  } else if (IsKeyPressed(KEY_S)) {
+    if (state.current_state.MoveBlockAt(sel_x, sel_y, Direction::SOU)) {
+      sel_y++;
+    }
+  } else if (IsKeyPressed(KEY_D)) {
+    if (state.current_state.MoveBlockAt(sel_x, sel_y, Direction::EAS)) {
+      sel_x++;
+    }
+  } else if (IsKeyPressed(KEY_P)) {
+    std::cout << "State: " << state.current_state.state << std::endl;
+    Block sel = state.current_state.GetBlock(sel_x, sel_y);
+    int idx = state.current_state.GetIndex(sel.x, sel.y) - 5;
+    if (sel.IsValid()) {
+      std::cout << "Sel:   " << state.current_state.state.substr(0, 5)
+                << std::string(idx, '.') << sel.ToString()
+                << std::string(state.current_state.state.length() - idx - 7,
+                               '.')
+                << std::endl;
+    }
+  } else if (IsKeyPressed(KEY_N)) {
+    block_add_x = -1;
+    block_add_y = -1;
+    has_block_add_xy = false;
+    state.PreviousPreset();
+  } else if (IsKeyPressed(KEY_M)) {
+    block_add_x = -1;
+    block_add_y = -1;
+    has_block_add_xy = false;
+    state.NextPreset();
+  } else if (IsKeyPressed(KEY_R)) {
+    state.ResetState();
+  } else if (IsKeyPressed(KEY_G)) {
+    state.FillGraph();
+  } else if (IsKeyPressed(KEY_C)) {
+    state.ClearGraph();
+  } else if (IsKeyPressed(KEY_I)) {
+    renderer.mark_solutions = !renderer.mark_solutions;
+  } else if (IsKeyPressed(KEY_O)) {
+    renderer.connect_solutions = !renderer.connect_solutions;
+  } else if (IsKeyPressed(KEY_F)) {
+    state.current_state.ToggleRestricted();
+    state.ClearGraph();
+    state.edited = true;
+  } else if (IsKeyPressed(KEY_T)) {
+    state.current_state.ToggleTarget(sel_x, sel_y);
+    state.ClearGraph();
+    state.edited = true;
+  } else if (IsKeyPressed(KEY_LEFT) && state.current_state.width > 1) {
+    state.current_state = state.current_state.RemoveColumn();
+    state.ClearGraph();
+    state.edited = true;
+  } else if (IsKeyPressed(KEY_RIGHT) && state.current_state.width < 9) {
+    state.current_state = state.current_state.AddColumn();
+    state.ClearGraph();
+    state.edited = true;
+  } else if (IsKeyPressed(KEY_UP) && state.current_state.height > 1) {
+    state.current_state = state.current_state.RemoveRow();
+    state.ClearGraph();
+    state.edited = true;
+  } else if (IsKeyPressed(KEY_DOWN) && state.current_state.height < 9) {
+    state.current_state = state.current_state.AddRow();
+    state.ClearGraph();
+    state.edited = true;
+  }
+}
+
+auto InputHandler::HandleInput() -> void {
+  HandleMouseHover();
+  HandleMouse();
+  HandleKeys();
+}
diff --git a/src/main.cpp b/src/main.cpp
index 70532a1..cd6c3f3 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -3,10 +3,10 @@
 #include <raymath.h>
 
 #include "config.hpp"
-#include "klotski.hpp"
+#include "input.hpp"
 #include "mass_springs.hpp"
 #include "renderer.hpp"
-#include "states.hpp"
+#include "state.hpp"
 
 #ifndef WEB
 #include <omp.h>
@@ -26,54 +26,11 @@
 //       - Click states to display them in the board
 //       - Find shortest path to any winning state and mark it in the graph
 //         - Also mark the next move along the path on the board
-
-auto apply_state(MassSpringSystem &mass_springs, StateGenerator generator)
-    -> State {
-  mass_springs.springs.clear();
-  mass_springs.masses.clear();
-
-  State s = generator();
-  mass_springs.AddMass(MASS, Vector3Zero(), false, s.state);
-
-  return s;
-};
-
-auto populate_masssprings(MassSpringSystem &mass_springs,
-                          const State &current_state) -> void {
-  std::pair<std::unordered_set<std::string>,
-            std::vector<std::pair<std::string, std::string>>>
-      closure = current_state.Closure();
-  for (const auto &state : closure.first) {
-    Vector3 pos =
-        Vector3(static_cast<float>(GetRandomValue(-10000, 10000)) / 1000.0,
-                static_cast<float>(GetRandomValue(-10000, 10000)) / 1000.0,
-                static_cast<float>(GetRandomValue(-10000, 10000)) / 1000.0);
-
-    mass_springs.AddMass(MASS, pos, false, state);
-  }
-  for (const auto &[from, to] : closure.second) {
-    mass_springs.AddSpring(from, to, SPRING_CONSTANT, DAMPENING_CONSTANT,
-                           REST_LENGTH);
-  }
-  std::cout << "Inserted " << mass_springs.masses.size() << " masses and "
-            << mass_springs.springs.size() << " springs." << std::endl;
-  std::cout << "Consuming "
-            << sizeof(decltype(*mass_springs.masses.begin())) *
-                   mass_springs.masses.size()
-            << " Bytes for masses." << std::endl;
-  std::cout << "Consuming "
-            << sizeof(decltype(*mass_springs.springs.begin())) *
-                   mass_springs.springs.size()
-            << " Bytes for springs." << std::endl;
-}
-
-auto clear_masssprings(MassSpringSystem &masssprings,
-                       const State &current_state) -> std::string {
-  masssprings.masses.clear();
-  masssprings.springs.clear();
-  masssprings.AddMass(MASS, Vector3Zero(), false, current_state.state);
-  return current_state.state;
-}
+// TODO: Don't tie the simulation step resolution to the FPS (frametime)
+//       - This breaks the simulation on slower systems
+//       - Add a modifiable speed setting?
+//       - Clamp the frametime?
+//       - Use a fixed step size and control how often it runs per frame?
 
 auto main(int argc, char *argv[]) -> int {
   // if (argc < 2) {
@@ -85,34 +42,21 @@ auto main(int argc, char *argv[]) -> int {
   std::cout << "OpenMP: " << omp_get_max_threads() << " threads." << std::endl;
 #endif
 
+  // RayLib window setup
   SetTraceLogLevel(LOG_ERROR);
-
-  // SetTargetFPS(60);
   SetConfigFlags(FLAG_VSYNC_HINT);
   SetConfigFlags(FLAG_MSAA_4X_HINT);
   SetConfigFlags(FLAG_WINDOW_RESIZABLE);
   SetConfigFlags(FLAG_WINDOW_ALWAYS_RUN);
-
   InitWindow(INITIAL_WIDTH * 2, INITIAL_HEIGHT + MENU_HEIGHT, "MassSprings");
 
-  // Rendering configuration
+  // Game setup
   Renderer renderer;
-
-  // Klotski configuration
-  int current_preset = 0;
-  MassSpringSystem masssprings;
-  State current_state = apply_state(masssprings, generators[current_preset]);
+  MassSpringSystem mass_springs;
+  StateManager state = StateManager(mass_springs);
+  InputHandler input = InputHandler(state, renderer);
 
   // Game loop
-  float frametime;
-  bool edited = false;
-  bool has_block_add_xy = false;
-  int block_add_x = -1;
-  int block_add_y = -1;
-  int hov_x = 0;
-  int hov_y = 0;
-  int sel_x = 0;
-  int sel_y = 0;
 #ifdef PRINT_TIMINGS
   double last_print_time = GetTime();
   std::chrono::duration<double, std::milli> physics_time_accumulator =
@@ -122,186 +66,10 @@ auto main(int argc, char *argv[]) -> int {
   int time_measure_count = 0;
 #endif
   while (!WindowShouldClose()) {
-    frametime = GetFrameTime();
-    std::string previous_state = current_state.state;
-
-    // Mouse handling
-    const int board_width = GetScreenWidth() / 2.0 - 2 * BOARD_PADDING;
-    const int board_height =
-        GetScreenHeight() - MENU_HEIGHT - 2 * BOARD_PADDING;
-    int block_size = std::min(board_width / current_state.width,
-                              board_height / current_state.height) -
-                     2 * BLOCK_PADDING;
-    int x_offset =
-        (board_width - (block_size + 2 * BLOCK_PADDING) * current_state.width) /
-        2.0;
-    int y_offset = (board_height -
-                    (block_size + 2 * BLOCK_PADDING) * current_state.height) /
-                   2.0;
-    Vector2 m = GetMousePosition();
-    if (m.x < x_offset) {
-      hov_x = 100;
-    } else {
-      hov_x = (m.x - x_offset) / (block_size + 2 * BLOCK_PADDING);
-    }
-    if (m.y - MENU_HEIGHT < y_offset) {
-      hov_y = 100;
-    } else {
-      hov_y = (m.y - MENU_HEIGHT - y_offset) / (block_size + 2 * BLOCK_PADDING);
-    }
-    if (IsMouseButtonPressed(MOUSE_BUTTON_LEFT)) {
-      // If we clicked a block...
-      if (current_state.GetBlock(hov_x, hov_y).IsValid()) {
-        sel_x = hov_x;
-        sel_y = hov_y;
-      }
-      // If we clicked empty space...
-      else {
-        // Select a position
-        if (!has_block_add_xy) {
-          if (hov_x >= 0 && hov_x < current_state.width && hov_y >= 0 &&
-              hov_y < current_state.height) {
-            block_add_x = hov_x;
-            block_add_y = hov_y;
-            has_block_add_xy = true;
-          }
-        }
-        // If we have already selected a position
-        else {
-          int block_add_width = hov_x - block_add_x + 1;
-          int block_add_height = hov_y - block_add_y + 1;
-          if (block_add_width <= 0 || block_add_height <= 0) {
-            block_add_x = -1;
-            block_add_y = -1;
-            has_block_add_xy = false;
-          } else if (block_add_x >= 0 &&
-                     block_add_x + block_add_width <= current_state.width &&
-                     block_add_y >= 0 &&
-                     block_add_y + block_add_height <= current_state.height) {
-            bool success = current_state.AddBlock(
-                Block(block_add_x, block_add_y, block_add_width,
-                      block_add_height, false));
-
-            if (success) {
-              block_add_x = -1;
-              block_add_y = -1;
-              has_block_add_xy = false;
-              previous_state = clear_masssprings(masssprings, current_state);
-              edited = true;
-            }
-          }
-        }
-      }
-    } else if (IsMouseButtonPressed(MOUSE_BUTTON_RIGHT)) {
-      if (current_state.RemoveBlock(hov_x, hov_y)) {
-        previous_state = clear_masssprings(masssprings, current_state);
-        edited = true;
-      } else if (has_block_add_xy) {
-        block_add_x = -1;
-        block_add_y = -1;
-        has_block_add_xy = false;
-      }
-    }
-
-    // Key handling
-    if (IsKeyPressed(KEY_W)) {
-      if (current_state.MoveBlockAt(sel_x, sel_y, Direction::NOR)) {
-        sel_y--;
-      }
-    } else if (IsKeyPressed(KEY_A)) {
-      if (current_state.MoveBlockAt(sel_x, sel_y, Direction::WES)) {
-        sel_x--;
-      }
-    } else if (IsKeyPressed(KEY_S)) {
-      if (current_state.MoveBlockAt(sel_x, sel_y, Direction::SOU)) {
-        sel_y++;
-      }
-    } else if (IsKeyPressed(KEY_D)) {
-      if (current_state.MoveBlockAt(sel_x, sel_y, Direction::EAS)) {
-        sel_x++;
-      }
-    } else if (IsKeyPressed(KEY_P)) {
-      std::cout << "State: " << current_state.state << std::endl;
-      Block sel = current_state.GetBlock(sel_x, sel_y);
-      int idx = current_state.GetIndex(sel.x, sel.y) - 5;
-      if (sel.IsValid()) {
-        std::cout << "Sel:   " << current_state.state.substr(0, 5)
-                  << std::string(idx, '.') << sel.ToString()
-                  << std::string(current_state.state.length() - idx - 7, '.')
-                  << std::endl;
-      }
-    } else if (IsKeyPressed(KEY_N)) {
-      block_add_x = -1;
-      block_add_y = -1;
-      has_block_add_xy = false;
-      current_preset =
-          (generators.size() + current_preset - 1) % generators.size();
-      current_state = apply_state(masssprings, generators[current_preset]);
-      previous_state = current_state.state;
-      edited = false;
-    } else if (IsKeyPressed(KEY_M)) {
-      block_add_x = -1;
-      block_add_y = -1;
-      has_block_add_xy = false;
-      current_preset = (current_preset + 1) % generators.size();
-      current_state = apply_state(masssprings, generators[current_preset]);
-      previous_state = current_state.state;
-      edited = false;
-    } else if (IsKeyPressed(KEY_R)) {
-      current_state = generators[current_preset]();
-      if (edited) {
-        // We also need to clear the graph, in case the state has been edited.
-        // Then the graph would contain states that are impossible.
-        previous_state = clear_masssprings(masssprings, current_state);
-        edited = false;
-      }
-    } else if (IsKeyPressed(KEY_G)) {
-      previous_state = clear_masssprings(masssprings, current_state);
-      populate_masssprings(masssprings, current_state);
-      renderer.UpdateWinningStates(masssprings, win_conditions[current_preset]);
-    } else if (IsKeyPressed(KEY_C)) {
-      previous_state = clear_masssprings(masssprings, current_state);
-    } else if (IsKeyPressed(KEY_I)) {
-      renderer.mark_solutions = !renderer.mark_solutions;
-    } else if (IsKeyPressed(KEY_O)) {
-      renderer.connect_solutions = !renderer.connect_solutions;
-    } else if (IsKeyPressed(KEY_F)) {
-      current_state.ToggleRestricted();
-      previous_state = clear_masssprings(masssprings, current_state);
-      edited = true;
-    } else if (IsKeyPressed(KEY_T)) {
-      current_state.ToggleTarget(sel_x, sel_y);
-      previous_state = clear_masssprings(masssprings, current_state);
-      edited = true;
-    } else if (IsKeyPressed(KEY_LEFT) && current_state.width > 1) {
-      current_state = current_state.RemoveColumn();
-      previous_state = clear_masssprings(masssprings, current_state);
-      edited = true;
-    } else if (IsKeyPressed(KEY_RIGHT) && current_state.width < 9) {
-      current_state = current_state.AddColumn();
-      previous_state = clear_masssprings(masssprings, current_state);
-      edited = true;
-    } else if (IsKeyPressed(KEY_UP) && current_state.height > 1) {
-      current_state = current_state.RemoveRow();
-      previous_state = clear_masssprings(masssprings, current_state);
-      edited = true;
-    } else if (IsKeyPressed(KEY_DOWN) && current_state.height < 9) {
-      current_state = current_state.AddRow();
-      previous_state = clear_masssprings(masssprings, current_state);
-      edited = true;
-    }
-
-    if (previous_state != current_state.state) {
-      masssprings.AddMass(
-          MASS,
-          Vector3(static_cast<float>(GetRandomValue(-1000, 1000)) / 1000.0,
-                  static_cast<float>(GetRandomValue(-1000, 1000)) / 1000.0,
-                  static_cast<float>(GetRandomValue(-1000, 1000)) / 1000.0),
-          false, current_state.state);
-      masssprings.AddSpring(current_state.state, previous_state,
-                            SPRING_CONSTANT, DAMPENING_CONSTANT, REST_LENGTH);
-      renderer.AddWinningState(current_state, win_conditions[current_preset]);
-    }
+    // Input update
+    state.previous_state = state.current_state;
+    input.HandleInput();
+    state.UpdateGraph();
 
     // Physics update
 #ifdef PRINT_TIMINGS
@@ -309,13 +77,13 @@ auto main(int argc, char *argv[]) -> int {
         std::chrono::high_resolution_clock::now();
 #endif
     for (int i = 0; i < UPDATES_PER_FRAME; ++i) {
-      masssprings.ClearForces();
-      masssprings.CalculateSpringForces();
-      masssprings.CalculateRepulsionForces();
+      mass_springs.ClearForces();
+      mass_springs.CalculateSpringForces();
+      mass_springs.CalculateRepulsionForces();
 #ifdef VERLET_UPDATE
-      masssprings.VerletUpdate(frametime / UPDATES_PER_FRAME * SIM_SPEED);
+      mass_springs.VerletUpdate(GetFrameTime() / UPDATES_PER_FRAME * SIM_SPEED);
 #else
-      mass_springs.EulerUpdate(frametime * SIM_SPEED);
+      mass_springs.EulerUpdate(GetFrameTime() * SIM_SPEED);
 #endif
     }
 #ifdef PRINT_TIMINGS
@@ -329,12 +97,15 @@ auto main(int argc, char *argv[]) -> int {
     std::chrono::high_resolution_clock::time_point rs =
         std::chrono::high_resolution_clock::now();
 #endif
-    renderer.UpdateCamera(masssprings, current_state);
+    renderer.UpdateCamera(mass_springs, state.current_state);
     renderer.UpdateTextureSizes();
-    renderer.DrawMassSprings(masssprings, current_state);
-    renderer.DrawKlotski(current_state, hov_x, hov_y, sel_x, sel_y, block_add_x,
-                         block_add_y, win_conditions[current_preset]);
-    renderer.DrawMenu(masssprings, current_preset, current_state);
+    renderer.DrawMassSprings(mass_springs, state.current_state,
+                             state.winning_states);
+    renderer.DrawKlotski(state.current_state, input.hov_x, input.hov_y,
+                         input.sel_x, input.sel_y, input.block_add_x,
+                         input.block_add_y, state.CurrentWinCondition());
+    renderer.DrawMenu(mass_springs, state.current_preset, state.current_state,
+                      state.winning_states);
     renderer.DrawTextures();
 #ifdef PRINT_TIMINGS
     std::chrono::high_resolution_clock::time_point re =
diff --git a/src/mass_springs.cpp b/src/mass_springs.cpp
index 0512b71..017321c 100644
--- a/src/mass_springs.cpp
+++ b/src/mass_springs.cpp
@@ -77,26 +77,26 @@ auto Spring::CalculateSpringForce() const -> void {
 }
 
 auto MassSpringSystem::AddMass(float mass, Vector3 position, bool fixed,
-                               const std::string &state) -> void {
-  if (!masses.contains(state)) {
-    masses.insert(std::make_pair(state, Mass(mass, position, fixed)));
+                               const State &state) -> void {
+  if (!masses.contains(state.state)) {
+    masses.insert(std::make_pair(state.state, Mass(mass, position, fixed)));
   }
 }
 
-auto MassSpringSystem::GetMass(const std::string &state) -> Mass & {
-  return masses.at(state);
+auto MassSpringSystem::GetMass(const State &state) -> Mass & {
+  return masses.at(state.state);
 }
 
-auto MassSpringSystem::AddSpring(const std::string &massA,
-                                 const std::string &massB,
+auto MassSpringSystem::AddSpring(const State &massA, const State &massB,
                                  float spring_constant,
                                  float dampening_constant, float rest_length)
     -> void {
   std::string states;
-  if (std::hash<std::string>{}(massA) < std::hash<std::string>{}(massB)) {
-    states = std::format("{}{}", massA, massB);
+  if (std::hash<std::string>{}(massA.state) <
+      std::hash<std::string>{}(massB.state)) {
+    states = std::format("{}{}", massA.state, massB.state);
   } else {
-    states = std::format("{}{}", massB, massA);
+    states = std::format("{}{}", massB.state, massA.state);
   }
 
   if (!springs.contains(states)) {
@@ -109,6 +109,7 @@ auto MassSpringSystem::AddSpring(const std::string &massA,
 auto MassSpringSystem::Clear() -> void {
   masses.clear();
   springs.clear();
+  InvalidateGrid();
 }
 
 auto MassSpringSystem::ClearForces() -> void {
@@ -216,21 +217,6 @@ auto MassSpringSystem::CalculateRepulsionForces() -> void {
 
     mass->force = Vector3Add(mass->force, force);
   }
-
-  // Old method
-  // for (auto &[state, mass] : masses) {
-  //   for (auto &[s, m] : masses) {
-  //     Vector3 dx = Vector3Subtract(mass.position, m.position);
-  //
-  //     // This can be accelerated with a spatial data structure
-  //     if (Vector3Length(dx) >= 3 * REST_LENGTH) {
-  //       continue;
-  //     }
-  //
-  //     mass.force = Vector3Add(
-  //         mass.force, Vector3Scale(Vector3Normalize(dx), REPULSION_FORCE));
-  //   }
-  // }
 }
 
 auto MassSpringSystem::EulerUpdate(float delta_time) -> void {
@@ -245,3 +231,12 @@ auto MassSpringSystem::VerletUpdate(float delta_time) -> void {
     mass.VerletUpdate(delta_time);
   }
 }
+
+auto MassSpringSystem::InvalidateGrid() -> void {
+  mass_vec.clear();
+  indices.clear();
+  cell_ids.clear();
+  last_build = REPULSION_GRID_REFRESH;
+  last_masses_count = 0;
+  last_springs_count = 0;
+}
diff --git a/src/renderer.cpp b/src/renderer.cpp
index 5422057..5fc3e27 100644
--- a/src/renderer.cpp
+++ b/src/renderer.cpp
@@ -3,6 +3,7 @@
 #include <format>
 #include <raylib.h>
 #include <raymath.h>
+#include <unordered_set>
 
 #include "config.hpp"
 #include "klotski.hpp"
@@ -87,27 +88,6 @@ auto OrbitCamera3D::Update(const Mass &current_mass) -> void {
   camera.target = target;
 }
 
-auto Renderer::UpdateWinningStates(const MassSpringSystem &masssprings,
-                                   const WinCondition win_condition) -> void {
-  winning_states.clear();
-  winning_states.reserve(masssprings.masses.size());
-  for (const auto &[state, mass] : masssprings.masses) {
-    if (win_condition(state)) {
-      winning_states.insert(state);
-    }
-  }
-
-  std::cout << "Found " << winning_states.size() << " winning states."
-            << std::endl;
-}
-
-auto Renderer::AddWinningState(const State &state,
-                               const WinCondition win_condition) -> void {
-  if (win_condition(state)) {
-    winning_states.insert(state);
-  }
-}
-
 auto Renderer::UpdateCamera(const MassSpringSystem &masssprings,
                             const State &current) -> void {
   const Mass &c = masssprings.masses.at(current.state);
@@ -132,7 +112,9 @@ auto Renderer::UpdateTextureSizes() -> void {
 }
 
 auto Renderer::DrawMassSprings(const MassSpringSystem &masssprings,
-                               const State &current) -> void {
+                               const State &current,
+                               const std::unordered_set<State> &winning_states)
+    -> void {
   BeginTextureMode(render_target);
   ClearBackground(RAYWHITE);
 
@@ -260,7 +242,9 @@ auto Renderer::DrawKlotski(const State &state, int hov_x, int hov_y, int sel_x,
 }
 
 auto Renderer::DrawMenu(const MassSpringSystem &masssprings, int current_preset,
-                        const State &current_state) -> void {
+                        const State &current_state,
+                        const std::unordered_set<State> &winning_states)
+    -> void {
   BeginTextureMode(menu_target);
   ClearBackground(RAYWHITE);
 
diff --git a/src/state.cpp b/src/state.cpp
new file mode 100644
index 0000000..7a151b6
--- /dev/null
+++ b/src/state.cpp
@@ -0,0 +1,106 @@
+#include "state.hpp"
+#include "presets.hpp"
+
+#include <raymath.h>
+
+auto StateManager::LoadPreset(int preset) -> void {
+  current_state = generators[preset]();
+  previous_state = current_state;
+  ClearGraph();
+  current_preset = preset;
+  edited = false;
+}
+
+auto StateManager::ResetState() -> void {
+  current_state = generators[current_preset]();
+  previous_state = current_state;
+  if (edited) {
+    // We also need to clear the graph, in case the state has been edited.
+    // Then the graph would contain states that are impossible to reach.
+    ClearGraph();
+    edited = false;
+  }
+}
+
+auto StateManager::PreviousPreset() -> void {
+  LoadPreset((generators.size() + current_preset - 1) % generators.size());
+}
+
+auto StateManager::NextPreset() -> void {
+  LoadPreset((current_preset + 1) % generators.size());
+}
+
+auto StateManager::FillGraph() -> void {
+  ClearGraph();
+
+  std::pair<std::unordered_set<std::string>,
+            std::vector<std::pair<std::string, std::string>>>
+      closure = current_state.Closure();
+  for (const auto &state : closure.first) {
+    // TODO: Insert masses in the spring loop and choose the position based on
+    //       the existing mass
+    Vector3 pos =
+        Vector3(static_cast<float>(GetRandomValue(-10000, 10000)) / 1000.0,
+                static_cast<float>(GetRandomValue(-10000, 10000)) / 1000.0,
+                static_cast<float>(GetRandomValue(-10000, 10000)) / 1000.0);
+
+    mass_springs.AddMass(MASS, pos, false, state);
+  }
+  for (const auto &[from, to] : closure.second) {
+    mass_springs.AddSpring(from, to, SPRING_CONSTANT, DAMPENING_CONSTANT,
+                           REST_LENGTH);
+  }
+  std::cout << "Inserted " << mass_springs.masses.size() << " masses and "
+            << mass_springs.springs.size() << " springs." << std::endl;
+  FindWinningStates();
+  std::cout << "Consuming "
+            << sizeof(decltype(*mass_springs.masses.begin())) *
+                   mass_springs.masses.size()
+            << " Bytes for masses." << std::endl;
+  std::cout << "Consuming "
+            << sizeof(decltype(*mass_springs.springs.begin())) *
+                   mass_springs.springs.size()
+            << " Bytes for springs." << std::endl;
+}
+
+auto StateManager::UpdateGraph() -> void {
+  if (previous_state != current_state.state) {
+    mass_springs.AddMass(
+        MASS,
+        // TODO: Add beside previous_state
+        Vector3(static_cast<float>(GetRandomValue(-1000, 1000)) / 1000.0,
+                static_cast<float>(GetRandomValue(-1000, 1000)) / 1000.0,
+                static_cast<float>(GetRandomValue(-1000, 1000)) / 1000.0),
+        false, current_state.state);
+    mass_springs.AddSpring(current_state.state, previous_state, SPRING_CONSTANT,
+                           DAMPENING_CONSTANT, REST_LENGTH);
+    if (win_conditions[current_preset](current_state)) {
+      winning_states.insert(current_state);
+    }
+  }
+}
+
+auto StateManager::ClearGraph() -> void {
+  winning_states.clear();
+  mass_springs.Clear();
+  mass_springs.AddMass(MASS, Vector3Zero(), false, current_state);
+
+  // The previous_state is no longer in the graph
+  previous_state = current_state;
+}
+
+auto StateManager::FindWinningStates() -> void {
+  winning_states.clear();
+  for (const auto &[state, mass] : mass_springs.masses) {
+    if (win_conditions[current_preset](state)) {
+      winning_states.insert(state);
+    }
+  }
+
+  std::cout << "Found " << winning_states.size() << " winning states."
+            << std::endl;
+}
+
+auto StateManager::CurrentWinCondition() -> WinCondition {
+  return win_conditions[current_preset];
+}