content: {-# LANGUAGE ViewPatterns #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE PartialTypeSignatures #-}
  {-# LANGUAGE InstanceSigs #-}
  
  
  
  import Algebra.Geometric.Cl3
  import Data.Massiv.Array as A
  import Control.Monad (replicateM_,when)
  import Data.Time.Clock.POSIX (getPOSIXTime)
  
  import Foreign.Ptr (Ptr, castPtr, plusPtr)
  import Foreign.Storable (Storable(..))
  
  import Data.IORef (IORef(),newIORef,modifyIORef',readIORef)
  
  
  main :: IO ()
  main = do
    arrayProgrammingAlgorithm
    bitTwiddlingAlgorithm
    
  arrayProgrammingAlgorithm :: IO ()
  arrayProgrammingAlgorithm = 
    do tStart <- getPOSIXTime
       n <- return (50000000 :: Int)
       planetsPos :: Array S Ix1 Cl3_V3 <- fromListsM Seq planetspos  -- generate immutible Massiv array of the planets position
       planetsVel :: Array S Ix1 Cl3_V3 <- fromListsM Seq planetsvel  -- generate immutible Massiv array of the planets velocity
       planetsMass :: Array S Ix1 Cl3_R <- fromListsM Seq planetsmass  -- generate immutible Massiv array of the planets mass
       mutablePos <- thaw planetsPos  -- copy to a mutible Massiv array of the planets position
       mutableVel <- thaw planetsVel  -- copy to a mutible Massiv array of the planets velocity
       let (Sz1 c) = size planetsVel  -- Size of the velocity array
           c' = (c^(2 :: Int) - c) `div` 2  -- size of the nChoose2With array
       mutableScratch <- newMArray (Sz1 c') (toCl3_V3 $ V3 0 0 0) -- nChoose2With sized scratch pad
       mutableScratch2D <- newMArray (Sz (c :. (c-1))) (toCl3_V3 $ V3 0 0 0) -- 2D scratch pad for holding acceleration components
       offsetMomentum mutableVel planetsMass -- modify the Sun's velocity such that the total momentum is zero
       printEnergy mutablePos mutableVel planetsMass
       replicateM_ n (advance mutablePos mutableVel planetsMass mutableScratch mutableScratch2D)
       printEnergy mutablePos mutableVel planetsMass
       tEnd <- getPOSIXTime
       print $ show (tEnd - tStart) ++ " (sec) for " ++ show n ++ " iterations." -- Relative time measurement
  
  -- | 'bitTwiddlingAlgorithm' with Ϟ symbolizing the bit twiddling algorithm versions
  bitTwiddlingAlgorithm :: IO ()
  bitTwiddlingAlgorithm = 
    do tStart <- getPOSIXTime
       n <- return (50000000 :: Int)
       (solarSystem,sz) <- offsetMomentumϞ [sun,jupiter,saturn,uranus,neptune] -- returns a zero momentum solar system
       printEnergyϞ solarSystem sz
       replicateM_ n (advanceϞ solarSystem sz)
       printEnergyϞ solarSystem sz
       tEnd <- getPOSIXTime
       print $ show (tEnd - tStart) ++ " (sec) for " ++ show n ++ " iterations." -- Relative time measurement
  
  offsetMomentumϞ :: [Body] -> IO (SolarSystem,Int) -- A SolarSystem and the Size : size: sz :: Int
  offsetMomentumϞ bodies = do
    initBodies :: Array S Ix1 Body <- A.fromListsM Seq bodies
    -- mutate the solar system so that the sun's velocity gives zero total momentum
    let velocity :: Array D Ix1 Cl3 = A.map (fromCl3_V3.vel) initBodies
        mass :: Array D Ix1 Cl3 = A.map (fromCl3_R.mas) initBodies
    ss :: SolarSystem <- thaw initBodies
    let sz = totalElem.sizeOfMArray $ ss
    bs <- readM ss 0
    let sunV = fromCl3_V3.vel $ bs
        sunM = fromCl3_R.mas $ bs
    writeM ss 0 (bs{vel = toCl3_V3 (sunV - ((velocity `dot_V3` mass) / sunM))})
    return (ss,sz)
  
  
  printEnergyϞ :: SolarSystem -> Int -> IO ()
  printEnergyϞ ss sz = do
    acc :: IORef Cl3 <- newIORef (R 0)
    ssFroz <- freeze Seq ss
    iforIO_ ssFroz $ \idx bi ->
      do let bi_pos :: Cl3 = fromCl3_V3.pos $ bi
             bi_vel :: Cl3 = fromCl3_V3.vel $ bi
             bi_mas :: Cl3 = fromCl3_R.mas $ bi
             ke :: Cl3 = 0.5 * bi_mas * (toR $ bi_vel^(2::Int))
         modifyIORef' acc (+ ke)  -- accumulate the kenetic energy for each body
         -- now calcualate the potential energy
         let szMinus1 = sz-1
         when (idx < szMinus1) $ do
           brest <- extractM (idx+1) (Sz1 (szMinus1-idx)) ssFroz
           forIO_ brest $ \bj ->
             do let bj_pos :: Cl3 = fromCl3_V3.pos $ bj
                    bj_mas :: Cl3 = fromCl3_R.mas $ bj
                    distance :: Cl3 = sqrt.toR $ (bi_pos - bj_pos)^(2::Int)
                    pe_ij :: Cl3 = (bi_mas * bj_mas) / distance
                modifyIORef' acc (\ac -> ac - pe_ij)
    energy <- readIORef acc
    print energy
  
  advanceϞ :: SolarSystem -> Int -> IO ()
  advanceϞ ss sz = do
    updateVelocityϞ ss sz
    updatePositionϞ ss
  
  updateVelocityϞ :: SolarSystem -> Int -> IO ()
  updateVelocityϞ ss sz = do
    ssFroz <- freeze Seq ss
    iforIO_ ssFroz $ \idx bi ->
      when (idx < sz-1) $ do
        let bi_pos :: Cl3 = fromCl3_V3.pos $ bi
            bi_mas :: Cl3 = fromCl3_R.mas $ bi
        brest <- extractM (idx+1) (Sz1 (sz-idx-1)) ssFroz
        iforIO_ brest $ \jdx bj ->
          do let bj_pos :: Cl3 = fromCl3_V3.pos $ bj
                 bj_mas :: Cl3 = fromCl3_R.mas $ bj
                 (absMag,dposUnit) = abssignum $ bi_pos - bj_pos
                 mag = dt / ((absMag)^(2::Int))
                 aij = negate $ dposUnit * bj_mas * mag
                 aji = dposUnit * bi_mas * mag
             modify_ ss (modVel aij) idx
             modify_ ss (modVel aji) (idx+jdx+1)
  
  modVel :: PrimMonad m => Cl3 -> Body -> m Body
  modVel a bod = do
    let vbod = fromCl3_V3 $ vel bod
    return bod{vel = toCl3_V3 $ vbod + a}
  
  updatePositionϞ :: SolarSystem -> IO ()
  updatePositionϞ ss = do
    iforPrimM_ ss $ \idx b ->
      do let p :: Cl3 = fromCl3_V3.pos $ b
             v :: Cl3 = fromCl3_V3.vel $ b
         writeM ss idx b{pos = toCl3_V3 (p + dt * v)}
          
  
  advance :: MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix1 Cl3_V3 -> Array S Ix1 Cl3_R -> MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix2 Cl3_V3 -> IO ()
  advance sysPos sysVel sysMass scratchPad scratchPad2D =
    do updateVelocity sysPos sysVel sysMass scratchPad scratchPad2D
       updatePosition sysPos sysVel
  
  updateVelocity :: MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix1 Cl3_V3 -> Array S Ix1 Cl3_R -> MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix2 Cl3_V3 -> IO ()
  updateVelocity sysPos sysVel sysMass scratchPad scratchPad2D =
    do -- Calculate the inverse square distance and unit vector between the bodies
       frozPos <- freeze Seq sysPos
       let ps :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozPos
           scaledPij = nChoose2With (\ i j -> let (a,s) = abssignum $ i - j in (recip (a^(2 :: Int))) * s) ps -- probably should use a computeSource here, because scaledPij is used twice, or verify streams are still working.
       computeInto scratchPad (A.map toCl3_V3 scaledPij)  -- Compute result to the scratch pad
       frozScratch <- freeze Seq scratchPad  -- Freeze the result for future use
       -- Scale the previous inverse square distance by the mass of the other body
       let scratch :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozScratch
           ms :: Array D Ix1 Cl3 = A.map fromCl3_R sysMass
           upper = nChoose2With (\ _ j -> negate j) ms !*! scratch
           lower = nChoose2With (\ i _ -> i) ms !*! scratch
       computeInto scratchPad (A.map toCl3_V3 upper)
       frozScratchUpper <- freeze Seq scratchPad
       computeInto scratchPad (A.map toCl3_V3 lower)
       frozScratchLower <- freeze Seq scratchPad
       -- Load each body acceleration contribution into a 2D array 
       frozVel <- freeze Seq sysVel
       let vs :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozVel
           rows = elemsCount vs
       iforPrimM_ scratchPad2D (\ix _ -> write_ scratchPad2D ix (genAcc rows frozScratchUpper frozScratchLower ix))
       frozScratch2D <- freeze Seq scratchPad2D
       --  Sum each row to calcuate the total acceleration and then update the velocity
       let accComponents :: Array D Ix2 Cl3 = A.map fromCl3_V3 frozScratch2D
           as = makeArray Seq (size vs) (\ ix -> foldlS (+) (V3 0 0 0) (accComponents !> ix))
           vs' = vs !+! (dt *. as)
       computeInto sysVel (A.map toCl3_V3 vs')
  
  genAcc :: Int -> Array S Ix1 Cl3_V3 -> Array S Ix1 Cl3_V3 -> Ix2 -> Cl3_V3
  genAcc rows upper lower (row :. col) | row <= col = upper ! unSplitIdx rows (row,col+1)
                                       | otherwise  = lower ! unSplitIdx rows (col,row)
  
  unSplitIdx :: Int -> (Int,Int) -> Int
  unSplitIdx c (row,col) = go 0 (c-1) 0 + (col-1 - row)
    where
      go n c1 acc | n == row = acc
                  | otherwise = go (n+1) (c1-1) (acc+c1)
  
  
  
  updatePosition :: MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix1 Cl3_V3 -> IO ()
  updatePosition sysPos sysVel =
    do frozPos <- freeze Seq sysPos
       frozVel <- freeze Seq sysVel
       let ps :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozPos
           vs :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozVel
           ps' = ps !+! (dt *. vs)
       computeInto sysPos (A.map toCl3_V3 ps')
  
  
  offsetMomentum :: MArray RealWorld S Ix1 Cl3_V3 -> Array S Ix1 Cl3_R -> IO ()
  offsetMomentum sysVel sysMass =
    do frozVel <- freeze Seq sysVel
       planetsvs :: Array D Ix1 Cl3 <- extractM 1 4 (A.map fromCl3_V3 frozVel)  -- select just the planets and not the sun
       planetms :: Array D Ix1 Cl3 <- extractM 1 4 (A.map fromCl3_R sysMass)  -- select just the planets and not the sun
       (fromCl3_R -> sunMassCl3) <- sysMass !? 0
       writeM sysVel 0 (toCl3_V3 $ negate (planetsvs `dot_V3` planetms) / sunMassCl3)
       
  
  -- dot_V3 is to work around the typical sum in the dot product that has an aculumulator of a Real 0 value
  dot_V3 :: Array D Ix1 Cl3 -> Array D Ix1 Cl3 -> Cl3
  dot_V3 v1 v2 = foldlS (+) (V3 0 0 0) (A.zipWith (*) v1 v2)
  
  
  printEnergy :: MArray RealWorld S Ix1 Cl3_V3 -> MArray RealWorld S Ix1 Cl3_V3 -> Array S Ix1 Cl3_R -> IO ()
  printEnergy sysPos sysVel sysMass =
    do frozPos <- freeze Seq sysPos
       frozVel <- freeze Seq sysVel
       let ps :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozPos
           vs :: Array D Ix1 Cl3 = A.map fromCl3_V3 frozVel
           ms :: Array D Ix1 Cl3 = A.map fromCl3_R sysMass
           ke = 0.5 * (ms !.! A.map toR (vs.^2)) -- Kenetic Energy O(n)
           pe = negate.A.sum $ nChoose2With (*) ms !/! nChoose2With (\ x y -> abs $ x - y) ps -- Gravitational Potential Energy O((n^2 - n)/2)
       print $ ke + pe -- print the total energy
  
  -- This function is a bit of a combinatorial explosion so it is marked to be evaluated in Seqallel, probably should convert to a sgenerate :: Sz1 -> (Ix1 -> e) -> Vector DS e 
  nChoose2With :: (Cl3 -> Cl3 -> Cl3) -> Array D Ix1 Cl3 -> Array D Ix1 Cl3
  nChoose2With f bs =
    let (Sz1 c) = size bs
        c' = (c^(2 :: Int) - c) `div` 2
        bs' :: Array BN Ix1 Cl3 = computeSource bs
    in makeArray Seq (Sz1 c') (genElements f bs' c)
  
  genElements :: (Cl3 -> Cl3 -> Cl3) -> Array BN Ix1 Cl3 -> Int -> Ix1 -> Cl3
  genElements f bs c idx =
    let (ix1,ix2) = splitIdx c idx
    in f (bs ! ix1) (bs ! ix2)
  
  splitIdx :: Int -> Int -> (Int,Int)
  splitIdx c = go (0,1) (c-2)
    where
      go (row,col) x ix | ix - x <= 0 = (row, col + ix)
                        | otherwise = go (row+1,col+1) (x-1) (ix-x-1)
  
  planetspos :: [Cl3_V3]
  planetspos = pos <$> [sun,jupiter,saturn,uranus,neptune]
  
  planetsvel :: [Cl3_V3]
  planetsvel = vel <$> [sun,jupiter,saturn,uranus,neptune]
  
  planetsmass :: [Cl3_R]
  planetsmass = mas <$> [sun,jupiter,saturn,uranus,neptune]
  
  type SolarSystem = MArray RealWorld S Ix1 Body
  
  data Body = Body
    {pos :: Cl3_V3
    ,vel :: Cl3_V3
    ,mas :: Cl3_R}
  
  instance Storable Body where
    sizeOf _ = sizeOf (undefined :: Cl3_V3) + sizeOf (undefined :: Cl3_V3) + sizeOf (undefined :: Cl3_R)
    alignment _ = sizeOf (undefined :: Double) -- alignment must be a 2^x
    peekElemOff ptr idx = peek (castPtr ptr `plusPtr` (idx * sizeOf (undefined :: Body)))
    pokeElemOff ptr idx e = poke (castPtr ptr `plusPtr` (idx * sizeOf (undefined :: Body))) e
    peek :: Ptr Body -> IO Body
    peek ptr = do
      p :: Cl3_V3 <- peek (castPtr ptr `plusPtr` 0)
      v :: Cl3_V3 <- peek (castPtr ptr `plusPtr` (sizeOf p))
      m :: Cl3_R <- peek (castPtr ptr `plusPtr` (sizeOf p) `plusPtr` (sizeOf v))
      return (Body p v m)
    poke :: Ptr Body -> Body -> IO ()
    poke ptr (Body p v m) = do
      poke (castPtr ptr `plusPtr` 0) p
      poke (castPtr ptr `plusPtr` (sizeOf p)) v
      poke (castPtr ptr `plusPtr` (sizeOf p) `plusPtr` (sizeOf v)) m
  
  sun :: Body
  sun = Body (toCl3_V3 $ V3 0 0 0) -- Sun Position
             (toCl3_V3 $ V3 0 0 0) -- Sun Velocity 
             (toCl3_R $ R solarMass) -- Sun Mass
  jupiter :: Body
  jupiter = Body (toCl3_V3 $ V3 (4.84143144246472090e+00) (-1.16032004402742839e+00) (-1.03622044471123109e-01))
                 (toCl3_V3 $ V3 (1.66007664274403694e-03*dp) (7.69901118419740425e-03*dp) (-6.90460016972063023e-05*dp))
                 (toCl3_R $ R (9.54791938424326609e-04 * solarMass))
  saturn :: Body
  saturn = Body (toCl3_V3 $ V3 (8.34336671824457987e+00) (4.12479856412430479e+00) (-4.03523417114321381e-01))
                (toCl3_V3 $ V3 (-2.76742510726862411e-03*dp) (4.99852801234917238e-03*dp) (2.30417297573763929e-05*dp))
                (toCl3_R $ R (2.85885980666130812e-04 * solarMass))
  uranus :: Body
  uranus = Body (toCl3_V3 $ V3 (1.28943695621391310e+01) (-1.51111514016986312e+01) (-2.23307578892655734e-01))
                (toCl3_V3 $ V3 (2.96460137564761618e-03*dp) (2.37847173959480950e-03*dp) (-2.96589568540237556e-05*dp))
                (toCl3_R $ R (4.36624404335156298e-05 * solarMass))
  neptune :: Body
  neptune = Body (toCl3_V3 $ V3 (1.53796971148509165e+01) (-2.59193146099879641e+01) (1.79258772950371181e-01))
                 (toCl3_V3 $ V3 (2.68067772490389322e-03*dp) (1.62824170038242295e-03*dp) (-9.51592254519715870e-05*dp))
                 (toCl3_R $ R (5.15138902046611451e-05 * solarMass))
  
  
  daysPerYear :: Double
  daysPerYear = 365.24
  
  solarMass :: Double
  solarMass = 4 * pi^(2 :: Int)
  dp :: Double
  dp = daysPerYear
  dt :: Cl3
  dt = R 0.01
  