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

-- 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))

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

data State : Nat -> Type where
  MkState : (head : (Int, Int)) -> (tails : Vect n (Int, Int)) -> State n
%name State state, state1, state2

emptyState : {n : Nat} -> State n
emptyState = MkState (0,0) (replicate _ (0,0))

Show (State n) where
  show (MkState head tails) = "Head: " ++ show head ++ " Tails: " ++ show tails

stateTails : State n -> Vect n (Int, Int)
stateTails (MkState head xs) = xs

sign : Int -> Int
sign i = if i == 0 then 0 else div i (abs i)

mapP : (a -> b) -> (a, a) -> (b, b)
mapP f (x, y) = (f x, f y)

zipP : (a -> b -> c) -> (a, a) -> (b, b) -> (c, c)
zipP f (x, y) (z, w) = (f x z, f y w)

-- Increment the tail
incTail : (head : (Int, Int)) -> (tail : (Int, Int)) -> (Int, Int)
incTail head tail =
  let diff = zipP (-) head tail
      absDiff = mapP abs diff
      signs = mapP sign diff
  in if touching head tail
     then tail
     else zipP (+) tail signs

incTails : {n: Nat} -> (head : (Int, Int)) -> (tails : Vect n (Int, Int)) -> Vect n (Int, Int)
incTails head [] = []
incTails head (x :: xs) =
  let newTail = incTail head x in
  newTail :: incTails newTail xs

moveHead' : {n: Nat} -> Direction -> State n -> State n
moveHead' direction (MkState head tails) =
  let newHead = applyDirection direction head in
  MkState newHead (incTails newHead tails)

applyMotion' : {m: Nat} -> (motion : Motion) -> State m -> (State m, Vect (S (distance motion)) (State m))
applyMotion' motion state = helper (breakdown motion) state
  where helper : {m: Nat} -> (directions : Vect n Direction) -> State m -> (State m, Vect (S n) (State m))
        helper [] state = (state, [state])
        helper (x :: xs) state =
          let newState = moveHead' x state
              (resState, rest) = helper xs newState
          in (resState, state :: rest)

applyMotions' : {m : Nat} -> (motions : List Motion) -> State m -> (State m, List (State m))
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
  """

complex = """
  R 5
  U 8
  L 8
  D 3
  R 17
  D 10
  L 25
  U 20
  """

testPart : Nat -> (input : String) -> IO ()
testPart Z input = putStrLn "No Tail!"
testPart (S m) input =
  case parseMotions input of
    Nothing => putStrLn "Failed to parse motions"
    Just motions =>
      let (state, states) = applyMotions' motions (emptyState {n = (S m)})
          tails = map (last . stateTails) states
          uniqueTails = nub tails
      in do putStrLn "Movements:"
            traverse_ printLn motions
            putStrLn "\nSteps:"
            traverse_ printLn states
            putStrLn "\nUnique Tails:"
            printLn $ length uniqueTails

testPart1 : IO ()
testPart1 = testPart 1 simple

testPart2 : IO ()
testPart2 =
  do testPart 9 simple
     testPart 9 complex

part : Nat -> String -> Maybe Nat
part 0 str = ?part_rhs_0
part (S k) input =
  case parseMotions input of
    Nothing => Nothing
    Just motions =>
      let (state, states) = applyMotions' motions (emptyState {n = (S k)})
          tails = map (last . stateTails) states
      in Just . length . nub $ tails


part1 : String -> Maybe Nat
part1 input = part 1 input

part2 : String -> Maybe Nat
part2 input = part 9 input

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 =>
            do putStrLn ("Part 1: " ++ show part1Count)
               case part2 contents of
                 Nothing => printLn "Error in part 2"
                 Just part2Count => putStrLn ("Part 2: " ++ show part2Count)