{-# LANGUAGE LambdaCase, TupleSections, ViewPatterns #-}
module Solve where

import System.Console.ANSI
import Control.Concurrent
import Control.Monad

import System.IO.Unsafe

import Data.Map.Strict (Map(..))
import qualified Data.Map.Strict as M

import Data.Maybe
import Data.List

type Vec2 = (Int, Int)

data Choice = TurnLeft | Straight | TurnRight deriving (Show, Eq, Enum)

data Cart = Cart { position :: Vec2
                 , velocity :: Vec2
                 , choice :: [Choice] } deriving Eq

instance Show Cart where
    show (Cart pos velocity (choice:choices)) = "Cart@" ++ show pos ++ " going " ++ show velocity ++ " nextchoice:" ++ show choice

type Track = Map (Int, Int) TrackPiece

data TrackPiece = Vertical
                | Horizontal
                | Bidirectional
                | LeftDown
                | RightUp
                deriving Show

showPiece :: TrackPiece -> Char
showPiece Vertical      = '|'
showPiece Horizontal    = '-'
showPiece Bidirectional = '+'
showPiece LeftDown      = '\\'
showPiece RightUp       = '/'

showTrack :: Track -> String
showTrack m = show $ M.toList m

showCart :: Cart -> Char
showCart (Cart _ ( 1,  0) _) = '>'
showCart (Cart _ (-1,  0) _) = '<'
showCart (Cart _ ( 0,  1) _) = 'v'
showCart (Cart _ ( 0, -1) _) = '^'

data State = State { carts :: [Cart]
                   , track :: Track } deriving Show

parseTrackPiece :: Char -> Maybe TrackPiece
parseTrackPiece '|'  = Just Vertical
parseTrackPiece '-'  = Just Horizontal
parseTrackPiece '+'  = Just Bidirectional
parseTrackPiece '\\' = Just LeftDown
parseTrackPiece '/'  = Just RightUp
parseTrackPiece c   | c `elem` "<>^v" = Just $ trackFromCart $ cartFromC (0, 0) c
                    | otherwise = Nothing

trackFromCart :: Cart -> TrackPiece
trackFromCart (Cart _ (_, 0) _) = Horizontal
trackFromCart (Cart _ (0, _) _) = Vertical

cartFromC :: (Int, Int) -> Char -> Cart
cartFromC pos = \case '^' -> Cart pos ( 0, -1) choices
                      'v' -> Cart pos ( 0,  1) choices
                      '>' -> Cart pos ( 1,  0) choices
                      '<' -> Cart pos (-1,  0) choices
    where choices = cycle $ enumFromTo TurnLeft TurnRight

move :: Cart -> Cart
move (Cart (x, y) (xvel, yvel) c) = Cart (x + xvel, y + yvel) (xvel, yvel) c

rotLeft (x, y) = (y, -x)
rotRight (x, y) = (-y, x)

modifyCart :: TrackPiece -> Cart -> Cart
modifyCart Vertical      c = move c
modifyCart Horizontal    c = move c
modifyCart Bidirectional (Cart p vel (choice:choices)) = case choice of
    TurnLeft -> move (Cart p (rotLeft vel) choices)
    Straight -> move (Cart p vel choices)
    TurnRight -> move (Cart p (rotRight vel) choices)
modifyCart LeftDown      (Cart (x, y) (xv, yv) choice) = move $ Cart (x, y) (yv, xv) choice
modifyCart RightUp       (Cart (x, y) (xv, yv) choice) = move $ Cart (x, y) (-yv, -xv) choice


getTrackUnderCart :: Cart -> Track -> TrackPiece
getTrackUnderCart (Cart (x, y) _ _) track = track M.! (x, y)

stepCart :: Cart -> Track -> Cart
stepCart cart track = modifyCart piece cart
    where piece = getTrackUnderCart cart track

parseInput :: String -> State
parseInput str = State (parseCarts str) $ parseLines str

parseCarts :: String -> [Cart]
parseCarts = concatMap (uncurry getCartsLine) . zip [0..] . lines

getCartsLine :: Int -> String -> [Cart]
getCartsLine y cs = [ cartFromC (x, y) c
                    | (x, c) <- zip [0..] cs
                    , c `elem` "<>^v" ]

parseLine :: Int -> String -> Track
parseLine y = M.fromList . mapMaybe (\(x, char) -> ((x, y),) <$> parseTrackPiece char) . zip [0..]

parseLines :: String -> Track
parseLines = foldr M.union M.empty . zipWith parseLine [0 ..] . lines

printTrack :: (Int, Int) -> Track -> IO ()
printTrack (xo, yo) track = forM_ (M.assocs track) $ \((x, y), piece) -> do
                    setCursorPosition (y + yo) (x + xo)
                    putChar $ showPiece piece

printCarts :: (Int, Int) -> [Cart] -> IO ()
printCarts (xo, yo) carts = forM_ carts $ \cart@(Cart (x, y) (_, _) _) -> do
                                setCursorPosition (y + yo) (x + xo)
                                putChar $ showCart cart


printTrackSubset :: (Int, Int) -> Int -> (Int, Int) -> Track -> IO ()
printTrackSubset (xo, yo) sz (tx, ty) track = forM_ (M.assocs track) $ \((x, y), piece) -> do
    guard (x > (tx + sz) || x < (tx - sz) || y > (ty + sz) || y < (ty - sz))
    setCursorPosition (y + yo) (x + xo)
    putChar $ showPiece piece

printState :: (Int, Int) -> State -> IO ()
printState (x, y) (State c t) = printTrack (x, y) t >> printCarts (x, y) c

hasCollided :: Cart -> [Cart] -> Bool
hasCollided c = isJust . getCollision c


getCollision :: Cart -> [Cart] -> Maybe Cart
getCollision c = find (collide c)

collide :: Cart -> Cart -> Bool
collide a b = position a == position b

-- step' :: ([Cart], State) -> ([Cart], State)
-- step' (collisions, State (sortPositions position -> carts) track) = (newCollisions, State newCarts track)
--     where (newCollisions, newCartsUnfiltered) = foldr (reductor' track) (collisions, []) $ tails carts
--           newCarts = filter (`notElem` newCollisions) newCartsUnfiltered

stepMe :: ([Cart], State) -> ([Cart], State)
stepMe (collisions, State (sortPositions position -> carts) track) = (newCollisions, State newCarts track)
    where (newCollisions, newCarts) = go track carts (collisions, [])
          go track [] s = s
          go track (((`stepCart` track) -> sCart):carts) (currCollisions, newCarts)
            | hasCollided sCart newCarts =
                go track carts (sCart : currCollisions, filter (not . collide sCart) newCarts)
            | otherwise = case getCollision sCart carts of
                Nothing -> go track carts (currCollisions, sCart : newCarts)
                Just other -> go track (filter (not . collide sCart) carts) (other : sCart : collisions, newCarts)
                -- Nothing -> go track carts (collisions, sCart : newCarts)


-- step'' :: ([Cart], State) -> ([Cart], State)
-- step'' (collisions, State (sortPositions position -> carts) track) = (newCollisions, State newCarts track)
--     where (newCollisions, newCarts) = reductor'' track carts (collisions, [])

-- reductor'' :: Track -> [Cart] -> ([Cart], [Cart]) -> ([Cart], [Cart])
-- reductor'' track [] (collisions, newCarts) = (collisions, newCarts)
-- reductor'' track (cart:carts) (collisions, newCarts)
--     | hasCollided (stepCart cart track) newCarts =
--     reductor'' track carts (stepCart cart track : collisions, filter (not . collide (stepCart cart track)) newCarts)
--     | hasCollided (stepCart cart track) carts =
--     case getCollision (stepCart cart track) carts of
--         Just other -> reductor'' track carts (other : stepCart cart track : collisions, newCarts)
--     | otherwise = reductor'' track carts (collisions, stepCart cart track : newCarts)

-- reductor' :: Track -> [Cart] -> ([Cart], [Cart]) -> ([Cart], [Cart])
-- reductor' track [] (collisions, newCarts) = (collisions, newCarts)
-- reductor' track (cart:carts) (collisions, newCarts) | hasCollided (stepCart cart track) newCarts =
--                                                       (stepCart cart track : collisions, filter (not . collide cart) newCarts)
--                                                     | hasCollided (stepCart cart track) carts =
--                                                       case getCollision (stepCart cart track) carts of
--                                                         Just other -> (other : stepCart cart track : collisions, newCarts)
--                                                     | otherwise = (collisions, stepCart cart track : newCarts)


-- reductor :: Track -> [Cart] -> Cart -> ([Cart], [Cart]) -> ([Cart], [Cart])
-- reductor track carts v (collisions, newCarts) | hasCollided (stepCart v track) carts =
--                                                 (stepCart v track : collisions, filter (not . collide v) $ filter (not . collide (stepCart v track)) newCarts)
--                                               | otherwise = (collisions, stepCart v track : newCarts)

-- step :: ([Cart], State) -> ([Cart], State)
-- step (collisions, State (sortPositions position -> carts) track) = (newCollisions, State newCarts track)
--     where (newCollisions, newCarts) = foldr (reductor track carts) (collisions, []) carts


sortPositions :: (a -> Vec2) -> [a] -> [a]
sortPositions f = concat
                . sortOn ((snd . f) . head)
                . groupBy (\a b -> (fst . f) a == (fst . f) b)
                . sortOn (fst . f)

solve :: String -> IO ()
solve path = do
    f <- readFile path

    clearScreen

    let s = parseInput f

    let start = ([], s) :: ([Cart], State)

    let loop (coll, st) i c = do
        -- when (i == 100) $ do
            -- clearScreen
            -- printState (0, 8) st
            -- setCursorPosition 500 0

        -- putStr $ show (length . carts $ st) ++ "; "

        -- print coll
        -- print $ head $ carts st
        -- threadDelay 1000000

        when (c /= length coll) $ do
             let collision = head coll

             printTrackSubset (10, 10) 5 (position collision) (track st)

        threadDelay 1000000




        -- when (not . null $ ) $
            -- print coll

        if length (carts st) > 1
        then
            loop (stepMe (coll, st)) (succ i) (length coll)
        else print (fst . stepMe $ (coll, st))


    loop ([], s) 0 0

    pure ()