module Day4.Part1

import Data.String
import Data.Vect

%default total

--- TYPES

Board : Nat -> Nat -> Type
Board h w = Vect h (Vect w Char)

Pos : Nat -> Nat -> Type
Pos h w = (Fin h, Fin w)


data Direction = U | UR | R | DR | D | DL | L | UL

directions : List Direction
directions = [U, UR, R, DR, D, DL, L, UL]

--- PARSING

parseInput : String -> Maybe (h ** w ** Board h w)
parseInput inp =
  let ls@(row :: _) = unpack <$> lines inp
        | [] => Just (0 ** 0 ** [])
      width = length row
  in map (\l => (length l ** width ** fromList l))
       (traverse (toVect width) ls)

--- DATA

index : Pos h w -> Board h w -> Char
index (i, j) = index j . index i

predFin : Fin n -> Maybe (Fin n)
predFin FZ = Nothing
predFin (FS x) = Just $ weaken x

moveDir : {h, w : _} -> Direction -> Pos h w -> Maybe (Pos h w)
moveDir U (i, j) = (,j) <$> predFin i
moveDir UR (i, j) = (,) <$> predFin i <*> strengthen (FS j)
moveDir R (i, j) = (i,) <$> strengthen (FS j)
moveDir DR (i, j) = (,) <$> strengthen (FS i) <*> strengthen (FS j)
moveDir D (i, j) = (,j) <$> strengthen (FS i)
moveDir DL (i, j) = (,) <$> strengthen (FS i) <*> predFin j
moveDir L (i, j) = (i,) <$> predFin j
moveDir UL (i, j) = (,) <$> predFin i <*> predFin j

checkString : {h, w : _}
              -> Board h w -> String -> Pos h w -> Direction -> Bool
checkString board str = go (unpack str)
  where
    go : List Char -> Pos h w -> Direction -> Bool
    go [] _ _ = True
    go (c :: cs) pos dir =
      let Just next = moveDir dir pos
          | Nothing => False
      in index next board == c && go cs next dir


findChar : {h, w : _} -> Board h w -> Char -> List (Pos h w)
findChar board c = do
  i <- toList $ range {len=h}
  j <- toList $ range {len=w}
  guard (index (i, j) board == c)
  pure (i, j)

countString : {h, w : _} -> Board h w -> String -> Nat
countString board str with (strUncons str)
  _ | Nothing = 0
  _ | Just (c, str') =
    count (uncurry $ checkString board str')
      $ (,) <$> findChar board c <*> directions

--- SOLUTION

export
solution : String -> Maybe Nat
solution inp =
  let Just (_ ** _ ** board) = parseInput inp
      | Nothing => Nothing
  in Just $ countString board "XMAS"