Icosphere.hs

module Geometry.Icosphere
  ( generateIndexed
  , generateIndexedCW

  , icofaces
  , icopoints
  ) where

import RIO

import Geomancy.Vec3 (Vec3, vec3)
import Geomancy.Vec3 qualified as Vec3
import Data.Vector qualified as Vector
import Data.Vector.Mutable qualified as Mutable
import RIO.Vector.Partial ((!))
import RIO.Map qualified as Map
import RIO.Vector.Storable qualified as Storable
import Control.Monad.State.Strict (get, put, runState)
import Vulkan.NamedType ((:::))

import Geometry.Face (Face(..))

{- | Subdivide a D20 iteratively to get an indexed mesh.

Normalize positions from the input to keep them on a unit sphere.

The triangles would be placed in a counter-clockwise order.

@
(positions, attrs, indices) =
  Icosphere.generateIndexed
    3
    mkInitialAttrs
    mkMidpointAttrs
    mkVertices
  where
    mkInitialAttrs :: Vec3 -> ()
    mkInitialAttrs _pos = ()

    mkMidpointAttrs :: Float -> Vec3 -> () -> () -> ()
    mkMidpointAttrs _scale _midPos _attr1 _attr2 = ()

    mkVertices points _faces = do
      (rawPos, ()) <- points
      let normPos = Vec3.normalize rawPos
      pure
        ( Vec3.Packed normPos
        , Vec4.fromVec3 (normPos / 2 + 0.5) 1
        )
@
-}
generateIndexed
  :: ( Fractional scale
     , Storable pos
     , Storable vertexAttr
     )
  => "subdivisions"  ::: Natural
  -> "initial"       ::: (Vec3 -> pointAttr)
  -> "midpoint"      ::: (scale -> Vec3 -> pointAttr -> pointAttr -> pointAttr)
  -> "vertex"        ::: (Vector (Vec3, pointAttr) -> [Face Int] -> Vector (pos, vertexAttr))
  -> "model vectors" ::: (Storable.Vector pos, Storable.Vector vertexAttr, Storable.Vector Word32)
generateIndexed = generateIndexedEx True

-- | Like, `generateIndexed`, but the triangle winding order is clockwise.
--
-- The resulting object will appear inside-out unless the pipeline order is set to CLOCKWISE too.
generateIndexedCW
  :: ( Fractional scale
     , Storable pos
     , Storable vertexAttr
     )
  => "subdivisions"  ::: Natural
  -> "initial"       ::: (Vec3 -> pointAttr)
  -> "midpoint"      ::: (scale -> Vec3 -> pointAttr -> pointAttr -> pointAttr)
  -> "vertex"        ::: (Vector (Vec3, pointAttr) -> [Face Int] -> Vector (pos, vertexAttr))
  -> "model vectors" ::: (Storable.Vector pos, Storable.Vector vertexAttr, Storable.Vector Word32)
generateIndexedCW = generateIndexedEx False

generateIndexedEx
  :: ( Fractional scale
     , Storable pos
     , Storable vertexAttr
     )
  => "clockwise"   ::: Bool
  -> "subdivisions"  ::: Natural
  -> "initial"       ::: (Vec3 -> pointAttr)
  -> "midpoint"      ::: (scale -> Vec3 -> pointAttr -> pointAttr -> pointAttr)
  -> "vertex"        ::: (Vector (Vec3, pointAttr) -> [Face Int] -> Vector (pos, vertexAttr))
  -> "model vectors" ::: (Storable.Vector pos, Storable.Vector vertexAttr, Storable.Vector Word32)
generateIndexedEx windingCCW details mkInitialAttrs mkMidpointAttrs mkVertices =
  ( Storable.convert pv
  , Storable.convert av
  , Storable.fromList iv
  )
  where
    (pv, av) = Vector.unzip $ mkVertices finalPoints faces

    iv = do
      Face a b c <- faces
      vert <- if windingCCW then [c, b, a] else [a, b, c]
      pure $ fromIntegral vert

    (faces, (_midpoints, finalPoints, _finalPointsCount)) =
      runState
        (go icofaces details details)
        ( mempty
        , initialPoints
        , Vector.length icopoints
        )

    maxPoints = length icofaces * (4 ^ details) - 8

    initialPoints = Vector.create do
      v <- Mutable.new maxPoints
      Vector.imapM_
        (Mutable.unsafeWrite v)
        (Vector.map (id &&& mkInitialAttrs) icopoints)
      pure v

    go curFaces maxLevel curLevel = do
      -- traceShowM $ "Inflating level " <> textShow (maxLevel - curLevel)
      case curLevel of
        0 ->
          pure curFaces
        _ -> do
          let scale = fromIntegral curLevel / fromIntegral maxLevel
          next <- traverse (subdivideFace scale) curFaces
          go (mconcat next) maxLevel (curLevel - 1)

    subdivideFace scale (Face a b c) = do
      (mids, points, numPoints) <- get

      let
        extras = mempty
        (midsAB, extrasAB, ab) = midpoint scale mids   extras   points numPoints (a, b)
        (midsBC, extrasBC, bc) = midpoint scale midsAB extrasAB points numPoints (b, c)
        (midsCA, extrasCA, ca) = midpoint scale midsBC extrasBC points numPoints (c, a)

      put
        ( midsCA
        , runST do
            old <- Vector.unsafeThaw points

            Vector.imapM_
              ( \i point ->
                  Mutable.unsafeWrite old (numPoints + i) point
              )
              extrasCA

            Vector.unsafeFreeze old
        , numPoints + Vector.length extrasCA
        )

      pure
        [ Face ab bc ca
        , Face ca a ab
        , Face ab b bc
        , Face bc c ca
        ]

    midpoint scale mids extras points numPoints parents =
      case Map.lookup parents mids of
        Just knownIx ->
          ( mids
          , extras
          , knownIx
          )
        Nothing ->
          let
            (pos1, attr1) = points ! fst parents
            (pos2, attr2) = points ! snd parents
            midPos = Vec3.lerp 0.5 pos1 pos2

            newIx =
              numPoints +
              Vector.length extras

            point =
              ( midPos
              , mkMidpointAttrs scale midPos attr1 attr2
              )
          in
            ( Map.insert parents newIx mids
            , Vector.snoc extras point
            , newIx
            )

icofaces :: [Face Int]
icofaces =
  [ -- faces around point 0
    Face  5 11 0
  , Face  1  5 0
  , Face  7  1 0
  , Face 10  7 0
  , Face 11 10 0

    -- 5 adjacent faces
  , Face 9  5  1
  , Face 4 11  5
  , Face 2 10 11
  , Face 6  7 10
  , Face 8  1  7

    -- 5 adjacent faces around point 3
  , Face 4 9 3
  , Face 2 4 3
  , Face 6 2 3
  , Face 8 6 3
  , Face 9 8 3

    -- 5 adjacent faces
  , Face  5 9 4
  , Face 11 4 2
  , Face 10 2 6
  , Face  7 6 8
  , Face  1 8 9
  ]

icopoints :: Vector Vec3
icopoints = Vector.fromList
  [ vec3 (-1) 0   t
  , vec3   1  0   t
  , vec3 (-1) 0 (-t)
  , vec3   1  0 (-t)

  , vec3   0 (-t) (-1)
  , vec3   0 (-t)   1
  , vec3   0   t  (-1)
  , vec3   0   t    1

  , vec3   t     1  0
  , vec3   t   (-1) 0
  , vec3 (-t)    1  0
  , vec3 (-t)  (-1) 0
  ]
  where
    t = (1.0 + sqrt 5.0) / 2.0