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