import Data.List
import Data.Array
import Data.Array.ST
import Control.Monad
import Control.Monad.ST
import Debug.Trace

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


part1 grid =
  let start = head $ findOnGrid grid 'S'
      path = findPath grid start
      connections = sum $ map (numShortcuts 2 100) (init (tails path))
  in connections

part2 grid =
  let start = head $ findOnGrid grid 'S'
      path = findPath grid start
      connections = sum $ map (numShortcutsX 20 100) (init (tails path))
  in connections

-- save at least n steps, shortcut dist is <= d
numShortcutsX d n l =
  let start = head l
      rest = zip [1..] (tail l)
  in length $ filter (\(i,end) -> let md = manhattanDist start end
                                  in md <= d && md+n <= i) rest

-- save at least n steps, shortcut dist is d
numShortcuts d n l =
  let start = head l
      rest = drop (d+n) l
  in length $ filter (\end -> manhattanDist start end == d) rest

manhattanDist (a1,a2) (b1,b2) = abs(b1-a1) + abs(b2-a2)

findPath grid pos = find [] pos
  where find path pos =
          let n = getNeighbours grid pos
              next = filter (\p -> grid!p /= '#' && p `notElem` path) n
          in case next of
               [] -> reverse (pos:path)
               [p] -> find (pos:path) p

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

findOnGrid grid char = map fst . filter ((== char) . snd) . assocs $ grid

type Grid = Array (Int,Int) Char
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]]