import Data.String
import Data.Vect
import Data.List1
import System.File.ReadWrite

%default total

data Direction = Up | Right | Left | Down
%name Direction direction, direction2, direction3

Show Direction where
  show Up = "Up"
  show Right = "Right"
  show Left = "Left"
  show Down = "Down"

parseDirection : String -> Maybe Direction
parseDirection input =
  case trim input of
    "U" => Just Up
    "R" => Just Right
    "L" => Just Left
    "D" => Just Down
    x => Nothing

invert : Direction -> Direction
invert Up = Down
invert Right = Left
invert Left = Right
invert Down = Up

-- Apply a direction to a pair of Ints
applyDirection : Direction -> (Int, Int) -> (Int, Int)
applyDirection Up (x, y) = (x, y + 1)
applyDirection Right (x, y) = (x + 1, y)
applyDirection Left (x, y) = (x - 1, y)
applyDirection Down (x, y) = (x, y - 1)

data Motion = Move Direction Nat
%name Motion motion, motion2, motion3

Show Motion where
  show (Move direction count) = show direction ++ " " ++ show count

parseMotion : String -> Maybe Motion
parseMotion input =
  let components = split (== ' ') input in
    case forget components of
      [x, y] => do direction <- parseDirection x
                   ammount <- parsePositive y
                   pure (Move direction ammount)
      _ => Nothing

parseMotions : String -> Maybe (List Motion)
parseMotions input =
  let inputLines = lines input
  in traverse parseMotion inputLines

-- Get the distance moved by a motion
distance : Motion -> Nat
distance (Move _ i) = i

-- Break a motion down into a list of moves-by-one
breakdown : (motion : Motion) -> Vect (distance motion) Direction
breakdown (Move direction 0) = []
breakdown input@(Move direction (S k)) = direction :: breakdown (assert_smaller input (Move direction k))

-- The state of a head and tail
data State: Type where
  MkState : (head : (Int, Int)) -> (tail : (Int, Int)) -> State
%name State state, state1, state2

Show State where
  show (MkState head tail) = "Head: " ++ show head ++ " Tail: " ++ show tail

emptyState : State
emptyState = MkState (0,0) (0,0)

tail : State -> (Int, Int)
tail (MkState head x) = x

touching : (Int, Int) -> (Int, Int) -> Bool
touching (x, y) (z, w) =
  let (u, v) = (abs (x - z), abs (y - w)) in
  u <= 1 && v <= 1

-- Move the head in a particular direction, causing the tail to follow if needed
moveHead : Direction -> State -> State
moveHead direction (MkState head tail) =
  let newHead = applyDirection direction head in
  if touching newHead tail
    then MkState newHead tail
    else let newTail = applyDirection (invert direction) newHead in
         MkState newHead newTail

-- Apply a movement to the head, returning the sequence of tail positions, including the starting
-- position
applyMotion : (motion : Motion) -> State -> (State, Vect (S (distance motion)) State)
applyMotion motion state = helper (breakdown motion) state
  where helper : (directions : Vect n Direction) -> State -> (State, Vect (S n) State)
        helper [] state = (state, [state])
        helper (x :: xs) state =
          let newState = moveHead x state
              (resState, rest) = helper xs newState
          in (resState, state :: rest)

-- Apply all the motions in a list to the head, returning the full sequence of positions
applyMotions : (motions : List Motion) -> State -> (State, List State)
applyMotions [] state = (state, [state])
applyMotions (x :: xs) state =
  let (newState, motions) = applyMotion x state
      (outputState, rest) = applyMotions xs newState
  in (outputState, (toList motions) ++ rest)

simple = """
  R 4
  U 4
  L 3
  D 1
  R 4
  D 1
  L 5
  R 2
  """

testPart1 : IO ()
testPart1 =
  case parseMotions simple of
    Nothing => putStrLn "Failed to parse motions"
    Just motions =>
      let (state, states) = applyMotions motions emptyState
          tails = map tail states
          uniqueTails = nub tails
      in do putStrLn "Movements:"
            traverse_ printLn motions
            putStrLn "\nSteps:"
            traverse_ printLn states
            putStrLn "\nUnique Tails:"
            printLn $ length uniqueTails

part1 : String -> Maybe Nat
part1 input =
  case parseMotions input of
    Nothing => Nothing
    Just motions =>
      let (state, states) = applyMotions motions emptyState
          tails = map tail states
      in Just . length . nub $ tails

partial main : IO ()
main =
  do file <- readFile "input"
     case file of
       Left err => printLn err
       Right contents =>
        case part1 contents of
          Nothing => printLn "Error in part 1"
          Just part1Count =>
            putStrLn ("Part 1: " ++ show part1Count)