--
-- Fluid simulation
--

module Fluid (

  Simulation, fluid

) where

import Type

import Data.Array.Accelerate ( Z(..), (:.)(..), Exp, Acc, Scalar, Vector, (?|), (==*) )
import qualified Data.Array.Accelerate  as A


type Simulation
  =  Acc ( Scalar Timestep              -- time to evolve the simulation
         , DensitySource                -- locations to add density sources
         , VelocitySource               -- locations to add velocity sources
         , DensityField                 -- the current density field
         , VelocityField )              -- the current velocity field
  -> Acc ( DensityField, VelocityField )


-- A fluid simulation
--
fluid :: Viscosity -> Diffusion -> Simulation
fluid dp dn inputs =
  let (dt, ds, vs, df, vf)      = A.unlift inputs
      df'                       = density  dn dt ds vf df
      vf'                       = velocity dp dt vs vf
  in
  A.lift (df', vf')

-- The velocity over a timestep evolves due to three causes:
--   1. the addition of forces
--   2. viscous diffusion
--   3. self-advection
--
velocity
    :: Viscosity
    -> Acc (Scalar Timestep)
    -> Acc VelocitySource
    -> Acc VelocityField
    -> Acc VelocityField
velocity dp dt vs vf0
  = project
  . advect dt vf0
  . project
  . diffuse dp dt
  $ inject vs vf0


-- Ensure the velocity field conserves mass
--
project :: Acc VelocityField -> Acc VelocityField
project vf = A.stencil2 poisson A.Mirror vf A.Mirror p
  where
    steps       = 20
    grad        = A.stencil divF A.Mirror vf
    p1          = A.stencil2 pF (A.Constant 0) grad A.Mirror
    p           = foldl1 (.) (replicate steps p1) grad

    poisson :: A.Stencil3x3 Velocity -> A.Stencil3x3 Float -> Exp Velocity
    poisson (_,(_,uv,_),_) ((_,t,_), (l,_,r), (_,b,_)) = uv .-. 0.5 .*. A.lift (r-l, t-b)

    divF :: A.Stencil3x3 Velocity -> Exp Float
    divF ((_,t,_), (l,_,r), (_,b,_)) = -0.5 * (A.fst r - A.fst l + A.snd t - A.snd b)

    pF :: A.Stencil3x3 Float -> A.Stencil3x3 Float -> Exp Float
    pF (_,(_,x,_),_) ((_,t,_), (l,_,r), (_,b,_)) = 0.25 * (x + l + t + r + b)



-- The density over a timestep evolves due to three causes:
--   1. the addition of source particles
--   2. self-diffusion
--   3. motion through the velocity field
--
density
    :: Diffusion
    -> Acc (Scalar Timestep)
    -> Acc DensitySource
    -> Acc VelocityField
    -> Acc DensityField
    -> Acc DensityField
density dn dt ds vf
  = advect dt vf
  . diffuse dn dt
  . inject ds


-- Inject sources into the field
--
-- TLM: sources should be a vector of (index, value) pairs, but no fusion means
--   that extracting the components for permute (via unzip) is extra work.
--
inject
    :: FieldElt e
    => Acc (Vector Index, Vector e)
    -> Acc (Field e)
    -> Acc (Field e)
inject source field =
  let (is, ps) = A.unlift source
  in A.size ps ==* 0
       ?| ( field, A.permute (.+.) field (is A.!) ps )


diffuse
    :: FieldElt e
    => Diffusion
    -> Acc (Scalar Timestep)
    -> Acc (Field e)
    -> Acc (Field e)
diffuse dn dt df0 =
  a ==* 0
    ?| ( df0 , foldl1 (.) (replicate steps diffuse1) df0 )
  where
    steps       = 20
    a           = A.the dt * A.constant dn * (A.fromIntegral (A.size df0))
    c           = 1 + 4*a

    diffuse1 df = A.stencil2 relax (A.Constant zero) df0 A.Mirror df

    relax :: FieldElt e => A.Stencil3x3 e -> A.Stencil3x3 e -> Exp e
    relax (_,(_,x0,_),_) ((_,t,_), (l,_,r), (_,b,_)) = (x0 .+. a .*. (l.+.t.+.r.+.b)) ./. c


advect
    :: FieldElt e
    => Acc (Scalar Timestep)
    -> Acc VelocityField
    -> Acc (Field e)
    -> Acc (Field e)
advect dt' vf df = A.generate sh backtrace
  where
    dt          = A.the dt'
    sh          = A.shape vf
    Z :. h :. w = A.unlift sh

    backtrace ix = s0.*.(t0.*.d00 .+. t1.*.d10) .+. s1.*.(t0.*.d01 .+. t1.*.d11)
      where
        Z:.j:.i = A.unlift ix
        (u, v)  = A.unlift (vf A.! ix)

        -- backtrack densities based on velocity field
        clamp z = A.max 0.5 . A.min (A.fromIntegral z - 1.5)
        x       = w `clamp` (A.fromIntegral i - dt * u)
        y       = h `clamp` (A.fromIntegral j - dt * v)

        -- discrete locations surrounding point
        i0      = A.truncate x
        j0      = A.truncate y
        i1      = i0 + 1
        j1      = j0 + 1

        -- weighting based on location between the discrete points
        s1      = x - A.fromIntegral i0
        t1      = y - A.fromIntegral j0
        s0      = 1 - s1
        t0      = 1 - t1

        -- read the density values surrounding the calculated advection point
        d00     = df A.! A.lift (Z :. j0 :. i0)
        d10     = df A.! A.lift (Z :. j1 :. i0)
        d01     = df A.! A.lift (Z :. j0 :. i1)
        d11     = df A.! A.lift (Z :. j1 :. i1)