module Main where

import Data.Char
import Data.List
import Data.Maybe

import Input


parseVec :: String -> [Int]
parseVec = map read . tokens (3 :: Int)
  where
    tokenf c = isDigit c || c == '-'
    token str = span tokenf (snd (break tokenf str))
    tokens 0 _ = []
    tokens n str = uncurry (:) (fmap (tokens (n - 1)) (token str))

data Moon1D = Moon1D { mPos :: {-# UNPACK #-} !Int, mVel :: {-# UNPACK #-} !Int } deriving (Show, Eq)
newtype System1D = System1D [Moon1D] deriving (Show, Eq)

gravity1D :: [Int] -> [Int]
gravity1D coords = [length (filter (> x) coords) - length (filter (< x) coords) | x <- coords]

-- -- If there are MANY moons, the below will be faster (n log n instead of n^2)
-- gravity1D :: [Int] -> [Int]
-- gravity1D inCoords =
--     let (coords, backIndices) = unzip (sortOn fst (zip inCoords [0::Int ..]))
--         cumBeforeOf cds = map fst (scanl (\(nb, ct) (a, b) ->
--                                               if a < b then (nb + ct, 1) else (nb, ct + 1))
--                                          (0, 1) (zip cds (tail cds)))
--         cumBefore = cumBeforeOf coords
--         cumAfter = reverse (cumBeforeOf (map negate (reverse coords)))
--         grav = zipWith (-) cumAfter cumBefore
--     in map snd (sortOn fst (zip backIndices grav))

physics1D :: System1D -> System1D
physics1D (System1D moons) =
    let gravity = gravity1D (map mPos moons)
        moons1 = zipWith (\(Moon1D p v) g -> Moon1D p (v + g)) moons gravity
        moons2 = [Moon1D (p + v) v | Moon1D p v <- moons1]
    in System1D moons2

energy :: [System1D] -> Int
energy systems =
    let moons = transpose [moons1d | System1D moons1d <- systems]
        moonEnergy moon1ds = sum (map (abs . mPos) moon1ds) * sum (map (abs . mVel) moon1ds)
    in sum (map moonEnergy moons)

main :: IO ()
main = do
    initSystems <- map System1D . transpose . map (map (\p -> Moon1D p 0) . parseVec) <$> getInput 12

    print $ energy [iterate physics1D s !! 1000 | s <- initSystems]

    -- Note: any cycle must contain the starting position, because the
    -- simulation step is uniquely reversible. -- Credits Bert
    let reps = [1 + fromJust (findIndex (== initSystem) (tail (iterate physics1D initSystem)))
               | initSystem <- initSystems]
    print $ foldl1 lcm reps