import           Data.List       (group, minimumBy, sort, sortBy)
import qualified Data.Map.Strict as Map
import           Data.Maybe      (mapMaybe)
import           Data.Ord        (comparing)
import qualified Data.Set        as Set
import           Parsing         (splitByString)

type Record = ([Set.Set Char], [Set.Set Char])

main :: IO ()
main = do
  input <- getContents
  let
    records = parseRecords input
  (putStrLn . show . solution1) records
  (putStrLn . show . solution2) records

solution1 :: [Record] -> Int
solution1 = sum . (map countEasy)

solution2 :: [Record] -> Int
solution2 = sum . (map getNum)
  where
    getNum :: Record -> Int
    getNum (wires, display) = sum (map (uncurry (*)) (zip [1000, 100, 10, 1] digits))
      where
        digits :: [Int]
        digits = map (\d -> (Map.!) mapping d) display
        mapping :: Map.Map (Set.Set Char) Int
        mapping = findSegments wires

countEasy :: Record -> Int
countEasy record = length (mapMaybe checkDigit (snd record))
  where
    checkDigit :: Set.Set Char -> Maybe Int
    checkDigit d
      | len == 2 = Just 1
      | len == 4 = Just 4
      | len == 3 = Just 7
      | len == 7 = Just 8
      | otherwise = Nothing
        where len = length d

findSegments :: [Set.Set Char] -> Map.Map (Set.Set Char) Int
findSegments w =
  let
    wires = sortBy (comparing Set.size) w
    unknown = (map snd (filter isKnown (zip [0..] wires)))
      where
        isKnown :: (Int, Set.Set Char) -> Bool
        isKnown (i, _)
          | i == 0 = False -- 1
          | i == 1 = False -- 7
          | i == 2 = False -- 4
          | i == 9 = False -- 8
          | otherwise = True
    segmentA = (wires !! 1) `Set.difference` (wires !! 0) -- 7 - 1
    segmentB = (wires !! 2) `Set.difference` ((wires !! 0) `Set.union` segmentD) -- 4 - 1 + d
    segmentC = least cf
      where
        abdeg = unionAll [segmentA, segmentB, segmentD, segmentE, segmentG]
        cf = filter notTwo (map (`Set.difference` abdeg) unknown)
        notTwo wire =
          (Set.size wire == 1)
    segmentD = head (filter notZero (map (Set.difference (wires !! 9)) unknown))
      where
        notZero wire =
          (Set.size wire == 1)
          && ((wire `Set.difference` segmentE) /= Set.empty)
          && ((wire `Set.difference` (wires !! 0) /= Set.empty))
    segmentE = (wires !! 9) `Set.difference`
               ((wires !! 1) `Set.union`
                (wires !! 2) `Set.union`
                segmentG) -- 8 - (4 + 7 + a)
    segmentF = (wires !! 0) `Set.difference` segmentC
    segmentG = (foldr1 Set.intersection unknown) `Set.difference` segmentA
  in
    Map.fromList [
    (unionAll [segmentA, segmentB, segmentC, segmentE, segmentF, segmentG], 0),
    (wires !! 0, 1),
    (unionAll [segmentA, segmentC, segmentD, segmentE, segmentG], 2),
    (unionAll [segmentA, segmentC, segmentD, segmentF, segmentG], 3),
    (wires !! 2, 4),
    (unionAll [segmentA, segmentB, segmentD, segmentF, segmentG], 5),
    (unionAll [segmentA, segmentB, segmentD, segmentE, segmentF, segmentG], 6),
    (wires !! 1, 7),
    (wires !! 9, 8),
    (unionAll [segmentA, segmentB, segmentC, segmentD, segmentF, segmentG], 9)
    ]

parseRecords :: String -> [Record]
parseRecords = map parseRecord . lines

parseRecord :: String -> Record
parseRecord = splitTuple . splitHalf
  where
    splitTuple :: (String, String) -> Record
    splitTuple (patterns, output) =
      (map Set.fromList (splitByString " " patterns), map Set.fromList (splitByString " " output))
    splitHalf :: String -> (String, String)
    splitHalf string =
      let
        halves = splitByString " | " string
      in
        ((head halves), (last halves))

least :: Ord a => [a] -> a
least = head . (minimumBy (comparing length)) . group . sort

unionAll :: Ord a => [Set.Set a] -> Set.Set a
unionAll = foldr1 Set.union

-- Tests

testInput1 = "acedgfb cdfbe gcdfa fbcad dab cefabd cdfgeb eafb cagedb ab | cdfeb fcadb cdfeb cdbaf"
testInput2 = unlines [
  "be cfbegad cbdgef fgaecd cgeb fdcge agebfd fecdb fabcd edb | fdgacbe cefdb cefbgd gcbe",
  "edbfga begcd cbg gc gcadebf fbgde acbgfd abcde gfcbed gfec | fcgedb cgb dgebacf gc",
  "fgaebd cg bdaec gdafb agbcfd gdcbef bgcad gfac gcb cdgabef | cg cg fdcagb cbg",
  "fbegcd cbd adcefb dageb afcb bc aefdc ecdab fgdeca fcdbega | efabcd cedba gadfec cb",
  "aecbfdg fbg gf bafeg dbefa fcge gcbea fcaegb dgceab fcbdga | gecf egdcabf bgf bfgea",
  "fgeab ca afcebg bdacfeg cfaedg gcfdb baec bfadeg bafgc acf | gebdcfa ecba ca fadegcb",
  "dbcfg fgd bdegcaf fgec aegbdf ecdfab fbedc dacgb gdcebf gf | cefg dcbef fcge gbcadfe",
  "bdfegc cbegaf gecbf dfcage bdacg ed bedf ced adcbefg gebcd | ed bcgafe cdgba cbgef",
  "egadfb cdbfeg cegd fecab cgb gbdefca cg fgcdab egfdb bfceg | gbdfcae bgc cg cgb",
  "gcafb gcf dcaebfg ecagb gf abcdeg gaef cafbge fdbac fegbdc | fgae cfgab fg bagce"
  ]

test1 = solution1 (parseRecords testInput2)
test2 = solution2 (parseRecords testInput2)