{-# LANGUAGE TupleSections #-}
{-# LANGUAGE BangPatterns #-}
module Shader.Poke
    ( pokeArrayOff
    , pokeVX
    , pokeVX'
    , pokePoint3s
    , pokePoint33s
    , pokeVerxs
    , pokeLayVerxs
    ) where
import Shader.Data
import Shader.Parameters
import Picture.Data
import Geometry.Data

--import Data.Maybe
--import Data.List
import Foreign
--import Control.Monad
--import qualified Control.Foldl as F
import qualified Data.IntMap.Strict as IM
import Control.Lens

getterPicShad :: VertexType -> PicShads a -> a
--{-# INLINE getterPicShad #-}
getterPicShad PolyV    = _psPoly
getterPicShad PolyzV{} = _psPolyz
getterPicShad BezV{}   = _psBez
getterPicShad TextV{}  = _psText
getterPicShad ArcV{}   = _psArc
getterPicShad EllV     = _psEll

setterPicShad :: VertexType -> ASetter (PicShads a) (PicShads a) a a
setterPicShad PolyV    = psPoly
setterPicShad PolyzV{} = psPolyz
setterPicShad BezV{}   = psBez
setterPicShad TextV{}  = psText
setterPicShad ArcV{}   = psArc
setterPicShad EllV     = psEll

pokeVX' :: PicShads VShader -> PicShads Int -> Verx -> IO (PicShads Int)
pokeVX' shads count vx = do
    let f = setterPicShad (_vxType vx)
        g = getterPicShad (_vxType vx)
        vshad = g shads
        n = g count
        theVBO = _vaoVBO $ _vshaderVAO vshad
        ptr = _vboPtr theVBO
        stride = _vboStride theVBO
    pokeArrayOff ptr (stride * n) (toFls' vx)
    return $ count & f +~ 1

pokeVerxs :: PicShads VBO -> [Verx] -> IO (PicShads Int)
pokeVerxs vbos vxs0 = go vxs0 (pure 0)
  where
    go [] count = return count
    go (!vx:vxs) count = pokeVerx vbos count vx >> go vxs (addCountVerx count vx)

-- this is very brittle, but want to optimise speed if possible here
pokeVerx :: PicShads VBO -> PicShads Int -> Verx -> IO ()
{-# INLINE pokeVerx #-}
pokeVerx vbos offsets Verx{_vxPos=thePos,_vxCol=theCol,_vxType=theType} = case theType of
    PolyV -> poke34 (plusPtr (_vboPtr $ _psPoly vbos) (_psPoly offsets * 7 * floatSize)) thePos theCol
    PolyzV x        -> poke34 thePtr thePos theCol
        >> pokeElemOff thePtr 7 x
      where
        thePtr = plusPtr (_vboPtr $ _psPolyz vbos) (_psPolyz offsets * 8 * floatSize)
    BezV (V4 x y z w)  -> poke34 thePtr thePos theCol
        >> pokeElemOff thePtr 7 x
        >> pokeElemOff thePtr 8 y
        >> pokeElemOff thePtr 9 z
        >> pokeElemOff thePtr 10 w
      where
        thePtr = plusPtr (_vboPtr $ _psBez vbos)  (_psBez offsets * 11 * floatSize)
    TextV (V2 x y)     -> poke34 thePtr thePos theCol
        >> pokeElemOff thePtr 7 x
        >> pokeElemOff thePtr 8 y
      where
        thePtr = plusPtr (_vboPtr $ _psText vbos) (_psText offsets * 9 * floatSize)
    ArcV (V3 x y z)    -> poke34 thePtr thePos theCol
        >> pokeElemOff thePtr 7 x
        >> pokeElemOff thePtr 8 y
        >> pokeElemOff thePtr 9 z
      where
        thePtr = plusPtr (_vboPtr $ _psArc vbos) (_psArc offsets * 10 * floatSize)
    EllV            -> poke34 (plusPtr (_vboPtr $ _psEll vbos) (_psEll offsets * 7 * floatSize)) thePos theCol

poke34 :: Ptr Float -> Point3 -> Point4 -> IO ()
{-# INLINE poke34 #-}
poke34 ptr (V3 a b c) (V4 d e f g) = do
    pokeElemOff ptr 0 a
    pokeElemOff ptr 1 b
    pokeElemOff ptr 2 c
    pokeElemOff ptr 3 d
    pokeElemOff ptr 4 e
    pokeElemOff ptr 5 f
    pokeElemOff ptr 6 g

addCountVerx :: PicShads Int -> Verx -> PicShads Int
{-# INLINE addCountVerx #-}
addCountVerx ps@PicShads 
    { _psPoly   = sPoly 
    , _psPolyz  = sPolyz
    , _psBez    = sBez  
    , _psText   = sText 
    , _psArc    = sArc  
    , _psEll    = sEll  
    }   Verx{_vxType=theType} = case theType of
    PolyV       -> ps {_psPoly = sPoly + 1} 
    PolyzV _    -> ps {_psPolyz = sPolyz + 1} 
    BezV _      -> ps {_psBez  = sBez  + 1}
    TextV _     -> ps {_psText = sText + 1}
    ArcV _      -> ps {_psArc  = sArc  + 1}
    EllV        -> ps {_psEll  = sEll  + 1}


pokeVX :: PicShads VShader -> IM.IntMap (PicShads Int) -> Verx -> IO (IM.IntMap (PicShads Int))
pokeVX shads laycounts vx = do
    let theType = _vxType vx
        offset = _vxLayer vx
        counts = laycounts IM.! offset
        getF = getterPicShad theType
        vshad = getF shads
        n = getF counts
        theVBO = _vaoVBO $ _vshaderVAO vshad
        ptr = _vboPtr theVBO
        stride = _vboStride theVBO
    pokeArrayOff ptr (stride * (n+offset * numSubElements)) (toFls' vx)
    return $ laycounts & ix offset . setterPicShad theType +~ 1


toFls' :: Verx -> [Float]
toFls' vx = flat3 (_vxPos vx) ++ flat4 (_vxCol vx) ++ f (_vxType vx)
  where
    f (PolyzV x) = [x]
    f (BezV x) = flat4 x
    f (TextV x) = flat2 x
    f (ArcV x) = flat3 x
    f _ = []

pokeArrayOff :: Storable a => Ptr a -> Int -> [a] -> IO ()
--{-# INLINE pokeArrayOff #-}
--pokeArrayOff ptr i = zipWithM_ (pokeElemOff ptr) [i..]
pokeArrayOff ptr i = pokeArray (plusPtr ptr (floatSize * i))

pokePoint33s :: Ptr Float -> [(Point3,Point3)] -> IO Int
pokePoint33s ptr vals0 = go vals0 0
  where
    go [] n = return n
    go ( ((V3 a b c),(V3 d e f)):vals) n = do
        pokeElemOff ptr (off 0) a
        pokeElemOff ptr (off 1) b
        pokeElemOff ptr (off 2) c
        pokeElemOff ptr (off 3) d
        pokeElemOff ptr (off 4) e
        pokeElemOff ptr (off 5) f
        go vals (n+1)
      where
        off i = n*6 + i

pokePoint3s :: Ptr Float -> [Point3] -> IO Int
pokePoint3s ptr vals0 = go vals0 0
  where
    go [] n = return n
    go ( (V3 a b c):vals) n = do
        pokeElemOff ptr (off 0) a
        pokeElemOff ptr (off 1) b
        pokeElemOff ptr (off 2) c
        go vals (n+1)
      where
        off i = n*3 + i

pokeLayVerxs :: PicShads VBO -> [Verx] -> IO (IM.IntMap (PicShads Int))
pokeLayVerxs vbos vxs0 = go vxs0 (IM.fromList $ (, pure 0) <$> [0..5])
  where
    go [] count = return count
    go (!vx:vxs) count = pokeLayVerx vbos count vx >> go vxs (addLayCountVerx count vx)

-- this is very brittle, but want to optimise speed if possible here
pokeLayVerx :: PicShads VBO -> IM.IntMap (PicShads Int) -> Verx -> IO ()
{-# INLINE pokeLayVerx #-}
pokeLayVerx vbos imoffsets Verx{_vxPos=thePos,_vxCol=theCol,_vxType=theType, _vxLayer = theLay} = case theType of
    PolyV -> poke34 (plusPtr (_vboPtr $ _psPoly vbos) ((_psPoly offsets + layOff) * 7 * floatSize)) thePos theCol
    PolyzV x        -> poke34 thePtr thePos theCol
        >> pokeElemOff thePtr 7 x
      where
        thePtr = plusPtr (_vboPtr $ _psPolyz vbos) ((_psPolyz offsets + layOff) * 8 * floatSize)
    BezV (V4 x y z w)  -> poke34 thePtr thePos theCol
        >> pokeElemOff thePtr 7 x
        >> pokeElemOff thePtr 8 y
        >> pokeElemOff thePtr 9 z
        >> pokeElemOff thePtr 10 w
      where
        thePtr = plusPtr (_vboPtr $ _psBez vbos)  ((_psBez offsets + layOff) * 11 * floatSize)
    TextV (V2 x y)     -> poke34 thePtr thePos theCol
        >> pokeElemOff thePtr 7 x
        >> pokeElemOff thePtr 8 y
      where
        thePtr = plusPtr (_vboPtr $ _psText vbos) ((_psText offsets + layOff) * 9 * floatSize)
    ArcV (V3 x y z)    -> poke34 thePtr thePos theCol
        >> pokeElemOff thePtr 7 x
        >> pokeElemOff thePtr 8 y
        >> pokeElemOff thePtr 9 z
      where
        thePtr = plusPtr (_vboPtr $ _psArc vbos) ((_psArc offsets + layOff) * 10 * floatSize)
    EllV            -> poke34 (plusPtr (_vboPtr $ _psEll vbos) ((_psEll offsets + layOff) * 7 * floatSize)) thePos theCol
  where
    layOff = theLay * numSubElements
    offsets = imoffsets IM.! theLay

addLayCountVerx :: IM.IntMap (PicShads Int) -> Verx -> IM.IntMap (PicShads Int)
{-# INLINE addLayCountVerx #-}
addLayCountVerx m vx = IM.adjust f (_vxLayer vx) m
  where
    f = flip addCountVerx vx