software-shadows.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}

module Main where

import           Control.Exception
import           Control.Monad.Codensity
import           Control.Monad.Fix
import           Control.Monad.IO.Class
import           Control.Monad.ST
import           Data.Bits
import           Data.Foldable
import           Data.StateVar (($=))
import qualified Data.Vector as V
import qualified Data.Vector.Storable.Mutable as Vsm
import qualified Data.Vector.Unboxed as Vu
import           Data.Word
import           Foreign.ForeignPtr (newForeignPtr_)
import           Foreign.Ptr (castPtr)
import           Linear
import           Linear.Affine
import           Reflex.Class
import           Reflex.Dynamic
import           Reflex.Host.Class
import           Reflex.Spider
import qualified SDL as Sdl
import           System.Clock
import           System.Environment
import qualified System.Random.MWC as Mwc


data Scene rx =
    Scene {
      sceneImage :: Behavior rx (V2 Double -> V3 Double),
      sceneQuit  :: Event rx ()
    }


linear :: (Fractional a, Integral d, Traversable f, RealFrac s, Vu.Unbox a, Vu.Unbox (f d)) => (f d -> a) -> (f s -> a)
linear pic = \p -> Vu.head (foldl' go (ps0 p) (fmap originCube p))
    where
    originCube x = x - fromIntegral (floor x)
    ps0 = Vu.map pic . Vu.convert . rasterPoints

    go ps x = Vu.zipWith (\l r -> realToFrac (1 - x) * l + realToFrac x * r) ps1 ps2
        where
        (ps1, ps2) = Vu.splitAt (Vu.length ps `div` 2) ps

    rasterPoints =
        traverse $ \x ->
            let x0 = floor x
                x1 = ceiling x
            in V.fromListN 2 [x0, x1]


myScene :: (MonadFix m, MonadHold rx m, Reflex rx) => Event rx Double -> Event rx Sdl.EventPayload -> m (Scene rx)
myScene tickEv sdlEv = do
    time <- foldDyn (+) 0 tickEv
    mousePos <- hold 0 mousePosEv

    pure (Scene {
            sceneImage =
                pure (\(V2 mx my) t p ->
                          let int = pure 0.25 + V3 (0.05 + 0.05 * sin (pi*t)) 0 0 in
                          light int (V3 mx my 1) (normalMap p) (heightMap p) (colourMap p) 0)
                <*> mousePos
                <*> current time,

            sceneQuit = fmapMaybe quitFilt sdlEv
          })

    where
    !shadowDistContrib = recip (fromIntegral (Vu.length shadowDist))
    !shadowDist =
        Vu.fromList . map (0.01 *^) . filter ((<= 1) . norm) $
        sequenceA (pure [1, 1/3, -1/3, 1])

    light int l' n p col c = col * (amb + (multishadow *^ (diff + spec)) ^/ distance p l')
        where
        l = normalize (l' - p)
        amb  = pure 0.02
        diff = max 0 (dot n l) *^ int
        spec = (max 0 (dot n (normalize (l + (c - p)))) ** 100) *^ int

        multishadow =
            Vu.foldl' (+) 0 .
            Vu.map (shadow . (l' +)) $
            shadowDist

        shadow sl = topShadow
            where
            V3 lx ly lz = sl
            V3 px py pz = p

            topShadow =
                let h = height
                    pzlz = pz - lz
                    m = V2 ((lx*pz - lz*px + h*px - h*lx) / pzlz)
                           ((ly*pz - lz*py + h*py - h*ly) / pzlz)
                    s = (lz - h) / pzlz
                    ml = norm m
                in if ml < radius && s > (-1) then 0 else shadowDistContrib

    height = 0.5
    radius = 0.2

    heightMap :: V2 Double -> V3 Double
    heightMap p@(V2 x y)
        | otherwise = V3 x y (max 0 . min height $ height - height*5*(pl - radius))
        where
        pl = norm p

    normalMap :: V2 Double -> V3 Double
    normalMap p@(V2 x y)
        | pl < radius || pl > radius + 0.2 = V3 0 0 1
        | otherwise = normalize $ cross (V3 0 0 height + 0.2 *^ normalize (V3 (-x) (-y) 0)) (V3 y (-x) 0)
        where
        pl = norm p

    colourMap :: V2 Double -> V3 Double
    colourMap =
        runST $ do
            rng <- Mwc.create
            let gen = sequenceA (pure (Mwc.uniformR (0.5, 1) rng))
            arr <- Vu.replicateM 64 gen

            let scale p = 3.5 * (p + 1)
                clamp   = fmap (max 0 . min 7)
                iget (V2 x y) = arr Vu.! (8*y + x)

            pure (linear (iget . clamp) . scale)

    mousePosEv = fmapMaybe filt sdlEv
        where
        filt (Sdl.MouseMotionEvent motion) =
            let P (V2 x y) = fromIntegral <$> Sdl.mouseMotionEventPos motion
            in Just (V2 (-1 + x / 399.5) (1 - y / 399.5))
        filt _ = Nothing

    quitFilt Sdl.QuitEvent = Just ()
    quitFilt (Sdl.KeyboardEvent (Sdl.KeyboardEventData _ _ _ (Sdl.Keysym _ Sdl.KeycodeQ _))) = Just ()
    quitFilt _ = Nothing


main :: IO ()
main =
    lowerCodensity $ do
        args <- map read <$> liftIO getArgs
        let txtSize@(V2 tw th) =
                case args of
                  tw' : th' : _ -> V2 tw' th'
                  _             -> V2 30 30
        cbracket_ (Sdl.initialize [Sdl.InitVideo]) Sdl.quit
        win <- cbracket (Sdl.createWindow "Test" winCfg) Sdl.destroyWindow
        rend <- cbracket (Sdl.createRenderer win (-1) Sdl.defaultRenderer) Sdl.destroyRenderer
        txt <- cbracket (Sdl.createTexture rend Sdl.RGBA8888 Sdl.TextureAccessStreaming (fromIntegral <$> txtSize)) Sdl.destroyTexture
        liftIO . runSpiderHost $ do
            Sdl.HintRenderScaleQuality $= Sdl.ScaleLinear
            (tickEv, tickRef) <- newEventWithTriggerRef
            (sdlEv, sdlRef) <- newEventWithTriggerRef
            scene <- runHostFrame (myScene tickEv sdlEv)
            quitH <- subscribeEvent (sceneQuit scene)

            let draw = do
                    img <- sample (sceneImage scene)
                    liftIO $ do
                        bracket (Sdl.lockTexture txt Nothing) (\_ -> Sdl.unlockTexture txt) $ \(ptr', _) -> do
                            ptr <- newForeignPtr_ (castPtr ptr')
                            let arr = Vsm.unsafeFromForeignPtr0 ptr (tw*th) :: Vsm.IOVector Word32
                                i1  = tw*th

                                go :: Int -> Int -> Int -> IO ()
                                go i x y
                                    | i >= i1 = pure ()
                                    | x >= tw = go i 0 (y + 1)
                                    | otherwise = do
                                        let xf = -1 + 2*fromIntegral x / fromIntegral (tw - 1)
                                            yf = 1 - 2*fromIntegral y / fromIntegral (th - 1)
                                            V3 r g b =
                                                (\c -> round (255*c)) . min 1 . max 0
                                                <$> img (V2 xf yf)
                                        Vsm.write arr i (shiftL r 24 .|. shiftL g 16 .|. shiftL b 8 .|. 255)
                                        go (i + 1) (x + 1) y
                            go 0 (0 :: Int) (0 :: Int)
                        Sdl.clear rend
                        Sdl.copy rend txt Nothing Nothing
                        Sdl.present rend

            t0 <- liftIO (getTime Monotonic)
            ($ t0) . fix $ \again t' -> do
                mev <- liftIO Sdl.pollEvent
                case mev of
                  Nothing -> do
                      t <- liftIO (getTime Monotonic)
                      let !dt = fromInteger (toNanoSecs (t - t')) / 1000000000
                      fireEventRefAndRead tickRef dt quitH >>=
                          maybe (draw >> again t) (\_ -> pure ())
                  Just ev ->
                      fireEventRefAndRead sdlRef (Sdl.eventPayload ev) quitH >>=
                      maybe (again t') (\_ -> pure ())

    where
    winCfg = Sdl.defaultWindow {
               Sdl.windowInitialSize = V2 800 800
             }

    cbracket :: IO a -> (a -> IO r) -> Codensity IO a
    cbracket create destroy = Codensity (bracket create destroy)

    cbracket_ :: IO a -> IO b -> Codensity IO a
    cbracket_ start stop = Codensity (bracket start (const stop))