module Geometry.Icosphere
  ( generateIndexed

  , icofaces
  , icopoints

  , icosphere
  , icotris_v1
  , subNormal
  , subdivide
  ) where

import RIO

import Data.List (iterate, (!!))
import Geomancy.Vec3 (Vec3, vec3)
import Geomancy.Vec3 qualified as Vec3
import Resource.Model qualified as Model
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(..))

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 details mkInitialAttrs mkMidpointAttrs mkVertices =
  ( Storable.convert pv
  , Storable.convert av
  , Storable.fromList iv
  )
  where
    (pv, av) = Vector.unzip $ mkVertices finalPoints faces

    iv = do
      face <- faces
      vert <- toList face
      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

icotris_v1 :: [[Vec3]]
icotris_v1 =
  [ -- faces around point 0
    [ icopoints ! 0
    , icopoints ! 11
    , icopoints ! 5
    ]

  , [ icopoints ! 0
    , icopoints ! 5
    , icopoints ! 1
    ]

  , [ icopoints ! 0
    , icopoints ! 1
    , icopoints ! 7
    ]
{-
  , [ icopoints ! 0
    , icopoints ! 1
    , icopoints ! 7
    ]
-}
  , [ icopoints ! 0
    , icopoints ! 7
    , icopoints ! 10
    ]

  , [ icopoints ! 0
    , icopoints ! 10
    , icopoints ! 11
    ]

    -- 5 adjacent faces

  , [ icopoints ! 1
    , icopoints ! 5
    , icopoints ! 9
    ]

  , [ icopoints ! 5
    , icopoints ! 11
    , icopoints ! 4
    ]

  , [ icopoints ! 11
    , icopoints ! 10
    , icopoints ! 2
    ]

  , [ icopoints ! 10
    , icopoints ! 7
    , icopoints ! 6
    ]

  , [ icopoints ! 7
    , icopoints ! 1
    , icopoints ! 8
    ]

    -- 5 adjacent faces around point 3

  , [ icopoints ! 3
    , icopoints ! 9
    , icopoints ! 4
    ]

  , [ icopoints ! 3
    , icopoints ! 4
    , icopoints ! 2
    ]

  , [ icopoints ! 3
    , icopoints ! 2
    , icopoints ! 6
    ]

  , [ icopoints ! 3
    , icopoints ! 6
    , icopoints ! 8
    ]

  , [ icopoints ! 3
    , icopoints ! 8
    , icopoints ! 9
    ]

    -- 5 adjacent faces

  , [ icopoints ! 4
    , icopoints ! 9
    , icopoints ! 5
    ]

  , [ icopoints ! 2
    , icopoints ! 4
    , icopoints ! 11
    ]

  , [ icopoints ! 6
    , icopoints ! 2
    , icopoints ! 10
    ]

  , [ icopoints ! 8
    , icopoints ! 6
    , icopoints ! 7
    ]

  , [ icopoints ! 9
    , icopoints ! 8
    , icopoints ! 1
    ]
  ]

icosphere :: (Vec3 -> attrs) -> Int -> [Model.Vertex Vec3.Packed attrs]
icosphere mkAttrs n = do
  v <- reverse . concat $ subNormal n icotris_v1
  pure Model.Vertex
    { vPosition = Vec3.Packed v
    , vAttrs    = mkAttrs v
    }

subNormal :: Int -> [[Vec3]] -> [[Vec3]]
subNormal nu tris = map (map Vec3.normalize) $ subdivide nu tris

subdivide :: Int -> [[Vec3]] -> [[Vec3]]
subdivide frequency tris =
  iterate (concatMap subdivideTri) tris !! (frequency - 1)
  where
    subdivideTri = \case
      [v1, v2, v3] ->
        let
          a = midpoint v1 v2
          b = midpoint v2 v3
          c = midpoint v3 v1
        in
          [ [ v1, a, c ]
          , [ v2, b, a ]
          , [ v3, c, b ]
          , [  a, b, c ]
          ]
      _ ->
        error "subdivideTri: not a triangle somehow"

    midpoint a b = (a + b) / 2