module Task2 (task2) where

import Common (Position, readWorld, multiMapFromList, pairPermutations)
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Control.Monad (join)
import Flow

task2 :: String -> Int
task2 input = case readWorld input of
  ((width, height), antennas) -> multiMapFromList antennas |> Map.elems |> map antinodesOfFrequency |> join |> Set.fromList |> length
    where
      antinodesOfFrequency :: [Position] -> [Position]
      antinodesOfFrequency nodes = pairPermutations nodes >>= antinodesOfPair

      antinodesOfPair :: (Position, Position) -> [Position]
      antinodesOfPair (a@(ax, ay), b@(bx, by)) =
        let diffX = bx - ax
            diffY = by - ay
        in castInWorld a (-diffX, -diffY) ++ castInWorld b (diffX, diffY)

      castInWorld (sx, sy) (ox, oy) =
        map (\f -> (sx + ox * f, sy + oy * f)) [0..] |> takeWhile isInWorld

      isInWorld :: Position -> Bool
      isInWorld (x, y)
        | x < 0 = False
        | x >= width = False
        | y < 0 = False
        | y >= height = False
        | otherwise = True