content: {-# LANGUAGE BangPatterns #-}
  
  module Stage.Project
          (project)
  where
  import Model
  import FieldElt
  import Stage.Linear
  import Data.Array.Repa          as R
  import Data.Array.Repa.Unsafe   as R
  import Prelude                  as P
  
  
  project :: Int -> Field (Float, Float) -> IO (Field (Float, Float))
  project iters field
   = {-# SCC project #-}
     field `deepSeqArray`
     do   let _ :. height :. width = extent field
  
          divergence <- {-# SCC "project.genDiv" #-}
                        computeUnboxedP
                     $  fromFunction (Z:. height :. width) (genDivergence width height field)
  
          p          <- {-# SCC "project.linearSolver" #-}
                        linearSolver divergence divergence 1 4 iters
  
          f'         <- {-# SCC "project.apply" #-}
                        computeUnboxedP
                  $     unsafeTraverse field id (projectElem width height p)
  
          return f'
  {-# NOINLINE project #-}
  
  
  -- | Subtract a gradient field from the regular field to
  --   create a mass-conserving field.
  projectElem
          :: Int                          -- ^ Width of model.
          -> Int                          -- ^ Height of model.
          -> Field Float
          -> (DIM2 -> (Float, Float))     -- ^ Get data from the regular field.
          -> DIM2                         -- ^ Compute the value at this point.
          -> (Float, Float)
  
  projectElem !width !height !p !get !pos@(Z:.j:.i)
   = get pos ~-~ (0.5 * width'  * (p0 - p1),
                  0.5 * height' * (p2 - p3))
   where
          !width'  = fromIntegral width
          !height' = fromIntegral height
          !p0      = useIf (i < width -  1) (p `unsafeIndex` (Z :. j   :. i+1))
          !p1      = useIf (i >          0) (p `unsafeIndex` (Z :. j   :. i-1))
          !p2      = useIf (j < height - 1) (p `unsafeIndex` (Z :. j+1 :. i  ))
          !p3      = useIf (j >          0) (p `unsafeIndex` (Z :. j-1 :. i  ))
  {-# INLINE projectElem #-}
  
  
  -- | Get an approximation of the gradient at this point.
  genDivergence :: Int -> Int -> VelocityField -> DIM2 -> Float
  genDivergence !width !height !f (Z :. j :. i)
   = (-0.5 * ((u0 - u1) + (v0 - v1))) / fromIntegral width
   where
        (u0,  _) = useIf (i < width -  1) (f `unsafeIndex` (Z:. j   :. i+1))
        (u1,  _) = useIf (i >          0) (f `unsafeIndex` (Z:. j   :. i-1))
        ( _, v0) = useIf (j < height - 1) (f `unsafeIndex` (Z:. j+1 :. i  ))
        ( _, v1) = useIf (j >          0) (f `unsafeIndex` (Z:. j-1 :. i  ))
  {-# INLINE genDivergence #-}
  
  
  