import Data.List
import Data.List.Split
import Linear.V3

main :: IO ()
main = do
  moons <- parseContent <$> readFile "input"
  let velocity = replicate 4 [0, 0, 0]
  print $ part1 (toV3 moons, toV3 velocity)
  print $ part2 (moons, velocity)

part1 :: ([V3 Int], [V3 Int]) -> Int
part1 x = energy $ iterate step x !! 1000

--I’m sure there is a better way to do this
toV3 :: [[Int]] -> [V3 Int]
toV3 = map (\[x, y, z] -> V3 x y z)

parseContent :: String -> [[Int]]
parseContent = map (map read . splitOn ",") . lines . removeJunk
  where
    removeJunk xs = [x | x <- xs, x `notElem` " <>xyz="]

step :: (Num a, Eq a) => ([a], [a]) -> ([a], [a])
step (moons, vel) = (zipWith (+) moons newVel, newVel)
  where
    dVel = gravity moons
    newVel = zipWith (+) vel dVel
    gravity xs = [sum [signum $ x - y | x <- xs, y /= x] | y <- xs]

energy :: ([V3 Int], [V3 Int]) -> Int
energy (x, y) = sum $ zipWith (*) (geten x) (geten y)
  where
    geten = map (sum . abs)

findPeriod :: (Num a, Eq a) => ([a], [a]) -> ([a], [a]) -> Int -> Int
findPeriod x a n =
  if x' == a
    then n
    else findPeriod x' a (n + 1)
  where
    x' = step x

part2 :: (Num a, Eq a) => ([[a]], [[a]]) -> Int
part2 (moons, vel) = lcm' periods
  where
    m = transpose moons
    v = transpose vel
    periods = zipWith (curry findPeriod') m v
    findPeriod' x = findPeriod x x 1

lcm' :: [Int] -> Int
lcm' xs = product $ zipWith (^) nums maxElems
  where
    nums = nub $ concat decomp
    decomp = map decomposition xs
    maxElems = [maximum $ map (length . elemIndices x) decomp | x <- nums]

--hyper optimized prime decomposition
decomposition :: Int -> [Int]
decomposition 1 = []
decomposition x =
  let n =
        if x `mod` 2 == 0
          then 2
          else 3
      maxiter = (floor . (sqrt :: Double -> Double) . fromIntegral) x
      findprim candidate current
        | (candidate <= maxiter) && (x `mod` candidate /= 0) =
          findprim (next current) (current + 1)
        | otherwise = candidate
      m = findprim n 1
   in if m <= maxiter
        then m : decomposition (x `div` m)
        else [x]
  where
    next k = k * 4 - k `div` 2 * 2 + 1