module Main where

import Data.Char
import Data.List

import Input


data Vec = Vec Int Int Int deriving (Show)

instance Num Vec where  -- elementwise
    Vec a b c + Vec x y z = Vec (a + x) (b + y) (c + z)
    Vec a b c * Vec x y z = Vec (a * x) (b * y) (c * z)
    abs (Vec a b c) = Vec (abs a) (abs b) (abs c)
    signum (Vec a b c) = Vec (signum a) (signum b) (signum c)
    negate (Vec a b c) = Vec (negate a) (negate b) (negate c)
    fromInteger n = let n' = fromInteger n in Vec n' n' n'

coordinates :: Vec -> [Int]
coordinates (Vec a b c) = [a, b, c]

fromCoordinates :: [Int] -> Vec
fromCoordinates [a, b, c] = Vec a b c
fromCoordinates _ = error "Invalid length of list in fromCoordinates"

parseVec :: String -> Vec
parseVec = fromCoordinates . 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 Moon = Moon { mPos :: Vec, mVel :: Vec } deriving (Show)
data System = System [Moon] deriving (Show)

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))

physics :: System -> System
physics (System moons) =
    let gravity = map fromCoordinates . transpose . map gravity1D . transpose . map (coordinates . mPos) $ moons
        moons1 = zipWith (\(Moon p v) g -> Moon p (v + g)) moons gravity
        moons2 = [Moon (p + v) v | Moon p v <- moons1]
    in System moons2

class Energy a where energy :: a -> Int
instance Energy Vec where energy = sum . coordinates . abs
instance Energy Moon where energy (Moon p v) = energy p * energy v
instance Energy System where energy (System moons) = sum (map energy moons)

main :: IO ()
main = do
    initSys <- System . map (\line -> Moon (parseVec line) 0) <$> getInput 12
    print $ energy ((iterate physics initSys) !! 1000)