{-# LANGUAGE MultiWayIf #-}
module Main where

import Data.Functor
import Data.Matrix hiding (trace)

data Pos = Pos { ch :: Char
               , ma :: Bool
               }

instance Show Pos where
  show (Pos '.' True) = show '#'
  show (Pos ch _)     = show ch

parse :: String -> Matrix Pos
parse = fromLists . (((`Pos` False) <$>) <$>) . lines

positionsOfChar :: Char -> Matrix Pos -> [(Int,Int)]
positionsOfChar c s = zip [1..] (toLists s) >>= \(y,l) -> zip [1..] l >>= \(x,c') ->
  if c == ch c' then pure (y,x) else mempty

getMarked :: Matrix Pos -> [(Int,Int)]
getMarked s = zip [1..] (toLists s) >>= \(y,l) -> zip [1..] l >>= \(x,c) ->
  if ma c then pure (y,x) else mempty

availableChars :: [Char]
availableChars = ['a'..'z'] ++ ['A'..'Z'] ++ ['0'..'9']

line :: (Int,Int) -> (Int,Int) -> (Int,Int) -> Bool
line (y1,x1) (y2,x2) (y3,x3) =
  if | x2 == x1 && x3 == x2 -> True
     | x2 == x1 || x3 == x2 || x1 == x3 -> False
     | otherwise ->
       ((y2-y1) `div` (x2-x1)) == ((y2-y3) `div` (x2-x3)) &&
       ((y2-y1) `div` (x2-x1)) == ((y1-y3) `div` (x1-x3))

combo :: Eq a => [a] -> [(a,a)]
combo ps = [(a,b) | a <- ps, b <- ps, a /= b]


getAntinodePos :: (Int,Int) -> [(Int,Int)] -> [(Int,Int)]
getAntinodePos (i,j) ps =
  let inbounds (y,x) = and [0 < y, y <= i, 0 < x, x <= j] in
    [ c | (p@(y,x),p'@(y',x')) <- combo ps,
                 c <- [(y + (y-y'),x+(x-x')),
                       (y - (y-y'),x+(x-x'))],
                 inbounds c,
                 line c p p']

getAntinodePos' :: (Int,Int) -> [(Int,Int)] -> [(Int,Int)]
getAntinodePos' (i,j) ps =
  let inbounds (y,x) = and [0 < y, y <= i, 0 < x, x <= j]
      a = combo ps >>= \((y,x),(y',x')) ->
        takeWhile inbounds [ (y + step * (y - y'), x + step * (x - x')) | step <- [1..]]
      b = combo ps >>= \((y,x),(y',x')) ->
        takeWhile inbounds [ (y + step * (y - y'), x + step * (x - x')) | step <- [0,-1..]] in
    a ++ b

mark :: Matrix Pos -> [(Int,Int)] -> Matrix Pos
mark = foldr (\p@(y,x) s -> let c = getElem y x s in setElem (c {ma = True}) p s)

solve1 :: Matrix Pos -> Int
solve1 s = length . getMarked $
  foldr (flip mark . getAntinodePos (nrows s, ncols s) . (`positionsOfChar` s)) s availableChars

solve2 :: Matrix Pos -> Int
solve2 s = length . getMarked $
  foldr (flip mark . getAntinodePos' (nrows s, ncols s) . (`positionsOfChar` s)) s availableChars

main :: IO ()
main = readFile "inputs/8" <&> parse >>= \i ->
       print (solve1 i) >>
       print (solve2 i)