{-# LANGUAGE FlexibleInstances #-}
--
-- Types used throughout the simulation
--

module Fluid.Type (

  Timestep, Viscosity, Diffusion, Index, Density, Velocity,
  Field, FieldElt(..), DensityField, VelocityField, DensitySource, VelocitySource, Pic,

  RGBA, Image

) where

import Data.Word
import Data.Array.Accelerate

type Timestep           = Float
type Viscosity          = Float
type Diffusion          = Float
type Index              = DIM2
type Density            = Float
type Velocity           = (Float, Float)

type Field a            = Array DIM2 a
type DensityField       = Field Density
type VelocityField      = Field Velocity

type Source a           = (Vector Index, Vector a)
type DensitySource      = Source Density
type VelocitySource     = Source Velocity

type RGBA               = Word32
type Image a            = Array DIM2 a
type Pic                = Image (Word8,Word8,Word8)


infixl 6 .+.
infixl 6 .-.
infixl 7 .*.
infixl 7 ./.

class Elt e => FieldElt e where
  zero  :: e
  (.+.) :: Exp e -> Exp e -> Exp e
  (.-.) :: Exp e -> Exp e -> Exp e
  (.*.) :: Exp Float -> Exp e -> Exp e
  (./.) :: Exp e -> Exp Float -> Exp e

instance FieldElt Density where
  zero  = 0
  (.+.) = (+)
  (.-.) = (-)
  (.*.) = (*)
  (./.) = (/)

{-
instance FieldElt Density where
  zero = (0,0,0)
  (.+.) = app3 (+)
  (.-.) = app3 (-)
  c  .*. xy = let (x,y,z) = unlift xy in lift (c*x, c*y, c*z)
  xy ./. c  = let (x,y,z) = unlift xy in lift (x/c, y/c, z/c)
-}

instance FieldElt Velocity where
  zero  = (0, 0)
  (.+.) = app2 (+)
  (.-.) = app2 (-)
  c  .*. xy = let (x,y) = unlift xy in lift (c*x, c*y)
  xy ./. c  = let (x,y) = unlift xy in lift (x/c, y/c)

app2 :: Elt e => (Exp e -> Exp e -> Exp e) -> Exp (e,e) -> Exp (e,e) -> Exp (e,e)
app2 f xu yv = let (x,u) = unlift xu
                   (y,v) = unlift yv
               in  lift (f x y, f u v)

app3 :: Elt e => (Exp e -> Exp e -> Exp e) -> Exp (e,e,e) -> Exp (e,e,e) -> Exp (e,e,e)
app3 f xu yv = let (x,i,u) = unlift xu
                   (y,j,v) = unlift yv
               in  lift (f x y, f i j, f u v)
