{-|
Module:         Day19
Description:    <https://adventofcode.com/2021/day/19 Day 19: Beacon Scanner>
-}
{-# LANGUAGE FlexibleContexts, OverloadedStrings, TypeFamilies #-}
module Day19 (day19) where

import Control.Monad (guard, mfilter, replicateM)
import Data.Bits (setBit, testBit)
import Data.Bool (bool)
import Data.Containers.ListUtils (nubOrd)
import Data.Function (on)
import qualified Data.IntSet as IntSet (delete, fromDistinctAscList, member, null)
import Data.List (permutations, sortOn)
import qualified Data.Map as Map (assocs, fromListWith)
import Data.Maybe (listToMaybe)
import Data.Ord (Down(..))
import qualified Data.Set as Set (elems, fromList, intersection, size)
import Data.String (IsString)
import Data.Text (Text)
import Data.Void (Void)
import Text.Megaparsec (MonadParsec, ParseErrorBundle, Token, Tokens, chunk, parse, sepBy1, sepEndBy1, single, skipSome)
import Text.Megaparsec.Char (alphaNumChar, newline)
import Text.Megaparsec.Char.Lexer (decimal, signed)

parser :: (MonadParsec e s m, IsString (Tokens s), Token s ~ Char, Num a) => m [[[a]]]
parser :: m [[[a]]]
parser = ([[[a]]] -> Bool) -> m [[[a]]] -> m [[[a]]]
forall (m :: * -> *) a. MonadPlus m => (a -> Bool) -> m a -> m a
mfilter [[[a]]] -> Bool
forall (t :: * -> *) a. Foldable t => [[t a]] -> Bool
checkSizes (m [[[a]]] -> m [[[a]]]) -> m [[[a]]] -> m [[[a]]]
forall a b. (a -> b) -> a -> b
$ m [[a]]
scanner m [[a]] -> m Char -> m [[[a]]]
forall (m :: * -> *) a sep. MonadPlus m => m a -> m sep -> m [a]
`sepBy1` m Char
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m (Token s)
newline where
    scanner :: m [[a]]
scanner = Tokens s -> m (Tokens s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
chunk Tokens s
"--- scanner " m (Tokens s) -> m () -> m ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> m Char -> m ()
forall (m :: * -> *) a. MonadPlus m => m a -> m ()
skipSome m Char
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m (Token s)
alphaNumChar m () -> m (Tokens s) -> m (Tokens s)
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Tokens s -> m (Tokens s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
chunk Tokens s
" ---" m (Tokens s) -> m Char -> m Char
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> m Char
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m (Token s)
newline m Char -> m [[a]] -> m [[a]]
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*>
        (m () -> m a -> m a
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m () -> m a -> m a
signed (() -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) m a
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m a
decimal m a -> m Char -> m [a]
forall (m :: * -> *) a sep. MonadPlus m => m a -> m sep -> m [a]
`sepBy1` Token s -> m (Token s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Token s -> m (Token s)
single Char
Token s
',') m [a] -> m Char -> m [[a]]
forall (m :: * -> *) a sep. MonadPlus m => m a -> m sep -> m [a]
`sepEndBy1` m Char
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m (Token s)
newline
    checkSizes :: [[t a]] -> Bool
checkSizes ((t a
x:[t a]
xs):[[t a]]
xss) = (t a -> Bool) -> [t a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Int -> Int -> Bool) -> (t a -> Int) -> t a -> t a -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) t a
x) ([t a]
xs [t a] -> [t a] -> [t a]
forall a. [a] -> [a] -> [a]
++ [[t a]] -> [t a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[t a]]
xss) Bool -> Bool -> Bool
&& Bool -> Bool
not (([t a] -> Bool) -> [[t a]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any [t a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[t a]]
xss)
    checkSizes [[t a]]
_ = Bool
False

allTransforms :: Int -> [[(Bool, Int)]]
allTransforms :: Int -> [[(Bool, Int)]]
allTransforms Int
n = ([(Bool, Int)] -> Bool) -> [[(Bool, Int)]] -> [[(Bool, Int)]]
forall a. (a -> Bool) -> [a] -> [a]
filter [(Bool, Int)] -> Bool
parity ([[(Bool, Int)]] -> [[(Bool, Int)]])
-> [[(Bool, Int)]] -> [[(Bool, Int)]]
forall a b. (a -> b) -> a -> b
$ [Bool] -> [Int] -> [(Bool, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Bool] -> [Int] -> [(Bool, Int)])
-> [[Bool]] -> [[Int] -> [(Bool, Int)]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Bool] -> [[Bool]]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n [Bool
False, Bool
True] [[Int] -> [(Bool, Int)]] -> [[Int]] -> [[(Bool, Int)]]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Int] -> [[Int]]
forall a. [a] -> [[a]]
permutations [Int
0..Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] where
    parity :: [(Bool, Int)] -> Bool
parity [(Bool, Int)]
t = (Int, Bool) -> Bool
forall a b. (a, b) -> b
snd ((Int, Bool) -> Bool) -> (Int, Bool) -> Bool
forall a b. (a -> b) -> a -> b
$ (Int -> (Int, Bool) -> (Int, Bool))
-> (Int, Bool) -> [Int] -> (Int, Bool)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Int -> (Int, Bool) -> (Int, Bool)
decomp (Int
0 :: Int, Bool
True) [Int
0..Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] where
        decomp :: Int -> (Int, Bool) -> (Int, Bool)
decomp Int
i k :: (Int, Bool)
k@(Int
bits, Bool
p) = if Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
bits Int
i then (Int, Bool)
k else Int -> Int -> Bool -> (Int, Bool)
decomp' Int
i Int
bits (Bool -> (Int, Bool)) -> Bool -> (Int, Bool)
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not Bool
p
        decomp' :: Int -> Int -> Bool -> (Int, Bool)
decomp' Int
i Int
bits Bool
p = if Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
bits Int
i then (Int
bits, Bool
p) else Int -> Int -> Bool -> (Int, Bool)
decomp' Int
j (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
setBit Int
bits Int
i) (Bool -> (Int, Bool)) -> Bool -> (Int, Bool)
forall a b. (a -> b) -> a -> b
$ Bool
p Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
q where
            (Bool
q, Int
j) = [(Bool, Int)]
t [(Bool, Int)] -> Int -> (Bool, Int)
forall a. [a] -> Int -> a
!! Int
i

applyTransform :: (Num a) => [(Bool, Int)] -> [a] -> [a]
applyTransform :: [(Bool, Int)] -> [a] -> [a]
applyTransform [(Bool, Int)]
t [a]
point = [(a -> a) -> (a -> a) -> Bool -> a -> a
forall a. a -> a -> Bool -> a
bool a -> a
forall a. a -> a
id a -> a
forall a. Num a => a -> a
negate Bool
s (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ [a]
point [a] -> Int -> a
forall a. [a] -> Int -> a
!! Int
i | (Bool
s, Int
i) <- [(Bool, Int)]
t]

day19 :: Text -> Either (ParseErrorBundle Text Void) (Maybe (Int, Int))
day19 :: Text -> Either (ParseErrorBundle Text Void) (Maybe (Int, Int))
day19 Text
input = do
    [[[Int]]]
scanners <- Parsec Void Text [[[Int]]]
-> String -> Text -> Either (ParseErrorBundle Text Void) [[[Int]]]
forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
parse Parsec Void Text [[[Int]]]
forall e s (m :: * -> *) a.
(MonadParsec e s m, IsString (Tokens s), Token s ~ Char, Num a) =>
m [[[a]]]
parser String
"" Text
input
    let size :: Int
size = [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ [[Int]] -> [Int]
forall a. [a] -> a
head ([[Int]] -> [Int]) -> [[Int]] -> [Int]
forall a b. (a -> b) -> a -> b
$ [[[Int]]] -> [[Int]]
forall a. [a] -> a
head [[[Int]]]
scanners
        deltas :: Map (Set [Int]) [(Int, [(Bool, Int)])]
deltas = ([(Int, [(Bool, Int)])]
 -> [(Int, [(Bool, Int)])] -> [(Int, [(Bool, Int)])])
-> [(Set [Int], [(Int, [(Bool, Int)])])]
-> Map (Set [Int]) [(Int, [(Bool, Int)])]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [(Int, [(Bool, Int)])]
-> [(Int, [(Bool, Int)])] -> [(Int, [(Bool, Int)])]
forall a. Semigroup a => a -> a -> a
(<>) ([(Set [Int], [(Int, [(Bool, Int)])])]
 -> Map (Set [Int]) [(Int, [(Bool, Int)])])
-> [(Set [Int], [(Int, [(Bool, Int)])])]
-> Map (Set [Int]) [(Int, [(Bool, Int)])]
forall a b. (a -> b) -> a -> b
$ do
            (Int
i, [[Int]]
scanner) <- [Int] -> [[[Int]]] -> [(Int, [[Int]])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] ([[[Int]]] -> [(Int, [[Int]])]) -> [[[Int]]] -> [(Int, [[Int]])]
forall a b. (a -> b) -> a -> b
$ [[[Int]]] -> [[[Int]]]
forall a. [a] -> [a]
tail [[[Int]]]
scanners
            let delta :: [[Int]]
delta = [[Int]] -> [[Int]]
forall a. Ord a => [a] -> [a]
nubOrd ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) ([Int] -> [Int] -> [Int]) -> [[Int]] -> [[Int] -> [Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[Int]]
scanner [[Int] -> [Int]] -> [[Int]] -> [[Int]]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [[Int]]
scanner
            [(Bool, Int)]
t <- Int -> [[(Bool, Int)]]
allTransforms Int
size
            (Set [Int], [(Int, [(Bool, Int)])])
-> [(Set [Int], [(Int, [(Bool, Int)])])]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([[Int]] -> Set [Int]
forall a. Ord a => [a] -> Set a
Set.fromList ([[Int]] -> Set [Int]) -> [[Int]] -> Set [Int]
forall a b. (a -> b) -> a -> b
$ [(Bool, Int)] -> [Int] -> [Int]
forall a. Num a => [(Bool, Int)] -> [a] -> [a]
applyTransform [(Bool, Int)]
t ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[Int]]
delta, [(Int
i, [(Bool, Int)]
t)])
        go :: Set [Int] -> [[Int]] -> IntSet -> [(Int, Int)]
go Set [Int]
beacons [[Int]]
positions IntSet
remaining
          | IntSet -> Bool
IntSet.null IntSet
remaining
          = (Int, Int) -> [(Int, Int)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Set [Int] -> Int
forall a. Set a -> Int
Set.size Set [Int]
beacons, [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [[Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Int
forall a. Num a => a -> a
abs (Int -> Int) -> [Int] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [Int]
x [Int]
y | [Int]
x <- [[Int]]
positions, [Int]
y <- [[Int]]
positions])
          | Bool
otherwise = do
                let delta0 :: Set [Int]
delta0 = [[Int]] -> Set [Int]
forall a. Ord a => [a] -> Set a
Set.fromList ([[Int]] -> Set [Int]) -> [[Int]] -> Set [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) ([Int] -> [Int] -> [Int]) -> [[Int]] -> [[Int] -> [Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set [Int] -> [[Int]]
forall a. Set a -> [a]
Set.elems Set [Int]
beacons [[Int] -> [Int]] -> [[Int]] -> [[Int]]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set [Int] -> [[Int]]
forall a. Set a -> [a]
Set.elems Set [Int]
beacons
                (Int
i, [(Bool, Int)]
t) <- ((Set [Int], [(Int, [(Bool, Int)])]) -> Down Int)
-> [(Set [Int], [(Int, [(Bool, Int)])])]
-> [(Set [Int], [(Int, [(Bool, Int)])])]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Int -> Down Int
forall a. a -> Down a
Down (Int -> Down Int)
-> ((Set [Int], [(Int, [(Bool, Int)])]) -> Int)
-> (Set [Int], [(Int, [(Bool, Int)])])
-> Down Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set [Int] -> Int
forall a. Set a -> Int
Set.size (Set [Int] -> Int)
-> ((Set [Int], [(Int, [(Bool, Int)])]) -> Set [Int])
-> (Set [Int], [(Int, [(Bool, Int)])])
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set [Int] -> Set [Int] -> Set [Int]
forall a. Ord a => Set a -> Set a -> Set a
Set.intersection Set [Int]
delta0 (Set [Int] -> Set [Int])
-> ((Set [Int], [(Int, [(Bool, Int)])]) -> Set [Int])
-> (Set [Int], [(Int, [(Bool, Int)])])
-> Set [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set [Int], [(Int, [(Bool, Int)])]) -> Set [Int]
forall a b. (a, b) -> a
fst) (Map (Set [Int]) [(Int, [(Bool, Int)])]
-> [(Set [Int], [(Int, [(Bool, Int)])])]
forall k a. Map k a -> [(k, a)]
Map.assocs Map (Set [Int]) [(Int, [(Bool, Int)])]
deltas) [(Set [Int], [(Int, [(Bool, Int)])])]
-> ((Set [Int], [(Int, [(Bool, Int)])]) -> [(Int, [(Bool, Int)])])
-> [(Int, [(Bool, Int)])]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Set [Int], [(Int, [(Bool, Int)])]) -> [(Int, [(Bool, Int)])]
forall a b. (a, b) -> b
snd
                Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Int -> IntSet -> Bool
IntSet.member Int
i IntSet
remaining
                let points :: [[Int]]
points = [(Bool, Int)] -> [Int] -> [Int]
forall a. Num a => [(Bool, Int)] -> [a] -> [a]
applyTransform [(Bool, Int)]
t ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[[Int]]]
scanners [[[Int]]] -> Int -> [[Int]]
forall a. [a] -> Int -> a
!! Int
i
                [Int]
position <- (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) ([Int] -> [Int] -> [Int]) -> [[Int]] -> [[Int] -> [Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set [Int] -> [[Int]]
forall a. Set a -> [a]
Set.elems Set [Int]
beacons [[Int] -> [Int]] -> [[Int]] -> [[Int]]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [[Int]]
points
                let beacons' :: Set [Int]
beacons' = [[Int]] -> Set [Int]
forall a. Ord a => [a] -> Set a
Set.fromList ([[Int]] -> Set [Int]) -> [[Int]] -> Set [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [Int]
position ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[Int]]
points
                Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Set [Int] -> Int
forall a. Set a -> Int
Set.size (Set [Int] -> Set [Int] -> Set [Int]
forall a. Ord a => Set a -> Set a -> Set a
Set.intersection Set [Int]
beacons Set [Int]
beacons') Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
12
                Set [Int] -> [[Int]] -> IntSet -> [(Int, Int)]
go (Set [Int]
beacons Set [Int] -> Set [Int] -> Set [Int]
forall a. Semigroup a => a -> a -> a
<> Set [Int]
beacons') ([Int]
position[Int] -> [[Int]] -> [[Int]]
forall a. a -> [a] -> [a]
:[[Int]]
positions) (IntSet -> [(Int, Int)]) -> IntSet -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ Int -> IntSet -> IntSet
IntSet.delete Int
i IntSet
remaining
    Maybe (Int, Int)
-> Either (ParseErrorBundle Text Void) (Maybe (Int, Int))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Int, Int)
 -> Either (ParseErrorBundle Text Void) (Maybe (Int, Int)))
-> Maybe (Int, Int)
-> Either (ParseErrorBundle Text Void) (Maybe (Int, Int))
forall a b. (a -> b) -> a -> b
$ [(Int, Int)] -> Maybe (Int, Int)
forall a. [a] -> Maybe a
listToMaybe ([(Int, Int)] -> Maybe (Int, Int))
-> [(Int, Int)] -> Maybe (Int, Int)
forall a b. (a -> b) -> a -> b
$ Set [Int] -> [[Int]] -> IntSet -> [(Int, Int)]
go ([[Int]] -> Set [Int]
forall a. Ord a => [a] -> Set a
Set.fromList ([[Int]] -> Set [Int]) -> [[Int]] -> Set [Int]
forall a b. (a -> b) -> a -> b
$ [[[Int]]] -> [[Int]]
forall a. [a] -> a
head [[[Int]]]
scanners) [Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
size Int
0] (IntSet -> [(Int, Int)]) -> IntSet -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$
        [Int] -> IntSet
IntSet.fromDistinctAscList [Int
1..[[[Int]]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[[Int]]]
scanners Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]