{-|
Module:         Day8
Description:    <https://adventofcode.com/2021/day/8 Day 8: Seven Segment Search>
-}
{-# LANGUAGE OverloadedStrings #-}
module Day8 (day8a, day8b) where

import Control.Arrow (first)
import Control.Monad (guard)
import Data.Bits ((.&.), (.|.), complement, popCount, setBit)
import Data.Char (ord)
import qualified Data.IntMap as IntMap ((!?), fromListWith)
import Data.List (elemIndex, foldl', partition)
import Data.Text (Text)
import qualified Data.Text as T (breakOnEnd, foldl', length, lines, stripSuffix, words)

day8a :: Text -> Int
day8a :: Text -> Int
day8a Text
input = [()] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([()] -> Int) -> [()] -> Int
forall a b. (a -> b) -> a -> b
$ do
    Text
word <- Text -> [Text]
T.lines Text
input [Text] -> (Text -> [Text]) -> [Text]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Text -> [Text]
T.words (Text -> [Text]) -> (Text -> Text) -> Text -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, Text) -> Text
forall a b. (a, b) -> b
snd ((Text, Text) -> Text) -> (Text -> (Text, Text)) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text -> (Text, Text)
T.breakOnEnd Text
" | "
    Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Text -> Int
T.length Text
word Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int
2, Int
4, Int
3, Int
7]

day8b :: Text -> Maybe Int
day8b :: Text -> Maybe Int
day8b Text
input = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> Maybe [Int] -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> Maybe Int) -> [Text] -> Maybe [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Text -> Maybe Int
handle (Text -> [Text]
T.lines Text
input)

handle :: Text -> Maybe Int
handle :: Text -> Maybe Int
handle Text
line
  | (Just Text
lhs, Text
rhs) <- (Text -> Maybe Text) -> (Text, Text) -> (Maybe Text, Text)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (Text -> Text -> Maybe Text
T.stripSuffix Text
" | ") ((Text, Text) -> (Maybe Text, Text))
-> (Text, Text) -> (Maybe Text, Text)
forall a b. (a -> b) -> a -> b
$ Text -> Text -> (Text, Text)
T.breakOnEnd Text
" | " Text
line = do
    let signals :: [Int]
signals = Text -> Int
bits (Text -> Int) -> [Text] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> [Text]
T.words Text
lhs
        outputs :: [Int]
outputs = Text -> Int
bits (Text -> Int) -> [Text] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> [Text]
T.words Text
rhs
        counts :: IntMap [Int]
counts = ([Int] -> [Int] -> [Int]) -> [(Int, [Int])] -> IntMap [Int]
forall a. (a -> a -> a) -> [(Int, a)] -> IntMap a
IntMap.fromListWith [Int] -> [Int] -> [Int]
forall a. Semigroup a => a -> a -> a
(<>) [(Int -> Int
forall a. Bits a => a -> Int
popCount Int
s, [Int
s]) | Int
s <- [Int]
signals]
    [Int
one] <- IntMap [Int]
counts IntMap [Int] -> Int -> Maybe [Int]
forall a. IntMap a -> Int -> Maybe a
IntMap.!? Int
2
    [Int
seven] <- IntMap [Int]
counts IntMap [Int] -> Int -> Maybe [Int]
forall a. IntMap a -> Int -> Maybe a
IntMap.!? Int
3
    [Int
four] <- IntMap [Int]
counts IntMap [Int] -> Int -> Maybe [Int]
forall a. IntMap a -> Int -> Maybe a
IntMap.!? Int
4
    ([Int
two], [Int]
threeFive) <- (Int -> Bool) -> [Int] -> ([Int], [Int])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2) (Int -> Bool) -> (Int -> Int) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Bits a => a -> a -> Int
popCountWithout (Int
four Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
seven)) ([Int] -> ([Int], [Int])) -> Maybe [Int] -> Maybe ([Int], [Int])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap [Int]
counts IntMap [Int] -> Int -> Maybe [Int]
forall a. IntMap a -> Int -> Maybe a
IntMap.!? Int
5
    ([Int
three], [Int
five]) <- ([Int], [Int]) -> Maybe ([Int], [Int])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (([Int], [Int]) -> Maybe ([Int], [Int]))
-> ([Int], [Int]) -> Maybe ([Int], [Int])
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> ([Int], [Int])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1) (Int -> Bool) -> (Int -> Int) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Bits a => a -> a -> Int
popCountWithout Int
two) [Int]
threeFive
    ([Int
six], [Int]
zeroNine) <- (Int -> Bool) -> [Int] -> ([Int], [Int])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (Int -> Bool) -> (Int -> Int) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
one Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&.) (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int
forall a. Bits a => a -> a
complement) ([Int] -> ([Int], [Int])) -> Maybe [Int] -> Maybe ([Int], [Int])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap [Int]
counts IntMap [Int] -> Int -> Maybe [Int]
forall a. IntMap a -> Int -> Maybe a
IntMap.!? Int
6
    ([Int
zero], [Int
nine]) <- ([Int], [Int]) -> Maybe ([Int], [Int])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (([Int], [Int]) -> Maybe ([Int], [Int]))
-> ([Int], [Int]) -> Maybe ([Int], [Int])
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> ([Int], [Int])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2) (Int -> Bool) -> (Int -> Int) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Bits a => a -> a -> Int
popCountWithout Int
three) [Int]
zeroNine
    [Int
eight] <- IntMap [Int]
counts IntMap [Int] -> Int -> Maybe [Int]
forall a. IntMap a -> Int -> Maybe a
IntMap.!? Int
7
    let digits :: [Int]
digits = [Int
zero, Int
one, Int
two, Int
three, Int
four, Int
five, Int
six, Int
seven, Int
eight, Int
nine]
    (Int -> Int -> Int) -> Int -> [Int] -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
x Int
y -> Int
10 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y) Int
0 ([Int] -> Int) -> Maybe [Int] -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> Maybe Int) -> [Int] -> Maybe [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Int -> [Int] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` [Int]
digits) [Int]
outputs
  | Bool
otherwise = Maybe Int
forall a. Maybe a
Nothing
  where
    bits :: Text -> Int
bits = (Int -> Char -> Int) -> Int -> Text -> Int
forall a. (a -> Char -> a) -> a -> Text -> a
T.foldl' (\Int
acc Char
c -> Int
acc Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`setBit` (Char -> Int
ord Char
c Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
31)) (Int
0 :: Int)
    popCountWithout :: a -> a -> Int
popCountWithout a
antimask = a -> Int
forall a. Bits a => a -> Int
popCount (a -> Int) -> (a -> a) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a
forall a. Bits a => a -> a -> a
.&. a -> a
forall a. Bits a => a -> a
complement a
antimask)