{-|
Module:         Day17
Description:    <https://adventofcode.com/2021/day/17 Day 17: Trick Shot>
-}
{-# LANGUAGE FlexibleContexts, OverloadedStrings, TypeFamilies #-}
module Day17 (day17) where

import Control.Monad (guard)
import Data.Ix (inRange)
import qualified Data.IntMap as IntMap ((!?), empty, insertWith)
import qualified Data.IntSet as IntSet (findMax, null, singleton, size, unions)
import Data.List (foldl', scanl')
import Data.List.NonEmpty (NonEmpty((:|)))
import Data.Maybe (catMaybes)
import Data.Semigroup (Max(..), Sum(..), sconcat)
import Data.String (IsString)
import Data.Text (Text)
import Data.Void (Void)
import Text.Megaparsec (MonadParsec, ParseErrorBundle, Token, Tokens, chunk, parse)
import Text.Megaparsec.Char.Lexer (decimal)

parser :: (MonadParsec e s m, IsString (Tokens s), Token s ~ Char, Num a, Ord a) => m ((a, a), (a, a))
parser :: m ((a, a), (a, a))
parser = do
    a
x0 <- Tokens s -> m (Tokens s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
chunk Tokens s
"target area: x=" m (Tokens s) -> m a -> m a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> m a
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m a
decimal
    a
x1 <- 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 a -> m a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> m a
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m a
decimal
    a
y0 <- Tokens s -> m (Tokens s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Tokens s -> m (Tokens s)
chunk Tokens s
", y=-" m (Tokens s) -> m a -> m a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> (a -> a
forall a. Num a => a -> a
negate (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m a
decimal)
    a
y1 <- 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 a -> m a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> (a -> a
forall a. Num a => a -> a
negate (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m a
decimal)
    ((a, a), (a, a)) -> m ((a, a), (a, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a
x0, a
y0), (a
x1, a
y1))

day17 :: Text -> Either (ParseErrorBundle Text Void) (Int, Int)
day17 :: Text -> Either (ParseErrorBundle Text Void) (Int, Int)
day17 Text
input = do
    ((Int
x0, Int
y0), (Int
x1, Int
y1)) <- Parsec Void Text ((Int, Int), (Int, Int))
-> String
-> Text
-> Either (ParseErrorBundle Text Void) ((Int, Int), (Int, Int))
forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
parse Parsec Void Text ((Int, Int), (Int, Int))
forall e s (m :: * -> *) a.
(MonadParsec e s m, IsString (Tokens s), Token s ~ Char, Num a,
 Ord a) =>
m ((a, a), (a, a))
parser String
"" Text
input
    let (Int
maxT, IntMap IntSet
dyHits) = ((Int, IntMap IntSet) -> Int -> (Int, IntMap IntSet))
-> (Int, IntMap IntSet) -> [Int] -> (Int, IntMap IntSet)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Int, IntMap IntSet) -> Int -> (Int, IntMap IntSet)
f (Int
0, IntMap IntSet
forall a. IntMap a
IntMap.empty) [Int
y0 .. -Int
y0]
        f :: (Int, IntMap IntSet) -> Int -> (Int, IntMap IntSet)
f (Int, IntMap IntSet)
k Int
dy = ((Int, IntMap IntSet) -> (Int, Int) -> (Int, IntMap IntSet))
-> (Int, IntMap IntSet) -> [(Int, Int)] -> (Int, IntMap IntSet)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Int -> (Int, IntMap IntSet) -> (Int, Int) -> (Int, IntMap IntSet)
g Int
dy) (Int, IntMap IntSet)
k ([(Int, Int)] -> (Int, IntMap IntSet))
-> [(Int, Int)] -> (Int, IntMap IntSet)
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] ([Int] -> [(Int, Int)]) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
y0) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 [Int
dy, Int
dy Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1..]
        g :: Int -> (Int, IntMap IntSet) -> (Int, Int) -> (Int, IntMap IntSet)
g Int
dy k :: (Int, IntMap IntSet)
k@(Int
maxT', IntMap IntSet
m) (Int
t, Int
y)
          | (Int, Int) -> Int -> Bool
forall a. Ix a => (a, a) -> a -> Bool
inRange (Int
y0, Int
y1) Int
y = (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
maxT' Int
t, (IntSet -> IntSet -> IntSet)
-> Int -> IntSet -> IntMap IntSet -> IntMap IntSet
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IntMap.insertWith IntSet -> IntSet -> IntSet
forall a. Semigroup a => a -> a -> a
(<>) Int
t (Int -> IntSet
IntSet.singleton Int
dy) IntMap IntSet
m)
          | Bool
otherwise = (Int, IntMap IntSet)
k
        (Max Int
maxDy, Sum Int
count) = NonEmpty (Max Int, Sum Int) -> (Max Int, Sum Int)
forall a. Semigroup a => NonEmpty a -> a
sconcat (NonEmpty (Max Int, Sum Int) -> (Max Int, Sum Int))
-> NonEmpty (Max Int, Sum Int) -> (Max Int, Sum Int)
forall a b. (a -> b) -> a -> b
$ (Int -> Max Int
forall a. a -> Max a
Max Int
0, Int -> Sum Int
forall a. a -> Sum a
Sum Int
0) (Max Int, Sum Int)
-> [(Max Int, Sum Int)] -> NonEmpty (Max Int, Sum Int)
forall a. a -> [a] -> NonEmpty a
:| do
            Int
dx <- [Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.25 :: Double) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
0.5)..Int
x1]
            let dys :: IntSet
dys = [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IntSet.unions ([IntSet] -> IntSet) -> [IntSet] -> IntSet
forall a b. (a -> b) -> a -> b
$ [Maybe IntSet] -> [IntSet]
forall a. [Maybe a] -> [a]
catMaybes
                  [ IntMap IntSet
dyHits IntMap IntSet -> Int -> Maybe IntSet
forall a. IntMap a -> Int -> Maybe a
IntMap.!? Int
t
                  | (Int
t, Int
x) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..Int
maxT] ([Int] -> [(Int, Int)]) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
x1) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ [Int
dx, Int
dx Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1..Int
1] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int]
forall a. a -> [a]
repeat Int
0
                  , (Int, Int) -> Int -> Bool
forall a. Ix a => (a, a) -> a -> Bool
inRange (Int
x0, Int
x1) Int
x
                  ]
            Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ IntSet -> Bool
IntSet.null IntSet
dys
            (Max Int, Sum Int) -> [(Max Int, Sum Int)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Max Int
forall a. a -> Max a
Max (Int -> Max Int) -> Int -> Max Int
forall a b. (a -> b) -> a -> b
$ IntSet -> Int
IntSet.findMax IntSet
dys, Int -> Sum Int
forall a. a -> Sum a
Sum (Int -> Sum Int) -> Int -> Sum Int
forall a b. (a -> b) -> a -> b
$ IntSet -> Int
IntSet.size IntSet
dys)
    (Int, Int) -> Either (ParseErrorBundle Text Void) (Int, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
maxDy Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
maxDy Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2, Int
count)