{-# OPTIONS_GHC -Wall #-}
import qualified Data.Text.IO as TextIO
import Control.Monad ((<$!>))
import qualified Data.Array.Unboxed as Array
import qualified Data.Text as Text
import qualified Data.Char as Char
import Data.Array.Unboxed (UArray)
import qualified Data.Foldable as Foldable
import Control.Arrow ((&&&))
import qualified Data.List as List

parse :: Text.Text -> UArray (Int, Int) Int
parse t = let
    banks = init $ Text.lines t
    bankSize = maybe 0 pred $ Text.findIndex (== '\n') t
    bankCount = Text.count "\n" t - 2
  in Array.listArray ((0, 0), (bankCount, bankSize)) $ List.concatMap (fmap Char.digitToInt . Text.unpack) banks

rowsOf :: UArray (Int, Int) Int -> Int
rowsOf = fst . snd . Array.bounds

colsOf :: UArray (Int, Int) Int -> Int
colsOf = snd . snd . Array.bounds

main :: IO ()
main = do
  batteryBanks <- parse <$!> TextIO.getContents
  print $ part1 batteryBanks
  print $ part2 batteryBanks

part1 :: UArray (Int, Int) Int -> Int
part1 batteryBanks = Foldable.sum $ pickBatteries 2 batteryBanks <$> [0.. rowsOf batteryBanks]

part2 :: UArray (Int, Int) Int -> Int
part2 banks = Foldable.sum $ pickBatteries 12 banks <$> [0.. rowsOf banks]

pickBatteries :: Int -> UArray (Int, Int) Int -> Int -> Int
pickBatteries batteryCount banks row = let
    width = colsOf banks
    getBattery col = banks Array.! (row, col)

    go acc 0 _      = acc
    go acc n offset = let
        effectiveEnd = width - pred n
        availableIndices = [offset .. effectiveEnd]
        batteryWithIndices = (id &&& getBattery) <$> availableIndices
        (offset', selectedBattery) = maximumOn snd batteryWithIndices
      in go (acc * 10 + selectedBattery) (pred n) (succ offset')
  in go 0 batteryCount 0

maximumOn :: (Foldable t, Ord b) => (a -> b) -> t a -> a
maximumOn f collection = case Foldable.toList collection of
  [] -> error "maximumOn: empty foldable"
  (x:xs) -> List.foldl selectMax x xs
  where
    selectMax a b = if f a < f b then b else a