advent12/src/advent12.hs

   1 import Data.Text (Text)
   2 import qualified Data.Text.IO as TIO
   3
   4 import Data.Void (Void)
   5
   6 import Text.Megaparsec hiding (State)
   7 import Text.Megaparsec.Char
   8 import qualified Text.Megaparsec.Char.Lexer as L
   9 import qualified Control.Applicative as CA
  10
  11 import Linear (V3(..), V1(..), (^+^), (^-^))
  12
  13 import qualified Data.Vector as V
  14
  15 class (Eq a) => NVec a where
  16     (^+^^) :: a -> a -> a
  17     (^-^^) :: a -> a -> a
  18     nvZero :: a
  19     nvSignum :: a -> a
  20     nvAbsSum :: a -> Integer
  21 instance NVec (V1 Integer) where
  22     x ^+^^ y  = x ^+^ y
  23     x ^-^^ y  = x ^-^ y
  24     nvZero = V1 0
  25     nvSignum (V1 x) = V1 (signum x)
  26     nvAbsSum (V1 x) = abs x
  27 instance NVec (V3 Integer) where
  28     x ^+^^ y  = x ^+^ y
  29     x ^-^^ y  = x ^-^ y
  30     nvZero = V3 0 0 0
  31     nvSignum (V3 x y z) = V3 (signum x) (signum y) (signum z)
  32     nvAbsSum (V3 x y z) = (abs x) + (abs y) + (abs z)
  33
  34
  35 data Planet a = Planet { _pos :: a , _vel :: a} deriving (Show, Eq, Ord)
  36 type Planets a = V.Vector (Planet a)
  37
  38
  39 main :: IO ()
  40 main = do
  41         text <- TIO.readFile "data/advent12.txt"
  42         let planetsT = successfulParse text
  43         let planets = enplanet planetsT
  44         -- print planets
  45         print $ part1 planets
  46         print $ part2 planets
  47
  48 part1 :: Planets (V3 Integer) -> Integer
  49 part1 = systemEnergy . head . drop 1000 . simulate
  50
  51 part2 :: Planets (V3 Integer) -> Integer
  52 part2 planets = period
  53     where planetDimensions = unzipPlanets planets
  54           simCounts = map countSimulate planetDimensions
  55           period = foldl lcm 1 simCounts
  56
  57
  58 enplanet :: (NVec a) => [a] -> Planets a
  59 enplanet = V.fromList . map mkPlanet
  60
  61 mkPlanet :: (NVec a) => a -> Planet a
  62 mkPlanet p = Planet {_pos = p, _vel = nvZero}
  63
  64
  65 unzipPlanets :: V.Vector (Planet (V3 Integer)) -> [V.Vector (Planet (V1 Integer))]
  66 unzipPlanets = dimensionSlice . V.map unzipPlanet
  67
  68 unzipPlanet :: Planet (V3 Integer) -> [Planet (V1 Integer)]
  69 unzipPlanet planet = map mkPlanet posVecs
  70     where posVecs = unzipVec $ _pos planet
  71
  72
  73 unzipVec :: V3 Integer -> [V1 Integer]
  74 unzipVec (V3 x y z) = [V1 x, V1 y, V1 z]
  75
  76 dimensionSlice :: (NVec a) => V.Vector [Planet a] -> [Planets a]
  77 dimensionSlice slicedPlanets = [sliceDim d | d <- [0..2]]
  78     where sliceDim d = V.map (!!d) slicedPlanets
  79
  80
  81
  82 simulate :: (NVec a) => Planets a -> [Planets a]
  83 simulate = iterate simulationStep
  84
  85 countSimulate :: (NVec a) => Planets a -> Integer
  86 countSimulate planets0 = go (simulationStep planets0) 1
  87     where go planets n
  88             | planets0 == planets = n
  89             | otherwise = go (simulationStep planets) (n + 1)
  90
  91
  92 simulationStep :: (NVec a) => Planets a -> Planets a
  93 simulationStep planets = planets''
  94     where   planets' = applyGravity planets
  95             planets'' = applyVelocity planets'
  96
  97
  98 gravity :: (NVec a) => a -> a
  99 gravity v = nvSignum v
 100
 101 applyGravity :: (NVec a) => Planets a -> Planets a
 102 applyGravity planets = V.map (applyGravityHere planets) planets
 103
 104 applyGravityHere :: (NVec a) => Planets a -> Planet a -> Planet a
 105 applyGravityHere planets here = V.foldl' updateGravity here planets
 106
 107 updateGravity :: (NVec a) => Planet a -> Planet a -> Planet a
 108 updateGravity here there = here { _vel = vel'}
 109     where   vel = _vel here
 110             vel' = vel ^+^^ gravity ((_pos there) ^-^^ (_pos here))
 111
 112 applyVelocity :: (NVec a) => Planets a -> Planets a
 113 applyVelocity = V.map applyVelocityHere
 114
 115 applyVelocityHere :: (NVec a) => Planet a  -> Planet a
 116 applyVelocityHere here = here {_pos = (_pos here) ^+^^ (_vel here)}
 117
 118
 119
 120
 121 -- absSum (Vec1 (V1 x)) = (abs x)
 122 -- absSum (Vec3 (V3 x y z)) = (abs x) + (abs y) + (abs z)
 123
 124 potentalEnergy planet = nvAbsSum $ _pos planet
 125 kineticEnergy planet = nvAbsSum $ _vel planet
 126 totalEnergy planet = (potentalEnergy planet) * (kineticEnergy planet)
 127
 128 systemEnergy = (V.foldl' (+) 0) . (V.map totalEnergy)
 129
 130
 131
 132 -- Parse the input file
 133 type Parser = Parsec Void Text
 134
 135 sc :: Parser ()
 136 sc = L.space (skipSome spaceChar) CA.empty CA.empty
 137 -- sc = L.space (skipSome (char ' ')) CA.empty CA.empty
 138
 139 lexeme  = L.lexeme sc
 140 integer = lexeme L.decimal
 141 signedInteger = L.signed sc integer
 142 symb = L.symbol sc
 143 equalP = symb "="
 144 commaP = symb ","
 145 identifierP = some alphaNumChar <* sc
 146 openBracketP = symb "<"
 147 closeBracketP = symb ">"
 148
 149 planetsP = many planetP
 150
 151 planetP = (between openBracketP closeBracketP) coordsP
 152
 153 coordsP = envector <$> (coordP `sepBy` commaP)
 154     where envector [x, y, z] = V3 x y z
 155 coordP = identifierP *> equalP *> signedInteger
 156
 157
 158 -- successfulParse :: Text -> [Vec]
 159 successfulParse :: Text -> [V3 Integer]
 160 successfulParse input =
 161         case parse planetsP "input" input of
 162                 Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 163                 Right planets -> planets