{-# Language BlockArguments, ImportQualifiedPost, QuasiQuotes #-}
{-|
Module      : Main
Description : Day 20 solution
Copyright   : (c) Eric Mertens, 2020
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2020/day/20>

Assembling a jigsaw puzzle and searching for sea serpents in the final image

-}
module Main (main) where

import Advent (arrIx, countBy, same)
import Advent.Coord
import Advent.Format (format)
import Data.Array.Unboxed qualified as A
import Data.Bits (setBit)
import Data.IntMap (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.List (foldl')

type Picture = [Coord]

-- <https://www.youtube.com/watch?v=EIyixC9NsLI>
snek :: Picture
snek :: Picture
snek =
  [String] -> Picture
toPicture
    [String
"                  # "
    ,String
"#    ##    ##    ###"
    ,String
" #  #  #  #  #  #   "]

-- | Rotate an image 90 degrees clockwise
rotate :: Picture -> Picture
rotate :: Picture -> Picture
rotate Picture
xs = (Coord -> Coord) -> Picture -> Picture
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> Int -> Coord
C Int
0 Int
n Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
+) (Coord -> Coord) -> (Coord -> Coord) -> Coord -> Coord
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Coord -> Coord
turnRight) Picture
xs
  where
    n :: Int
n = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((Coord -> Int) -> Picture -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Coord -> Int
coordRow Picture
xs)

-- | Generate all 8 rotations and flips of a picture
reorient :: Picture -> [Picture]
reorient :: Picture -> [Picture]
reorient Picture
xs =
  do Picture
xs' <- Int -> [Picture] -> [Picture]
forall a. Int -> [a] -> [a]
take Int
4 ((Picture -> Picture) -> Picture -> [Picture]
forall a. (a -> a) -> a -> [a]
iterate Picture -> Picture
rotate Picture
xs)
     [Picture
xs', (Coord -> Coord) -> Picture -> Picture
forall a b. (a -> b) -> [a] -> [b]
map Coord -> Coord
invert Picture
xs']

toPicture :: [String] -> Picture
toPicture :: [String] -> Picture
toPicture [String]
rs = [Int -> Int -> Coord
C Int
y Int
x | (Int
y,String
r) <- [Int] -> [String] -> [(Int, String)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [String]
rs, (Int
x,Char
'#') <- [Int] -> String -> [(Int, Char)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] String
r]

-- | Characterize image orientations by their left edge
edgeMap :: [(Int,Picture)] -> IntMap [(Int,Picture)]
edgeMap :: [(Int, Picture)] -> IntMap [(Int, Picture)]
edgeMap [(Int, Picture)]
xs =
  ([(Int, Picture)] -> [(Int, Picture)] -> [(Int, Picture)])
-> [(Int, [(Int, Picture)])] -> IntMap [(Int, Picture)]
forall a. (a -> a -> a) -> [(Int, a)] -> IntMap a
IntMap.fromListWith [(Int, Picture)] -> [(Int, Picture)] -> [(Int, Picture)]
forall a. [a] -> [a] -> [a]
(++)
    [(Picture -> Int
leftEdge Picture
pic', [(Int
i,Picture
pic')]) | (Int
i, Picture
pic) <- [(Int, Picture)]
xs, Picture
pic' <- Picture -> [Picture]
reorient Picture
pic]

-- |
-- >>> :main
-- 8581320593371
-- 2031
main :: IO ()
IO ()
main =
  do [(Int, Picture)]
inp <- ((Int, [String]) -> (Int, Picture))
-> [(Int, [String])] -> [(Int, Picture)]
forall a b. (a -> b) -> [a] -> [b]
map (([String] -> Picture) -> (Int, [String]) -> (Int, Picture)
forall a b. (a -> b) -> (Int, a) -> (Int, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [String] -> Picture
toPicture) ([(Int, [String])] -> [(Int, Picture)])
-> IO [(Int, [String])] -> IO [(Int, Picture)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [format|2020 20 (Tile %u:%n(%s%n)*%n)*|]

     -- group available pictures by their left-edge code
     let em :: IntMap [(Int, Picture)]
em = [(Int, Picture)] -> IntMap [(Int, Picture)]
edgeMap [(Int, Picture)]
inp

     -- pick a tile with a unique left and top edge to be the first corner
     let corner :: (Int, Picture)
corner = [(Int, Picture)] -> (Int, Picture)
forall a. HasCallStack => [a] -> a
head
                   [ (Int, Picture)
x | (Int, Picture)
x:[(Int, Picture)]
xs <- IntMap [(Int, Picture)] -> [[(Int, Picture)]]
forall a. IntMap a -> [a]
IntMap.elems IntMap [(Int, Picture)]
em
                       , let sameCodes :: [(Int, b)] -> Bool
sameCodes = [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => t a -> Bool
same ([Int] -> Bool) -> ([(Int, b)] -> [Int]) -> [(Int, b)] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, b) -> Int) -> [(Int, b)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, b) -> Int
forall a b. (a, b) -> a
fst -- edge codes that only match themselves
                       , [(Int, Picture)] -> Bool
forall {b}. [(Int, b)] -> Bool
sameCodes ((Int, Picture)
x(Int, Picture) -> [(Int, Picture)] -> [(Int, Picture)]
forall a. a -> [a] -> [a]
:[(Int, Picture)]
xs)
                       , [(Int, Picture)] -> Bool
forall {b}. [(Int, b)] -> Bool
sameCodes (IntMap [(Int, Picture)]
em IntMap [(Int, Picture)] -> Int -> [(Int, Picture)]
forall a. IntMap a -> Int -> a
IntMap.! Picture -> Int
topEdge ((Int, Picture) -> Picture
forall a b. (a, b) -> b
snd (Int, Picture)
x))]

     -- arrange all the tiles
     let image :: Array Coord (Int, Picture)
image = IntMap [(Int, Picture)]
-> (Int, Picture) -> Array Coord (Int, Picture)
place IntMap [(Int, Picture)]
em (Int, Picture)
corner

     -- print the product of the corner tile IDs
     Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> IO ()) -> Int -> IO ()
forall a b. (a -> b) -> a -> b
$ [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [(Int, Picture) -> Int
forall a b. (a, b) -> a
fst (Array Coord (Int, Picture)
image Array Coord (Int, Picture) -> Coord -> (Int, Picture)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
A.! Int -> Int -> Coord
C Int
y Int
x) | Int
y <- [Int
0, Int
11], Int
x <- [Int
0, Int
11]]

     -- assemble the complete image while removing borders
     let pixels :: A.UArray Coord Bool
         pixels :: UArray Coord Bool
pixels =
           (Bool -> Bool -> Bool)
-> Bool -> (Coord, Coord) -> [(Coord, Bool)] -> UArray Coord Bool
forall (a :: * -> * -> *) e i e'.
(IArray a e, Ix i) =>
(e -> e' -> e) -> e -> (i, i) -> [(i, e')] -> a i e
A.accumArray (\Bool
_ Bool
x -> Bool
x) Bool
False (Int -> Int -> Coord
C Int
0 Int
0, Int -> Int -> Coord
C (Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
12) (Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
12))
             [ (Int -> Int -> Coord
C (Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
yyInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
xxInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1), Bool
True)
             | (C Int
yy Int
xx, (Int
_,Picture
cell)) <- Array Coord (Int, Picture) -> [(Coord, (Int, Picture))]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
A.assocs Array Coord (Int, Picture)
image
             , C Int
y Int
x <- Picture
cell
             , Int
y Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0, Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0, Int
y Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
9, Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
9 -- remove edges
             ]

     -- count occurrences of the snake in the pixel set
     let n :: Int
n = [()] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length
             [ ()
             | Picture
s <- Picture -> [Picture]
reorient Picture
snek
             , Coord
d <- UArray Coord Bool -> Picture
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [i]
A.indices UArray Coord Bool
pixels
             , (Coord -> Bool) -> Picture -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Coord
x -> UArray Coord Bool -> Coord -> Maybe Bool
forall (a :: * -> * -> *) e i (f :: * -> *).
(IArray a e, Ix i, Alternative f) =>
a i e -> i -> f e
arrIx UArray Coord Bool
pixels (Coord
x Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
+ Coord
d) Maybe Bool -> Maybe Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True) Picture
s
             ]

     -- cut all the snakes out of the picture
     Int -> IO ()
forall a. Show a => a -> IO ()
print ((Bool -> Bool) -> [Bool] -> Int
forall (f :: * -> *) a. Foldable f => (a -> Bool) -> f a -> Int
countBy Bool -> Bool
forall a. a -> a
id (UArray Coord Bool -> [Bool]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
A.elems UArray Coord Bool
pixels) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Picture -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Picture
snek)

place ::
  IntMap [(Int, Picture)] ->
  (Int, Picture) ->
  A.Array Coord (Int, Picture) {- ^ arranged image -}
place :: IntMap [(Int, Picture)]
-> (Int, Picture) -> Array Coord (Int, Picture)
place IntMap [(Int, Picture)]
em (Int, Picture)
start = Array Coord (Int, Picture)
board
  where
    bnds :: (Coord, Coord)
bnds = (Int -> Int -> Coord
C Int
0 Int
0, Int -> Int -> Coord
C Int
11 Int
11)
    board :: Array Coord (Int, Picture)
board = (Coord, Coord) -> [(Int, Picture)] -> Array Coord (Int, Picture)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
A.listArray (Coord, Coord)
bnds (Coord -> (Int, Picture)
pickTile (Coord -> (Int, Picture)) -> Picture -> [(Int, Picture)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Coord, Coord) -> Picture
forall a. Ix a => (a, a) -> [a]
A.range (Coord, Coord)
bnds)

    pickTile :: Coord -> (Int, Picture)
pickTile Coord
c
      | Coord
c Coord -> Coord -> Bool
forall a. Eq a => a -> a -> Bool
== Coord
origin = (Int, Picture)
start

      | Coord -> Int
coordRow Coord
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
         [(Int, Picture)] -> (Int, Picture)
forall a. HasCallStack => [a] -> a
head [ (Int
tileId, Picture
pic)
              | let (Int
nbId, Picture
nbPic) = Array Coord (Int, Picture)
board Array Coord (Int, Picture) -> Coord -> (Int, Picture)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
A.! Coord -> Coord
left Coord
c
              , (Int
tileId, Picture
pic) <- IntMap [(Int, Picture)]
em IntMap [(Int, Picture)] -> Int -> [(Int, Picture)]
forall a. IntMap a -> Int -> a
IntMap.! Picture -> Int
rightEdge Picture
nbPic
              , Int
tileId Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nbId ]

      | Bool
otherwise =
         [(Int, Picture)] -> (Int, Picture)
forall a. HasCallStack => [a] -> a
head [ (Int
tileId, Picture -> Picture
rotate Picture
pic)
              | let (Int
nbId, Picture
nbPic) = Array Coord (Int, Picture)
board Array Coord (Int, Picture) -> Coord -> (Int, Picture)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
A.! Coord -> Coord
above Coord
c
              , (Int
tileId, Picture
pic) <- IntMap [(Int, Picture)]
em IntMap [(Int, Picture)] -> Int -> [(Int, Picture)]
forall a. IntMap a -> Int -> a
IntMap.! Picture -> Int
bottomEdge' Picture
nbPic
              , Int
tileId Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nbId ]

topEdge, leftEdge, bottomEdge', rightEdge :: [Coord] -> Int
topEdge :: Picture -> Int
topEdge     Picture
xs = [Int] -> Int
fromBits [  Int
i | C Int
0 Int
i <- Picture
xs]
leftEdge :: Picture -> Int
leftEdge    Picture
xs = [Int] -> Int
fromBits [  Int
i | C Int
i Int
0 <- Picture
xs]
bottomEdge' :: Picture -> Int
bottomEdge' Picture
xs = [Int] -> Int
fromBits [Int
9Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i | C Int
9 Int
i <- Picture
xs]
rightEdge :: Picture -> Int
rightEdge   Picture
xs = [Int] -> Int
fromBits [  Int
i | C Int
i Int
9 <- Picture
xs]

fromBits :: [Int] -> Int
fromBits :: [Int] -> Int
fromBits = (Int -> Int -> Int) -> Int -> [Int] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Int -> Int -> Int
forall a. Bits a => a -> Int -> a
setBit Int
0