{-# OPTIONS_GHC -fplugin Foreign.Storable.Generic.Plugin #-}

{-# LANGUAGE OverloadedLists #-}

module Render.DescSets.Sun
  ( Sun(..)
  , createSet0Ds
  , set0

  , pattern MAX_VIEWS

  , Buffer

  , SunInput(..)
  , initialSunInput

  , Process
  , spawn1
  , mkSun

  , Observer
  , newObserver1
  , observe1
  ) where

import RIO

import Control.Monad.Trans.Resource (MonadResource, ResourceT)
import Data.Tagged (Tagged(..))
import Data.Vector qualified as Vector
import Data.Vector.Storable qualified as VectorS
import Foreign.Storable.Generic (GStorable)
import Geomancy (Vec3, Vec4, vec3, vec4)
import Geomancy.Transform (Transform)
import Geomancy.Transform qualified as Transform
import Geomancy.Vec4 qualified as Vec4
import Vulkan.Core10 qualified as Vk
import Vulkan.CStruct.Extends (SomeStruct(..))
import Vulkan.NamedType ((:::))
import Vulkan.Zero (Zero(..))

import Engine.Camera qualified as Camera
import Engine.Types (StageRIO)
import Engine.Vulkan.DescSets ()
import Engine.Vulkan.Types (DsLayoutBindings, HasVulkan(..))
import Engine.Worker qualified as Worker
import Resource.Buffer qualified as Buffer
import Resource.Region qualified as Region
import Resource.Vulkan.DescriptorPool qualified as DescriptorPool

-- * Set0 data for light projection

-- | Maximum "guaranteed" amount for multiview passes
pattern MAX_VIEWS :: Int
pattern $bMAX_VIEWS :: Int
$mMAX_VIEWS :: forall {r}. Int -> ((# #) -> r) -> ((# #) -> r) -> r
MAX_VIEWS = 6

data Sun = Sun
  { Sun -> Transform
sunViewProjection :: Transform
  , Sun -> Vec4
sunShadow         :: Vec4 -- offsetx, offsety, index, size -- XXX: only index is used
  , Sun -> Vec4
sunPosition       :: Vec4 -- XXX: alpha available for stuff
  , Sun -> Vec4
sunDirection      :: Vec4 -- XXX: alpha available for stuff
  , Sun -> Vec4
sunColor          :: Vec4 -- XXX: RGB premultiplied, alpha is available for stuff
  }
  deriving (Int -> Sun -> ShowS
[Sun] -> ShowS
Sun -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Sun] -> ShowS
$cshowList :: [Sun] -> ShowS
show :: Sun -> String
$cshow :: Sun -> String
showsPrec :: Int -> Sun -> ShowS
$cshowsPrec :: Int -> Sun -> ShowS
Show, forall x. Rep Sun x -> Sun
forall x. Sun -> Rep Sun x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Sun x -> Sun
$cfrom :: forall x. Sun -> Rep Sun x
Generic)

instance GStorable Sun

instance Zero Sun where
  zero :: Sun
zero = Sun
    { $sel:sunViewProjection:Sun :: Transform
sunViewProjection = forall a. Monoid a => a
mempty
    , $sel:sunShadow:Sun :: Vec4
sunShadow         = Vec4
0
    , $sel:sunPosition:Sun :: Vec4
sunPosition       = Vec4
0
    , $sel:sunDirection:Sun :: Vec4
sunDirection      = Float -> Float -> Float -> Float -> Vec4
vec4 Float
0 Float
1 Float
0 Float
0
    , $sel:sunColor:Sun :: Vec4
sunColor          = Vec4
0
    }

-- * Shadow casting descriptor set

set0
  :: Tagged Sun DsLayoutBindings
set0 :: Tagged Sun DsLayoutBindings
set0 = forall {k} (s :: k) b. b -> Tagged s b
Tagged
  [ (DescriptorSetLayoutBinding
set0bind0, forall a. Zero a => a
zero)
  ]

set0bind0 :: Vk.DescriptorSetLayoutBinding
set0bind0 :: DescriptorSetLayoutBinding
set0bind0 = Vk.DescriptorSetLayoutBinding
  { $sel:binding:DescriptorSetLayoutBinding :: Word32
binding           = Word32
0
  , $sel:descriptorType:DescriptorSetLayoutBinding :: DescriptorType
descriptorType    = DescriptorType
Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
  , $sel:descriptorCount:DescriptorSetLayoutBinding :: Word32
descriptorCount   = Word32
1
  , $sel:stageFlags:DescriptorSetLayoutBinding :: ShaderStageFlags
stageFlags        = ShaderStageFlags
Vk.SHADER_STAGE_VERTEX_BIT
  , $sel:immutableSamplers:DescriptorSetLayoutBinding :: Vector Sampler
immutableSamplers = forall a. Monoid a => a
mempty
  }

-- * Setup

type Buffer = Buffer.Allocated 'Buffer.Coherent Sun

createSet0Ds
  :: Tagged '[Sun] Vk.DescriptorSetLayout
  -> ResourceT (StageRIO st)
      ( Tagged '[Sun] (Vector Vk.DescriptorSet)
      , Buffer
      )
createSet0Ds :: forall st.
Tagged '[Sun] DescriptorSetLayout
-> ResourceT
     (StageRIO st) (Tagged '[Sun] (Vector DescriptorSet), Buffer)
createSet0Ds (Tagged DescriptorSetLayout
set0layout) = do
  DescriptorPool
descPool <- forall (m :: * -> *) a.
MonadResource m =>
m (ReleaseKey, a) -> ResourceT m a
Region.local forall a b. (a -> b) -> a -> b
$
    forall env (m :: * -> *).
(MonadVulkan env m, MonadResource m) =>
Maybe Text
-> Word32
-> [(DescriptorType, Word32)]
-> m (ReleaseKey, DescriptorPool)
DescriptorPool.allocate (forall a. a -> Maybe a
Just Text
"Basic.Sun") Word32
1
      [ ( DescriptorType
Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
        , Word32
1 forall a. Num a => a -> a -> a
+ Word32
1
        )
      ]

  Vector DescriptorSet
descSets <- forall env (m :: * -> *).
MonadVulkan env m =>
DescriptorPool
-> Maybe Text
-> Vector DescriptorSetLayout
-> m (Vector DescriptorSet)
DescriptorPool.allocateSetsFrom DescriptorPool
descPool (forall a. a -> Maybe a
Just Text
"Basic.Sun") [DescriptorSetLayout
set0layout]

  Buffer
sunData <- forall (m :: * -> *) a.
MonadResource m =>
m (ReleaseKey, a) -> ResourceT m a
Region.local forall a b. (a -> b) -> a -> b
$
    forall env (m :: * -> *) a.
(MonadVulkan env m, MonadResource m, Storable a) =>
Maybe Text
-> BufferUsageFlagBits
-> Int
-> Vector a
-> m (ReleaseKey, Allocated 'Coherent a)
Buffer.allocateCoherent
      (forall a. a -> Maybe a
Just Text
"Basic.Sun.Data")
      BufferUsageFlagBits
Vk.BUFFER_USAGE_UNIFORM_BUFFER_BIT Int
MAX_VIEWS
      (forall a. Storable a => Int -> a -> Vector a
VectorS.replicate Int
MAX_VIEWS forall a. Zero a => a
zero)

  forall st.
Tagged '[Sun] (Vector DescriptorSet)
-> Buffer -> ResourceT (StageRIO st) ()
updateSet0Ds (forall {k} (s :: k) b. b -> Tagged s b
Tagged Vector DescriptorSet
descSets) Buffer
sunData

  pure (forall {k} (s :: k) b. b -> Tagged s b
Tagged Vector DescriptorSet
descSets, Buffer
sunData)

updateSet0Ds
  :: Tagged '[Sun] (Vector Vk.DescriptorSet)
  -> Buffer.Allocated 'Buffer.Coherent Sun
  -> ResourceT (StageRIO st) ()
updateSet0Ds :: forall st.
Tagged '[Sun] (Vector DescriptorSet)
-> Buffer -> ResourceT (StageRIO st) ()
updateSet0Ds (Tagged Vector DescriptorSet
ds) Buffer.Allocated{Buffer
$sel:aBuffer:Allocated :: forall (s :: Store) a. Allocated s a -> Buffer
aBuffer :: Buffer
aBuffer} = do
  App GlobalHandles st
context <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a. a -> a
id
  forall (io :: * -> *).
MonadIO io =>
Device
-> ("descriptorWrites" ::: Vector (SomeStruct WriteDescriptorSet))
-> ("descriptorCopies" ::: Vector CopyDescriptorSet)
-> io ()
Vk.updateDescriptorSets (forall a. HasVulkan a => a -> Device
getDevice App GlobalHandles st
context) [SomeStruct WriteDescriptorSet
writeSet0b0] forall a. Monoid a => a
mempty

  where
    destSet0 :: DescriptorSet
destSet0 = case forall (m :: * -> *) a. Monad m => Vector a -> m a
Vector.headM Vector DescriptorSet
ds of
      Maybe DescriptorSet
Nothing ->
        forall a. HasCallStack => String -> a
error String
"assert: descriptor sets promised to contain [Sun]"
      Just DescriptorSet
one ->
        DescriptorSet
one

    writeSet0b0 :: SomeStruct WriteDescriptorSet
writeSet0b0 = forall (a :: [*] -> *) (es :: [*]).
(Extendss a es, PokeChain es, Show (Chain es)) =>
a es -> SomeStruct a
SomeStruct forall a. Zero a => a
zero
      { $sel:dstSet:WriteDescriptorSet :: DescriptorSet
Vk.dstSet          = DescriptorSet
destSet0
      , $sel:dstBinding:WriteDescriptorSet :: Word32
Vk.dstBinding      = Word32
0
      , $sel:dstArrayElement:WriteDescriptorSet :: Word32
Vk.dstArrayElement = Word32
0
      , $sel:descriptorCount:WriteDescriptorSet :: Word32
Vk.descriptorCount = Word32
1
      , $sel:descriptorType:WriteDescriptorSet :: DescriptorType
Vk.descriptorType  = DescriptorType
Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
      , $sel:bufferInfo:WriteDescriptorSet :: Vector DescriptorBufferInfo
Vk.bufferInfo      = [DescriptorBufferInfo
set0bind0I]
      }
      where
        set0bind0I :: DescriptorBufferInfo
set0bind0I = Vk.DescriptorBufferInfo
          { $sel:buffer:DescriptorBufferInfo :: Buffer
Vk.buffer = Buffer
aBuffer
          , $sel:offset:DescriptorBufferInfo :: DeviceSize
Vk.offset = DeviceSize
0
          , $sel:range:DescriptorBufferInfo :: DeviceSize
Vk.range  = DeviceSize
Vk.WHOLE_SIZE
          }

data SunInput = SunInput
  { SunInput -> Vec4
siColor :: Vec4

  , SunInput -> Float
siInclination :: Float
  , SunInput -> Float
siAzimuth     :: Float
  , SunInput -> Float
siRadius      :: Float
  , SunInput -> Vec3
siTarget      :: Vec3

  , SunInput -> Float
siDepthRange :: Float
  , SunInput -> Float
siSize       :: Float
  , SunInput -> Float
siShadowIx   :: Float
  }

initialSunInput :: SunInput
initialSunInput :: SunInput
initialSunInput = SunInput
  { $sel:siColor:SunInput :: Vec4
siColor = Float -> Float -> Float -> Float -> Vec4
vec4 Float
1 Float
1 Float
1 Float
1

  , $sel:siInclination:SunInput :: Float
siInclination = Float
τforall a. Fractional a => a -> a -> a
/Float
8
  , $sel:siAzimuth:SunInput :: Float
siAzimuth     = -Float
τforall a. Fractional a => a -> a -> a
/Float
8
  , $sel:siRadius:SunInput :: Float
siRadius      = forall a. (Eq a, Num a) => a
Camera.PROJECTION_FAR forall a. Fractional a => a -> a -> a
/ Float
2
  , $sel:siTarget:SunInput :: Vec3
siTarget      = Vec3
0

  , $sel:siDepthRange:SunInput :: Float
siDepthRange = forall a. (Eq a, Num a) => a
Camera.PROJECTION_FAR
  , $sel:siSize:SunInput :: Float
siSize       = Float
512
  , $sel:siShadowIx:SunInput :: Float
siShadowIx   = -Float
1
  }

type Process = Worker.Cell SunInput ("bounding box" ::: Transform, Sun)

spawn1
  :: ( MonadResource m
     , MonadUnliftIO m
     )
  => SunInput
  -> m Process
spawn1 :: forall (m :: * -> *).
(MonadResource m, MonadUnliftIO m) =>
SunInput -> m Process
spawn1 = forall (m :: * -> *) input output.
(MonadUnliftIO m, MonadResource m) =>
(input -> output) -> input -> m (Cell input output)
Worker.spawnCell SunInput -> (Transform, Sun)
mkSun

mkSun :: SunInput -> ("bounding box" ::: Transform, Sun)
mkSun :: SunInput -> (Transform, Sun)
mkSun SunInput{Float
Vec3
Vec4
siShadowIx :: Float
siSize :: Float
siDepthRange :: Float
siTarget :: Vec3
siRadius :: Float
siAzimuth :: Float
siInclination :: Float
siColor :: Vec4
$sel:siShadowIx:SunInput :: SunInput -> Float
$sel:siSize:SunInput :: SunInput -> Float
$sel:siDepthRange:SunInput :: SunInput -> Float
$sel:siTarget:SunInput :: SunInput -> Vec3
$sel:siRadius:SunInput :: SunInput -> Float
$sel:siAzimuth:SunInput :: SunInput -> Float
$sel:siInclination:SunInput :: SunInput -> Float
$sel:siColor:SunInput :: SunInput -> Vec4
..} =
  ( Transform
bbTransform
  , Sun
      { $sel:sunViewProjection:Sun :: Transform
sunViewProjection = forall a. Monoid a => [a] -> a
mconcat [Transform]
vp
      , $sel:sunShadow:Sun :: Vec4
sunShadow         = Float -> Float -> Float -> Float -> Vec4
vec4 Float
0 Float
0 Float
siShadowIx Float
siSize
      , $sel:sunPosition:Sun :: Vec4
sunPosition       = forall a. Coercible a Vec3 => a -> Float -> Vec4
Vec4.fromVec3 Vec3
position Float
0
      , $sel:sunDirection:Sun :: Vec4
sunDirection      = forall a. Coercible a Vec3 => a -> Float -> Vec4
Vec4.fromVec3 Vec3
direction Float
0
      , $sel:sunColor:Sun :: Vec4
sunColor          = Vec4
siColor
      }
  )
  where
    vp :: [Transform]
vp =
      [ Float -> Transform
Transform.rotateY (-Float
siAzimuth)
      , Float -> Transform
Transform.rotateX (-Float
siInclination)

      , Float -> Float -> Float -> Transform
Transform.translate Float
0 Float
0 Float
siRadius

      -- XXX: some area beyond the near plane receives light, but not shadows
      , Float -> Float -> Float -> Transform
Transform.scale3
          (Float
1 forall a. Fractional a => a -> a -> a
/ Float
siSize)
          (Float
1 forall a. Fractional a => a -> a -> a
/ Float
siSize)
          (Float
1 forall a. Fractional a => a -> a -> a
/ Float
siDepthRange)
      ]

    position :: Vec3
position = Vec3 -> Transform -> Vec3
Transform.apply (Float -> Float -> Float -> Vec3
vec3 Float
0 Float
0 Float
siRadius) Transform
rotation

    direction :: Vec3
direction = Vec3 -> Transform -> Vec3
Transform.apply (Float -> Float -> Float -> Vec3
vec3 Float
0 Float
0 forall a b. (a -> b) -> a -> b
$ -Float
1) Transform
rotation

    bbTransform :: Transform
bbTransform = forall a. Monoid a => [a] -> a
mconcat
      [ -- XXX: orient wire box "green/near -> far/red"
        Float -> Transform
Transform.rotateX (Float
τforall a. Fractional a => a -> a -> a
/Float
4)
        -- XXX: the rest must be matched with VP flipped
      , Float -> Float -> Float -> Transform
Transform.translate Float
0 Float
0 Float
0.5                 -- XXX: shift origin to the near face

        -- XXX: reverse light transform
      , Float -> Float -> Float -> Transform
Transform.scale3 Float
siSize Float
siSize Float
siDepthRange -- XXX: size to projection volume
      , Float -> Float -> Float -> Transform
Transform.translate Float
0 Float
0 (-Float
siRadius)         -- XXX: translate near face to radius
      , Transform
rotation                                    -- XXX: apply sphere coords
      ]

    rotation :: Transform
rotation = forall a. Monoid a => [a] -> a
mconcat
      [ Float -> Transform
Transform.rotateX Float
siInclination
      , Float -> Transform
Transform.rotateY Float
siAzimuth
      ]

type Observer = Worker.ObserverIO (VectorS.Vector ("bounding box" ::: Transform))

newObserver1 :: MonadIO m => m Observer
newObserver1 :: forall (m :: * -> *). MonadIO m => m Observer
newObserver1 = forall (m :: * -> *) a. MonadIO m => a -> m (ObserverIO a)
Worker.newObserverIO forall a. Monoid a => a
mempty

observe1 :: MonadUnliftIO m => Process -> Observer -> Buffer -> m ()
observe1 :: forall (m :: * -> *).
MonadUnliftIO m =>
Process -> Observer -> Buffer -> m ()
observe1 Process
sunP Observer
sunOut Buffer
sunData =
  forall (m :: * -> *) output a.
(MonadUnliftIO m, HasOutput output) =>
output -> ObserverIO a -> (a -> GetOutput output -> m a) -> m ()
Worker.observeIO_ Process
sunP Observer
sunOut \Vector Transform
_oldBB (Transform
bb, Sun
sun) -> do
    -- XXX: must stay the same or descsets must be updated with a new buffer
    Buffer
_same <- forall (io :: * -> *) a.
(MonadUnliftIO io, Storable a) =>
Vector a -> Allocated 'Coherent a -> io (Allocated 'Coherent a)
Buffer.updateCoherent (forall a. Storable a => a -> Vector a
VectorS.singleton Sun
sun) Buffer
sunData
    pure $ forall a. Storable a => a -> Vector a
VectorS.singleton Transform
bb

τ :: Float
τ :: Float
τ = Float
2 forall a. Num a => a -> a -> a
* forall a. Floating a => a
pi