import Data.Char
import Data.List
import Data.Maybe

type Height = Int
type Coo = (Int, Int)
type Map = [[Height]]

data Trail = Point Coo Height Trail | Nowhere deriving Show

(!?) :: [[a]] -> Coo -> Maybe a
mp !? (row, col)
  | row < 0                   = Nothing
  | row >= length mp          = Nothing
  | col < 0                   = Nothing
  | col >= (length $ head mp) = Nothing
  | otherwise                 = Just (mp !! row !! col)

consume :: [Char] -> Map
consume file =
  map (map (digitToInt)) $ lines file

trailheads2 :: [Coo] -> Int -> Int -> [Height] -> [Coo]
trailheads2 acc _ _ [] =
  reverse acc
trailheads2 acc r c (h:rest)
  | h == 0 = trailheads2 ((r, c):acc) r (c + 1) rest
  | otherwise = trailheads2 acc r (c + 1) rest

trailheads1 :: [Coo] -> Int -> Map -> [Coo]
trailheads1 acc _ [] =
  acc
trailheads1 acc r (row:rest) =
  trailheads1 (acc ++ heads) (r + 1) rest
  where heads = trailheads2 [] r 0 row

trailheads :: Map -> [Coo]
trailheads = trailheads1 [] 0

trails2 :: Map -> Trail -> [Trail]
trails2 mp p@(Point (r,c) h _) =
  map pointify $
    catMaybes $
    map oneup
    [nxt (-1, 0), nxt (0, -1), nxt (0, 1), nxt (1, 0)]
  where
    pointify c = Point c (h+1) p
    oneup (nr, nc, Just nh) = if nh == h + 1 then Just (nr, nc) else Nothing
    oneup (_, _, Nothing) = Nothing
    nxt (dr, dc) = (r+dr, c+dc, mp !? (r+dr, c+dc))

trails1 :: Map -> [Trail] -> [Trail] -> [Trail]
trails1 _ acc@((Point _ 9 _):_) [] =
  acc
trails1 mp acc [] =
  trails1 mp [] acc
trails1 mp acc (th@(Point _ 9 _):rest) =
  trails1 mp (th:acc) rest
trails1 mp acc (th:rest) =
  trails1 mp (acc ++ trails2 mp th) rest

trails :: Map -> Coo -> [Trail]
trails mp th =
  trails1 mp [] [Point th 0 Nowhere]

score :: [Trail] -> Int
score th =
  length $ nub $ map untrail th
  where
    untrail (Point c _ _) = c

rating :: [Trail] -> Int
rating = length

part1 :: Map -> Int
part1 mp =
  sum $ map (score . (trails mp)) $ trailheads mp

part2 :: Map -> Int
part2 mp =
  sum $ map (rating . (trails mp)) $ trailheads mp

main = do
  file <- readFile "day10.input"
  let mp = consume file
  putStrLn ("Part 1: " ++ (show $ part1 mp))
  putStrLn ("Part 2: " ++ (show $ part2 mp))