module Part2 where

import Commons
import Data.Array (Ix (inRange), bounds, (!))
import qualified Data.Map as M


type CompactPaths = M.Map ((Int, Int), Direction) ((Int, Int), Direction, Int)


invertDirection :: Direction -> Direction
invertDirection North = South
invertDirection East = West
invertDirection South = North
invertDirection West = East

findCompactPaths :: Trails -> [(Int, Int)] -> (Int, Int) -> Direction -> Int -> CompactPaths
findCompactPaths _ [(yF, xF)] (y, x) d n =
  let iD = invertDirection d
  in M.insert ((yF, xF), South) ((y, x), d, n + 1) $ M.singleton ((y, x), iD) ((yF, xF), North, n + 1)
findCompactPaths tr ((yF, xF): t) (y, x) d n
  | let t = tr ! (yF, xF) in t /= Empty && t /= Rock =
      let Slope nD = tr ! (yF, xF)
          iD = invertDirection d
          iND = invertDirection nD
      in if n <= 2 then findCompactPaths tr t (yF, xF) nD 1
         else M.union (M.insert ((yF, xF), nD) ((y, x), d, n + 1) $ M.singleton ((y, x), iD) ((yF, xF), iND, n + 1))
                      $ findCompactPaths tr t (yF, xF) nD 1
  | otherwise = findCompactPaths tr t (y, x) d (n + 1)

getCompactPaths :: Trails -> [[(Int, Int)]] -> CompactPaths
getCompactPaths _ [] = M.empty
getCompactPaths tr (h: t) = let (y, x) = head h
                            in M.union (getCompactPaths tr t) $ findCompactPaths tr h (y + 1, x) North 0


cleanNext :: Trails -> CompactPaths -> [(Int, Int)] -> [((Int, Int), Direction)] -> [((Int, Int), Direction, Int)]
cleanNext _ _ _ [] = []
cleanNext tr cp hist ((ch, dh): t) =
  if inRange (bounds tr) ch && (tr ! ch /= Rock) && notElem ch hist
  then let actualNext = M.lookup (ch, dh) cp
       in case actualNext of
               Just v -> v: cleanNext tr cp hist t
               _      -> (ch, dh, 1): cleanNext tr cp hist t
  else cleanNext tr cp hist t

getNext :: Trails -> CompactPaths -> [(Int, Int)] -> Direction -> [((Int, Int), Direction, Int)]
getNext tr cp ((y, x): t) North = cleanNext tr cp t [((y - 1, x), North), ((y, x - 1), West), ((y, x + 1), East)]
getNext tr cp ((y, x): t) East = cleanNext tr cp t [((y - 1, x), North), ((y, x + 1), East), ((y + 1, x), South)]
getNext tr cp ((y, x): t) South = cleanNext tr cp t [((y + 1, x), South), ((y, x - 1), West), ((y, x + 1), East)]
getNext tr cp ((y, x): t) West = cleanNext tr cp t [((y - 1, x), North), ((y, x - 1), West), ((y + 1, x), South)]

getNextTrails' :: Trails -> CompactPaths -> (Int, Int) -> (Int, Int) -> [((Int, Int), Int)] ->
                  [((Int, Int), Direction, Int)] -> [[((Int, Int), Int)]]
getNextTrails' _ _ _ _ _ [] = []
getNextTrails' tr cp f s hist ((ch, dh, nh): t)
  | ch == f   = ((ch, nh): hist): getNextTrails' tr cp f ch hist t
  | otherwise = getNextTrails tr cp f ch ((ch, nh): hist) dh ++ getNextTrails' tr cp f ch hist t

getNextTrails :: Trails -> CompactPaths -> (Int, Int) -> (Int, Int) -> [((Int, Int), Int)] -> Direction ->
                 [[((Int, Int), Int)]]
getNextTrails tr cp f s hist d = let histWithoutDist = map fst hist
                                     next = getNext tr cp histWithoutDist d
                                 in getNextTrails' tr cp f s hist next

getAllTrails :: Trails -> CompactPaths -> [[((Int, Int), Int)]]
getAllTrails tr cp = let ((yMin, xMin), (yMax, xMax)) = bounds tr
                     in getNextTrails tr cp (yMax, xMax - 1) (yMin - 1, xMin + 1) [((yMin - 1, xMin + 1), 0)] South