import Data.Fin
import Data.String
import Data.Vect
import Data.List
import Data.List1
import System.File.ReadWrite
import Debug.Trace

%default total

simple = """
  30373
  25512
  65332
  33549
  35390
  """

-- Parse a grid of numbers from a string
parseGrid : String -> Maybe  (n ** m ** Vect n (Vect m Nat))
parseGrid input =
  let inputLines = map unpack . filter ((> 0) . strLength) . map trim . lines $ input
  in do digits <- traverse (traverse parsePositive) . map (map Data.String.singleton) $ inputLines
        let height = length digits
        width <- map length . head' $ digits
        inner <- traverse (toVect width) digits
        outer <- toVect height inner
        Just (height ** width ** outer)

-- Variant of `take` and `drop` that works with fin
takeFin : {n: Nat} -> (m: Fin n) -> Vect n a -> Vect (finToNat m) a
takeFin FZ xs = []
takeFin (FS x) (y :: xs) = y :: takeFin x xs

dropFin : {n: Nat} -> (m: Fin n) -> Vect n a -> Vect (minus n (finToNat m)) a
dropFin FZ xs = xs
dropFin (FS x) (y :: xs) = dropFin x xs

-- Drop the last element of a Vect
dropLast : {n : Nat} ->  Vect (S n) a -> Vect n a
dropLast {n = 0} (x :: xs) = xs
dropLast {n = (S k)} (x :: xs) = x :: dropLast xs

tagTrace : Show a => String -> a -> a
tagTrace str =
  traceValBy (\x => str ++ ": " ++ show x)

-- Check if a tree is visible from the edge of the grid
checkVisible : {n: Nat} -> {m: Nat} -> (x: Fin n) -> (y: Fin m) -> Vect n (Vect m Nat) -> Bool
-- If either of the coordinates are zero or the maximum then the tree is on an edge and thus visible
checkVisible FZ y xs = True
checkVisible x FZ xs = True
checkVisible x@(FS _) y@(FS _) xs =
  if x == last || y == last
  then True
  else let horizontal = index x xs
           vertical = index y (transpose xs)
       in checkSingle y horizontal || checkSingle x vertical
  where checkSingle : {o: Nat} -> (z: Fin o) -> Vect o Nat -> Bool
        checkSingle FZ xs = True
        checkSingle z@(FS _) xs =
          let before = takeFin z xs
              after = drop 1 (toList (dropFin z xs))
              compValue = index z xs
          in all (< compValue) before || all (< compValue) after

showGrid : Vect n (Vect m Nat) -> String
showGrid xs =
  let lines = map (concatMap show) xs
  in joinBy "\n" (toList lines)

finPair : (m : Nat) -> (n : Nat) -> (Nat, Nat) -> Maybe (Fin m, Fin n)
finPair m n (x,y) =
  do a <- natToFin x m
     b <- natToFin y n
     Just (a, b)

testPart1 : IO ()
testPart1 =
  case parseGrid simple of
    Nothing => putStrLn "Failed to parse grid"
    Just (m ** n ** grid) =>
      do putStrLn "Grid: "
         putStrLn (showGrid grid)
         putStrLn ""
         case coords m n of
           Nothing => putStrLn "Hardcoded coords were invalid"
           Just cordList => do let testingFn = singleTest m n grid
                               traverse_ testingFn cordList
                               printLn "Tests Done"
  where coords : (m : Nat) -> (n : Nat) -> Maybe (List ((Fin m, Fin n), Bool))
        coords m n =
          let raw : List ((Nat, Nat), Bool)
              raw = [((1,1), True), ((1, 2), True), ((1, 3), False), ((2, 1), True), ((2, 2), False), ((2, 3), True)]
          in traverse (\(x, y) => map (\z => (z, y)) (finPair m n x)) raw
        singleTest : (m : Nat) -> (n : Nat) -> Vect m (Vect n Nat) -> ((Fin m, Fin n), Bool) -> IO ()
        singleTest m n grid ((x, y), result) =
          do let row = index x grid
             let val = index y row
             putStrLn $ "Coordinate: " ++ show (x,y) ++ " Value: " ++ show val ++ " Expected Visible: " ++ show result
             putStrLn $ "Calculated Visible: " ++ show (checkVisible x y grid) ++ "\n"

finRange : {n : Nat} -> List (Fin n)
finRange {n = 0} = []
finRange {n = (S k)} = forget (allFins k)

part1 : {m : Nat} -> {n : Nat} -> Vect m (Vect n Nat) -> Nat
part1 {m} {n} grid =
  let rows = finRange {n=m}
      cols = finRange {n=n}
      pairs = concatMap (\x => map (\y => (x, y)) cols) rows
  in length . Prelude.List.filter (\(x,y) => checkVisible x y grid) $ pairs


partial main : IO ()
main =
  do file <- readFile "input"
     case file of
       Left err => printLn err
       Right fileContents =>
         case parseGrid fileContents of
           Nothing => putStrLn "Failed to parse grid"
           Just (m ** n ** grid) =>
             do putStrLn $ "Part 1: " ++ show (part1 grid)