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

%default total

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

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

takeWhilePlus: Ord a => (a -> Bool) -> List a -> List a
takeWhilePlus f [] = []
takeWhilePlus f (x :: []) = [x]
takeWhilePlus f (x :: rest) =
  if f x
  then x :: takeWhilePlus f rest
  else [x]

senicScore : {m : Nat} -> {n : Nat} -> (x: Fin m) -> (y: Fin n) -> Vect m (Vect n Nat) -> Nat
senicScore FZ y xs = 0
senicScore x FZ xs = 0
senicScore x@(FS _) y@(FS _) xs =
  if x == last || y == last
  then 0
  else let horizontal = index x xs
           vertical = index y (transpose xs)
       in scoreLinear y horizontal * scoreLinear x vertical
  where scoreLinear : {o : Nat} -> (z: Fin o) -> Vect o Nat -> Nat
        scoreLinear FZ xs = 0
        scoreLinear z@(FS _) xs =
          let before = reverse . toList . takeFin z $ xs
              after = drop 1 (toList (dropFin z xs))
              compValue = index z xs
              beforeVisible = length (takeWhilePlus (< compValue) before)
              afterVisible = length (takeWhilePlus (< compValue) after)
          in beforeVisible * afterVisible

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

part2 : {m : Nat} -> {n: Nat} -> Vect m (Vect n Nat) -> Nat
part2 {m} {n} grid =
  let rows = finRange {n=m}
      cols = finRange {n=n}
      pairs = concatMap (\x => map (\y => (x, y)) cols) rows
  in foldl max 0 . map (\(x, y) => senicScore 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)
                putStrLn $ "Part 2: " ++ show (part2 grid)