content: {-# LANGUAGE BangPatterns     #-}
  {-# LANGUAGE FlexibleContexts #-}
  
  module Stage.Linear
          ( linearSolver
          , unstableSolver)
  where
  import Model
  import FieldElt                         as E
  import Data.Array.Repa                  as R
  import Data.Array.Repa.Stencil          as R
  import Data.Array.Repa.Stencil.Dim2     as R
  import Data.Array.Repa.Eval             as R
  import Data.Vector.Unboxed
  import Prelude as P
  
  
  -------------------------------------------------------------------------------
  linearSolver
          :: (FieldElt a, Source U a, Unbox a, Elt a, Num a)
          => Field a      -- ^ Original field.
          -> Field a      -- ^ Current field.
          -> Float
          -> Float
          -> Int          -- ^ Number of iterations to apply.
          -> IO (Field a)
  
  linearSolver origField curField !a !c !iters
          -- If nothing would change by running the solver, then skip it.
          | 0 <- a       = return origField
  
          -- The solver has finished its loop
          | 0 <- iters   = return curField
  
          -- Do one iteration
          | otherwise
          = origField `deepSeqArray` curField `deepSeqArray`
            do    let !c' = 1/c
                  let {-# INLINE zipFunc #-}
                      zipFunc !orig !new
                          = (orig ~+~ (new ~*~ a)) ~*~ c'
  
                  newField <- {-# SCC "linearSolver.mapStencil" #-}
                             computeUnboxedP
                           $ R.szipWith zipFunc origField
                           $ mapStencil2 (BoundConst 0) linearSolverStencil curField
  
                  -- TODO: this boundConst thing is costing a fair bit.
                  -- Do something about the branches introduced into the core code.
  
                  linearSolver origField newField a c (iters - 1)
  
  {-# SPECIALIZE linearSolver
          :: Field Float
          -> Field Float
          -> Float -> Float -> Int
          -> IO (Field Float) #-}
  
  {-# SPECIALIZE linearSolver
          :: Field (Float, Float)
          -> Field (Float, Float)
          -> Float -> Float -> Int
          -> IO (Field (Float, Float)) #-}
  
  
  
  -- | Stencil function for the linear solver.
  linearSolverStencil
          :: FieldElt a
          => Stencil DIM2 a
  
  linearSolverStencil
   = StencilStatic (Z:.3:.3) E.zero
        (\ix val acc ->
           case linearSolverCoeffs ix of
              Nothing    -> acc
              Just coeff -> acc ~+~ (val ~*~ coeff))
  {-# INLINE linearSolverStencil #-}
  
  
  -- | Linear solver stencil kernel.
  linearSolverCoeffs :: DIM2 -> Maybe Float
  linearSolverCoeffs (Z:.j:.i)
     | i ==  1, j ==  0 = Just 1
     | i == -1, j ==  0 = Just 1
     | i ==  0, j ==  1 = Just 1
     | i ==  0, j == -1 = Just 1
     | otherwise        = Nothing
  {-# INLINE linearSolverCoeffs #-}
  
  
  -- Unstable -------------------------------------------------------------------
  unstableSolver
          :: (FieldElt a, Source U a, Unbox a, Elt a, Num a)
          => Field a      -- ^ Original field.
          -> Field a      -- ^ Current field.
          -> Float
          -> IO (Field a)
  
  unstableSolver !origField !curField !a
   = origField `deepSeqArray` curField `deepSeqArray`
     do
          let {-# INLINE zipFunc #-}
              zipFunc !orig !new
                  = orig ~+~ (new ~*~ a)
  
          newField
                  <- computeUnboxedP
                  $  R.szipWith zipFunc origField
                  $  mapStencil2 (BoundConst 0) unstableSolverStencil curField
  
          return newField
  
  
  {-# SPECIALIZE unstableSolver
          :: Field Float
          -> Field Float
          -> Float
          -> IO (Field Float) #-}
  
  {-# SPECIALIZE unstableSolver
          :: Field (Float, Float)
          -> Field (Float, Float)
          -> Float
          -> IO (Field (Float, Float)) #-}
  
  
  unstableSolverStencil
          :: FieldElt a
          => Stencil DIM2 a
  
  unstableSolverStencil
   = StencilStatic (Z:.3:.3) E.zero
        (\ix val acc ->
           case unstableSolverCoeffs ix of
              Nothing    -> acc
              Just coeff -> acc ~+~ (val ~*~ coeff))
  {-# INLINE unstableSolverStencil #-}
  
  
  unstableSolverCoeffs :: DIM2 -> Maybe Float
  unstableSolverCoeffs (Z:.j:.i)
      | i == 1,  j == 0   = Just 1
      | i == -1, j == 0   = Just 1
      | i == 0,  j == 1   = Just 1
      | i == 0,  j == -1  = Just 1
      | i == 0,  j == 0   = Just (-4)
      | otherwise         = Nothing
  {-# INLINE unstableSolverCoeffs #-}
  
  