#! /usr/bin/env stack
-- stack --resolver lts-12.20 --install-ghc script
{-# LANGUAGE ViewPatterns, BangPatterns, LambdaCase, TemplateHaskell #-}
module Main where

import Control.Applicative
import Control.Lens
import Control.Monad.State
import Data.Function (on)
import Data.Maybe (fromJust)
import Data.Tuple (swap)
import Data.Void
import System.Environment (getArgs)

import qualified Data.Array as A
import qualified Data.List  as L
import qualified Data.Map   as M
import qualified Data.Set   as S

import Debug.Trace

type CartID = Int
type Position = (Int, Int)

data Direction = DUp | DRight | DDown | DLeft
   deriving (Eq,Ord,Enum,Bounded,Show)

data Bias = TurnLeft | GoStraight | TurnRight
   deriving (Eq,Ord,Enum,Bounded,Show)

data Cart = Cart { _cartPosition  :: Position
                 , _cartDirection :: Direction
                 , _cartBias      :: Bias
                 } deriving (Eq,Ord,Show)
makeLenses ''Cart

type Carts  = M.Map CartID Cart

data Track = TEmpty | THorizontal | TVertical
           | TCurvedFW | TCurvedBW | TIntersection
   deriving (Eq,Ord,Enum,Bounded,Show)

type Tracks = A.Array Position Track

data Part = Part1 | Part2 deriving (Eq,Ord,Enum,Bounded,Show)

main = getArgs >>= \case
   []    -> interact $ solve Part1 . parseInput
   ["1"] -> interact $ solve Part1 . parseInput
   ["2"] -> interact $ solve Part2 . parseInput
   _     -> putStrLn "Invalid arguments."

parseInput :: String -> (Tracks, Carts)
parseInput input = (tracks, carts)
   where
      max_x = (maximum $ map length $ lines input) - 1
      max_y = (length $ lines input) - 1
      tracks = A.accumArray (flip const) TEmpty ((0,0),(max_x,max_y)) $
         [ ((x, y), case c of { '\\' -> TCurvedBW
                              ; '/'  -> TCurvedFW
                              ; '-'  -> THorizontal
                              ; '<'  -> THorizontal
                              ; '>'  -> THorizontal
                              ; '|'  -> TVertical
                              ; '^'  -> TVertical
                              ; 'v'  -> TVertical
                              ; 'V'  -> TVertical
                              ; '+'  -> TIntersection
                              })
         | (y, ln) <- zip [0..] $ lines input
         , (x, c)  <- zip [0..] ln
         , c /= ' '
         ]
      carts = M.fromList $ zip [0..] $
         [ Cart (x, y) direction TurnLeft
         | (y, ln) <- zip [0..] $ lines input
         , (x, c)  <- zip [0..] ln
         , direction <- case c of { '^' -> [DUp]
                                  ; '>' -> [DRight]
                                  ; 'v' -> [DDown]
                                  ; '<' -> [DLeft]
                                  ; _   -> []
                                  }
         ]

compareOrder :: Carts -> CartID -> CartID -> Ordering
compareOrder carts = compare `on` (\i -> carts ^? at i . to fromJust . cartPosition . to swap)

solve :: Part -> (Tracks, Carts) -> String
solve part (tracks, initialCarts) = result
   where
      atPosition :: Position -> Track
      atPosition pt
         | A.inRange (A.bounds tracks) pt = tracks A.! pt
         | otherwise = TEmpty
      result = either id absurd $ flip evalStateT initialCarts $ forever $ do
         let sortCartIDs carts = L.sortBy (compareOrder carts) (M.keys carts)
         sortedCarts <- sortCartIDs <$> get
         forM_ sortedCarts $ \cartID -> do
            removed <- M.notMember cartID <$> get
            unless removed $ do
               -- move cart
               direction <- use $ unsafeIx cartID . cartDirection 
               modifying (ix cartID . cartPosition) $ \(x,y) ->
                  case direction of
                     DUp    -> (x, y-1)
                     DRight -> (x+1, y)
                     DDown  -> (x, y+1)
                     DLeft  -> (x-1, y)
               -- turn cart
               position <- use $ unsafeIx cartID . cartPosition
               case atPosition position of
                  TIntersection -> do
                     bias <- use $ unsafeIx cartID . cartBias
                     ix cartID . cartDirection .= case (direction, bias) of
                        (DUp,    TurnLeft)   -> DLeft
                        (DUp,    GoStraight) -> DUp
                        (DUp,    TurnRight)  -> DRight
                        (DRight, TurnLeft)   -> DUp
                        (DRight, GoStraight) -> DRight
                        (DRight, TurnRight)  -> DDown
                        (DDown,  TurnLeft)   -> DRight
                        (DDown,  GoStraight) -> DDown
                        (DDown,  TurnRight)  -> DLeft
                        (DLeft,  TurnLeft)   -> DDown
                        (DLeft,  GoStraight) -> DLeft
                        (DLeft,  TurnRight)  -> DUp
                     ix cartID . cartBias .= case bias of
                        TurnLeft   -> GoStraight
                        GoStraight -> TurnRight
                        TurnRight  -> TurnLeft
                  TCurvedFW -> do  -- '/'
                     ix cartID . cartDirection .= case direction of
                        DUp    -> DRight
                        DRight -> DUp
                        DDown  -> DLeft
                        DLeft  -> DDown
                  TCurvedBW -> do -- '\'
                     ix cartID . cartDirection .= case direction of
                        DUp    -> DLeft
                        DRight -> DDown
                        DDown  -> DRight
                        DLeft  -> DUp
                  _ -> pure ()
               -- check for collisions
               here <- map fst . filter (\(i,c) -> c^.cartPosition == position) . M.assocs <$> get
               when (length here > 1) $ do
                  when (part == Part1) $ lift $ Left $ show position
                  forM_ here $ modify . M.delete
         when (part == Part2) $ do
            remaining <- toListOf (each . cartPosition) <$> get
            --traceM $ show remaining
            when (null (drop 1 remaining)) $ lift $ Left $ show remaining

unsafeIx :: (Functor f, Ixed t) => Index t
         -> Over (->) f t t (IxValue t) (IxValue t)
unsafeIx = unsafeSingular . ix