module Day05 where
import Data.Map (Map, insertWith, empty, elems)
import Data.List.Split (splitOn)

data Line = Line (Int,Int) (Int,Int) deriving Show
type Field = Map (Int,Int) Int
strToLine :: String -> Line
strToLine s = Line (x1,y1) (x2,y2) where
    [x1,y1] = map read $ splitOn "," p1
    [x2,y2] = map read $ splitOn "," p2
    [p1,"->",p2] = words s

coverPoint :: (Int,Int) -> Field -> Field
coverPoint point = Data.Map.insertWith (+) point 1

coverLine :: Line -> Field -> Field
coverLine (Line (x1,y1) (x2,y2)) field
    | y1 == y2 = let points = map (\x -> (x,y1)) [xMin..xMax] in foldr coverPoint field points 
    | x1 == x2 = let points = map (\y -> (x1,y)) [yMin..yMax] in foldr coverPoint field points
    | otherwise = field
    where xMin = min x1 x2; xMax = max x1 x2; yMin = min y1 y2; yMax = max y1 y2

coverLineDiag :: Line -> Field -> Field
coverLineDiag (Line (x1,y1) (x2,y2)) field
    | y1 == y2 = let points = map (\x -> (x,y1)) [xMin..xMax] in foldr coverPoint field points 
    | x1 == x2 = let points = map (\y -> (x1,y)) [yMin..yMax] in foldr coverPoint field points
    | x1 == xMin && y1 == yMin = let points = zip [x1..x2] [y1..y2] in foldr coverPoint field points
    | x1 == xMin && y1 == yMax = let points = zip [x1..x2] (reverse [y2..y1]) in foldr coverPoint field points
    | otherwise = coverLineDiag (Line (x2,y2) (x1,y1)) field
    where xMin = min x1 x2; xMax = max x1 x2; yMin = min y1 y2; yMax = max y1 y2

main :: IO ()
main = do
    lines <- getContents >>= ((return . map strToLine) . lines)
    
    putStr "part 1: "
    let covered = foldr coverLine Data.Map.empty lines
    let num_overlap = length $ filter (>=2) $ Data.Map.elems covered
    print num_overlap

    putStr "part 2: "
    let covered = foldr coverLineDiag Data.Map.empty lines
    let num_overlap = length $ filter (>=2) $ Data.Map.elems covered
    print num_overlap

    return ()