module Main (main) where

import Data.List (groupBy, nub, sort, sortBy, subsequences)
import Data.Maybe (catMaybes)

-- | Part 1
data Antenna = Antenna Char (Int, Int) deriving (Show, Eq)

coords :: Antenna -> (Int, Int)
coords (Antenna _ (x, y)) = (x, y)

findAntennas :: [[Char]] -> [Antenna]
findAntennas grid =
  let possibilities = [maybeAntenna (row, col) | row <- [0 .. length grid - 1], col <- [0 .. length (grid !! row) - 1]]
   in catMaybes possibilities
  where
    maybeAntenna :: (Int, Int) -> Maybe Antenna
    maybeAntenna (row, col)
      | (grid !! row) !! col == '.' = Nothing
      | otherwise = Just (Antenna ((grid !! row) !! col) (row, col))

calculateAntiNodes :: (Int, Int) -> (Int, Int) -> [(Int, Int)]
calculateAntiNodes (xa, ya) (xb, yb) =
  let (dx, dy) = (xb - xa, yb - ya)
   in [(xa - dx, ya - dy), (xb + dx, yb + dy)]

allAntiNodes :: [(Int, Int)] -> [(Int, Int)]
allAntiNodes antennas = concatMap (uncurry calculateAntiNodes . (\[a, b] -> (a, b))) $ filter ((2 ==) . length) $ subsequences antennas

outOfBounds :: [[a]] -> (Int, Int) -> Bool
outOfBounds g (row, col) = row < 0 || row >= length g || col < 0 || col >= length (g !! row)

part1 :: String -> IO ()
part1 s =
  let grid = lines s
      antennas = findAntennas grid
      groupedAntennas = (groupBy (\(Antenna a _) (Antenna b _) -> a == b) . sortBy (\(Antenna a _) (Antenna b _) -> compare a b)) antennas
      antiNodes = concatMap (filter (not . outOfBounds grid) . allAntiNodes . map coords) groupedAntennas
   in print $ length $ nub antiNodes

-- | Part 2
multiplesAN :: [[a]] -> (Int, Int) -> (Int, Int) -> [(Int, Int)]
multiplesAN grid (xa, ya) (xb, yb) =
  let (dx, dy) = (xb - xa, yb - ya)
   in (takeWhile (not . outOfBounds grid) [(xa - dx * i, ya - dy * i) | i <- [0 ..]]) ++ takeWhile (not . outOfBounds grid) [(xb + dx * i, yb + dy * i) | i <- [0 ..]]

allMultiplesAN :: [[a]] -> [(Int, Int)] -> [(Int, Int)]
allMultiplesAN grid antennas = foldr (\[a, b] l -> multiplesAN grid a b ++ l) [] $ filter ((2 ==) . length) $ subsequences antennas

part2 :: String -> IO ()
part2 s =
  let grid = lines s
      antennas = findAntennas grid
      groupedAntennas = (groupBy (\(Antenna a _) (Antenna b _) -> a == b) . sortBy (\(Antenna a _) (Antenna b _) -> compare a b)) antennas
      antiNodes = concatMap (allMultiplesAN grid . map coords) groupedAntennas
   in print $ length $ nub antiNodes

-- | Main
main :: IO ()
main = readFile "resources/day8.txt" >>= part2