import Data.Array
import Data.List

main = do
  input <- readFile "input"
  let grid = makeGrid $ lines input
--  print $ part1 grid
  print $ part2 grid

part1 grid = sum $ map (\r -> (regionArea r) * (regionPerimeter r)) $ findRegions grid

part2 grid = sum $ map (\r -> (regionArea r) * (regionSides r)) $ findRegions grid

regionSides r =
  let r1 = spans r
      r2 = spans $ map (\(a,b) -> (b,a)) r
  in countSides (map risingEdges r1) +
     countSides (map risingEdges r2) +
     countSides (map fallingEdges r1) +
     countSides (map fallingEdges r2)

countSides edges = count [] edges
  where count _ [] = 0
        count prev (e:rest) =
          (length $ filter (`notElem` prev) e) + (count e rest)

spans = map (map snd) . groupBy (\a b -> fst a == fst b) . sort

risingEdges xs = (head xs):(map snd . filter (\(a,b) -> (a+1) /= b) $ zip xs (tail xs))
fallingEdges = map negate . risingEdges . reverse . map negate

regionArea r = length r
regionPerimeter r = 4*(length r) - (sum $ map nneighbours r)
  where nneighbours (i,j) = length $ filter (`elem` r) [(i-1,j),(i+1,j),(i,j-1),(i,j+1)]

findRegions grid = doregion [] (indices grid)
  where doregion regions [] = regions
        doregion regions (p:next)
          | seen p regions = doregion regions next
          | otherwise = doregion ((findRegion grid p):regions) next
        seen p regions = or $ map (elem p) regions

findRegion grid start = expand [] [start]
  where plant = grid!start
        expand seen [] = seen
        expand seen (p:next)
          | p `elem` seen = expand seen next
          | otherwise = let n = filter (\i -> grid!i == plant)
                                $ getNeighbours grid p
                        in expand (p:seen) (next ++ n)

getNeighbours arr (i,j) = filter (inRange (bounds arr)) [(i-1,j),(i+1,j),(i,j-1),(i,j+1)]

makeGrid grid = array bounds elements
  where n = length grid
        m = length (head grid)
        bounds = ((0,0), (n-1,m-1))
        elements = [((i,j), (grid!!i!!j)) |
                    i <- [0..n-1], j <- [0..m-1]]