diff --git a/05/Main2.idr b/05/Main2.idr
index b32c082..c4b2f32 100644
--- a/05/Main2.idr
+++ b/05/Main2.idr
@@ -93,20 +93,35 @@ applyCommand multiple (Cmd count from to) stacks =
 applyCommands : {n : Nat} -> Bool -> List (Command n) -> Stacks n -> Stacks n
 applyCommands multiple commands stacks = foldl (\stack, command => applyCommand multiple command stack) stacks commands
 
+parseProblem : String -> Maybe (n ** (Stacks n, List (Command n)))
+parseProblem content =
+  case last' . takeWhile(/= "") . lines $ content of
+    Just columnsLine =>
+      let size = foldl max 0 . catMaybes . map parsePositive . words $ columnsLine
+          stacks = parseStacks {n = size} . takeWhile (/= "") . lines $ content
+          commands = catMaybes . map (parseCommand {n = size}) . drop 1 . dropWhile (/= "") . lines $ content
+      in Just (size ** (stacks, commands))
+    Nothing => Nothing
+
+
+doPart : {n : Nat} -> Bool -> Stacks n -> List (Command n) -> String
+doPart multiple stacks commands = tops $ applyCommands multiple commands stacks
+
+
 main : IO ()
 main =
   do file <- readFile "input"
      case file of
        Right content =>
-         let stacks = parseStacks {n = 9} . takeWhile (/= "") . lines $ content
-             commands = catMaybes . map (parseCommand {n = 9}) . drop 1 . dropWhile (/= "") . lines $ content
-             part1 = tops $ applyCommands False commands stacks
-             part2 = tops $ applyCommands True commands stacks
-         in do putStrLn "Input:"
-               printLn stacks
-               putStrLn ""
-               putStr "Part 1: "
-               putStrLn part1
-               putStr "Part 2: "
-               putStrLn part2
+         case parseProblem content of
+           Just (size ** (stacks, commands)) =>
+             let part1 = doPart False stacks commands
+                 part2 = doPart True stacks commands
+             in do putStrLn "Input:"
+                   printLn stacks
+                   putStr "\n\nPart 1: "
+                   putStrLn part1
+                   putStr "Part 2: "
+                   putStrLn part2
+           Nothing => putStrLn "Failed to parse problem"
        Left err => printLn err