# Pastebin wNWecKmJ
{-# LANGUAGE BangPatterns    #-}
{-# LANGUAGE DeriveGeneric   #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections   #-}
module Day4 where

import           Advent.AoC
import           Control.DeepSeq             (NFData (..))
import           Control.Lens
import           Control.Monad               (replicateM)
import           Control.Parallel.Strategies (parMap, rseq)
import           Data.IntMap                 (IntMap)
import qualified Data.IntMap.Strict          as IntMap
import           Data.List                   (sortOn)
import           Data.Maybe                  (catMaybes)
import           Data.Semigroup              (Max (..), Min (..))
import           Data.Tuple                  (swap)
import           GHC.Generics                (Generic (..))
import           Text.Megaparsec             (many, sepBy)
import           Text.Megaparsec.Char        (space)
import qualified Text.Megaparsec.Char.Lexer  as L


data Board = Board {
  _colFill   :: IntMap Int,
  _rowFill   :: IntMap Int,
  _boardNums :: [((Int,Int), Int)]
  } deriving (Show, Generic)

instance NFData Board

makeLenses ''Board

data Input = Input {
  drawn  :: IntMap Int,
  boards :: [Board]
  } deriving (Show, Generic)

instance NFData Input

winners :: Input -> [(Int, Int)]
winners Input{..} = catMaybes $ parMap rseq go boards
  where
    go Board{_boardNums=[]} = Nothing
    go !b@Board{_boardNums=(((x,y),n):xs)} =
      let b' = b & colFill . ix y -~ 1
                 & rowFill . ix x -~ 1
      in if 0 `elem` _colFill b' || 0 `elem` _rowFill b'
         then Just (n, (drawn IntMap.! n) * sum (((drawn IntMap.!) . snd) <$> xs))
         else go b'{_boardNums=xs}

parseInput :: Int -> Parser Input
parseInput size = do
  ns <- IntMap.fromList . flip zip [0..] <$> lexeme (L.decimal `sepBy` ",")
  bseq <- catMaybes . parMap rseq (traverse sequenceA) <$> many (board ns)
  pure $ Input (IntMap.fromList . fmap swap . IntMap.assocs $ ns) (Board fill fill <$> bseq)
    where
      fill = IntMap.fromList [(k,size) | k <- [1..size]]
      coords = [(x,y) | y <- [1..size], x <- [1..size]]
      lexeme = L.lexeme space
      board :: IntMap Int -> Parser [((Int,Int), Maybe Int)]
      board nm = sortOn snd . zip coords <$> replicateM (size*size) ((`IntMap.lookup` nm) <$> lexeme L.decimal)

getInput :: FilePath -> IO Input
getInput = parseFile (parseInput 5)

part1 :: Input -> Int
part1 = snd . minimum . winners

part2 ::Input ->  Int
part2 = snd . maximum . winners

biggun :: IO (Int,Int)
biggun = ex . foldMap b . winners <$> parseFile (parseInput 15) "input/day4.big"
  where
    b x = (Min x, Max x)
    ex (Min (_,x), Max(_,y)) = (x,y)