Haskell solution to Google Code Jam problem BotTrust

BotTrust.hs

-- Haskell solution to Google Code Jam 2011 problem, Bot Trust
-- http://code.google.com/codejam/contest/dashboard?c=975485#s=p0

import Data.List (find)
import Data.Maybe (fromMaybe)
import System.Environment (getArgs)

data Robot = Blue | Orange
             deriving Eq

instance Show Robot where
  show Blue = "B"
  show Orange = "O"

data ButtonPress = ButtonPress Robot Int

instance Show ButtonPress where
  show (ButtonPress robot number) = show robot ++ show number

data TestChamber = TestChamber
                   Int -- blue position
                   Int -- orange position
                   [ButtonPress] -- list of moves required
                   deriving Show

whichBot :: ButtonPress -> Robot
whichBot (ButtonPress bot _) = bot

position :: ButtonPress -> Int
position (ButtonPress _ p) = p

newChamber :: [ButtonPress] -> TestChamber
newChamber = TestChamber 1 1

-- Given a TestChamber, returns the number of moves required to solve it
solve :: TestChamber -> Int
solve (TestChamber _ _ []) = 0
solve tc = 1 + solve (step tc)

-- Given a TestChamber, returns a new TestChamber representing what the chamber would look like, one move later
step :: TestChamber -> TestChamber
step (TestChamber _ _ []) = error "The test is over. Please assume the party escort submission position."
step tc@(TestChamber _ _ bpbps@(ButtonPress nextBot x : bps)) =
  TestChamber b o newBs
  where 
    newBs = if x == positionOf nextBot tc then bps else bpbps
    b = nextPositionFor Blue   tc
    o = nextPositionFor Orange tc

-- returns the position the robot is trying to get to (or zero if they're done)
nextPressFor :: Robot -> TestChamber -> Int
nextPressFor robot (TestChamber _ _ bps) =
  case find ((robot ==) . whichBot) bps of
    Just bp -> position bp
    Nothing -> 0 -- It doesn't actually matter what we put here because this robot is done

-- returns the next position for the given robot. Equal to its current position, plus -1, 0, or 1
nextPositionFor :: Robot -> TestChamber -> Int
nextPositionFor robot tc
  | nextPress == pos = pos
  | nextPress < pos  = pos-1
  | nextPress > pos  = pos+1
  where pos = positionOf robot tc
        nextPress = nextPressFor robot tc
        
positionOf Blue   (TestChamber b _ _) = b
positionOf Orange (TestChamber _ o _) = o

main = do 
  args <- getArgs
  case args of
    [input,output] -> processFile input output
    _ -> putStrLn "error: exactly two arguments needed"

processFile input output = do
  fileContents <- readFile input
  writeFile output (produceSolutionOutput fileContents)

produceSolutionOutput fileContents = 
  unlines $ map showChamber numberedChambers
  where fileLines = tail $ lines fileContents
        chambers = map lineToChamber fileLines
        numberedChambers = zip [1..] chambers
        showChamber (i, chamber) = "Case #" ++ (show i) ++ ": " ++ (show $ solve chamber)

lineToChamber :: String -> TestChamber
lineToChamber = newChamber . stringsToButtonPresses . tail . words
  
stringsToButtonPresses :: [String] -> [ButtonPress]  
stringsToButtonPresses [] = []
stringsToButtonPresses (botString:posString:rest) =
  let bp = (ButtonPress (stringToRobot botString) ((read posString)::Int))
  in  bp : stringsToButtonPresses rest
stringsToButtonPresses _ = error "Error parsing input"

stringToRobot "O" = Orange
stringToRobot "B" = Blue
stringToRobot s   = error ("Expected 'O' or 'B', found '" ++ s ++ "'")

-- solve tc1 = 6
tc1 = newChamber [ButtonPress Orange 2, ButtonPress Blue 1, ButtonPress Blue 2, ButtonPress Orange 4]

-- solve tc2 = 100
tc2 = newChamber [ButtonPress Orange 5, ButtonPress Orange 8, ButtonPress Blue 100]

-- solve tc3 = 4
tc3 = newChamber [ButtonPress Blue 2, ButtonPress Blue 1]