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

--import qualified Streaming.Prelude as SP
--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
import qualified Data.Vector.Unboxed.Mutable as MV
import Control.Monad.Primitive

pokeVerxs'' :: PicShads VBO -> [Verx] -> IO (PicShads Int)
pokeVerxs'' vbos = F.foldM $ F.FoldM 
    (\count vx -> pokeVerx vbos count vx >> return (addCountVerx count vx))
    (pure $ pure 0)
    return

pokeVerxs' :: PicShads VBO -> [Verx] -> IO (MV.MVector (PrimState IO) Int)
pokeVerxs' vbos = F.foldM $ F.FoldM 
    (\count vx -> pokeVerx' vbos count vx >> addCountVerx' count vx >> return count)
    (MV.replicate 6 0)
    return

pokeVerxs :: PicShads VBO -> [Verx] -> IO (PicShads Int)
pokeVerxs vbos = F.foldM $ F.FoldM 
    (\count vx -> pokeVerx' vbos count vx >> addCountVerx' count vx >> return count)
    (MV.replicate 6 0)
    vToPicShad

vToPicShad :: MV.MVector (PrimState IO) Int -> IO (PicShads Int)
vToPicShad mv = do
    mapM (MV.read mv) $ PicShads 0 1 2 3 4 5

--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

pokeVerx' :: PicShads VBO -> MV.MVector (PrimState IO) Int -> Verx -> IO ()
{-# INLINE pokeVerx' #-}
pokeVerx' vbos offsets Verx{_vxPos=thePos,_vxCol=theCol,_vxType=theType} = case theType of
    PolyV -> do
        typeOff <- MV.read offsets 0
        poke34 (plusPtr (_vboPtr $ _psPoly vbos) (typeOff * 7 * floatSize)) thePos theCol
    PolyzV x -> do
        typeOff <- MV.read offsets 1
        let thePtr = plusPtr (_vboPtr $ _psPolyz vbos) (typeOff * 8 * floatSize)
        poke34 thePtr thePos theCol
        pokeElemOff thePtr 7 x
    BezV (V4 x y z w)  -> do
        typeOff <- MV.read offsets 2
        let thePtr = plusPtr (_vboPtr $ _psBez vbos)  (typeOff * 11 * floatSize)
        poke34 thePtr thePos theCol
        pokeElemOff thePtr 7 x
        pokeElemOff thePtr 8 y
        pokeElemOff thePtr 9 z
        pokeElemOff thePtr 10 w
    TextV (V2 x y)     -> do
        typeOff <- MV.read offsets 3
        let thePtr = plusPtr (_vboPtr $ _psText vbos) (typeOff * 9 * floatSize)
        poke34 thePtr thePos theCol
        pokeElemOff thePtr 7 x
        pokeElemOff thePtr 8 y
    ArcV (V3 x y z)    -> do
        typeOff <- MV.read offsets 4
        let thePtr = plusPtr (_vboPtr $ _psArc vbos) (typeOff * 10 * floatSize)
        poke34 thePtr thePos theCol
        pokeElemOff thePtr 7 x
        pokeElemOff thePtr 8 y
        pokeElemOff thePtr 9 z
    EllV            -> do
        typeOff <- MV.read offsets 5
        poke34 (plusPtr (_vboPtr $ _psEll vbos) (typeOff * 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}

addCountVerx' :: MV.MVector (PrimState IO) Int -> Verx -> IO ()
{-# INLINE addCountVerx' #-}
addCountVerx' v !Verx{_vxType=theType} = case theType of
    PolyV       -> MV.unsafeModify v (+ 1) 0
    PolyzV _    -> MV.unsafeModify v (+ 1) 1 
    BezV _      -> MV.unsafeModify v (+ 1) 2
    TextV _     -> MV.unsafeModify v (+ 1) 3
    ArcV _      -> MV.unsafeModify v (+ 1) 4
    EllV        -> MV.unsafeModify v (+ 1) 5

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

comLayVerx :: PicShads VBO -> IM.IntMap (PicShads Int) -> Verx -> IO (IM.IntMap (PicShads Int))
--{-# INLINE comLayVerx #-}
comLayVerx vbos counts vx = do
    poke34 thePtr thePos theCol
    pokeVerxType thePtr theType
    return $ addLayCountVerx counts theLayer theType
  where
    theLayer = _vxLayer vx
    layCounts = counts IM.! theLayer
    theType = _vxType vx
    thePos = _vxPos vx
    theCol = _vxCol vx
    basePtr = _vboPtr $ vboFromType vbos theType
    layOff = theLayer * numSubElements
    thePtr = plusPtr basePtr ((intFromType layCounts theType + layOff) * theStride * floatSize)
    theStride = pokeStride theType

pokeVerxType :: Ptr Float -> VertexType -> IO ()
{-# INLINE pokeVerxType #-}
pokeVerxType thePtr vt = case vt of
    PolyzV x            -> pokeElemOff thePtr 7 x
    BezV   (V4 x y z w) -> pokeElemOff thePtr 7 x >> pokeElemOff thePtr 8 y
                            >> pokeElemOff thePtr 9 z >> pokeElemOff thePtr 10 w
    TextV  (V2 x y)     -> pokeElemOff thePtr 7 x >> pokeElemOff thePtr 8 y
    ArcV   (V3 x y z)   -> pokeElemOff thePtr 7 x >> pokeElemOff thePtr 8 y
                            >> pokeElemOff thePtr 9 z
    _ -> return ()

intFromType :: PicShads Int -> VertexType -> Int
{-# INLINE intFromType #-}
intFromType ps vt = case vt of
    PolyV       -> _psPoly ps
    PolyzV{}    -> _psPolyz ps
    BezV{}      -> _psBez ps
    TextV{}     -> _psText ps
    ArcV{}      -> _psArc ps
    EllV        -> _psEll ps

vboFromType :: PicShads VBO -> VertexType -> VBO
{-# INLINE vboFromType #-}
vboFromType ps vt = case vt of
    PolyV       -> _psPoly ps
    PolyzV{}    -> _psPolyz ps
    BezV{}      -> _psBez ps
    TextV{}     -> _psText ps
    ArcV{}      -> _psArc ps
    EllV        -> _psEll ps

pokeLayVerxsFold :: PicShads VBO -> [Verx] -> IO (IM.IntMap (PicShads Int))
pokeLayVerxsFold vbos = F.foldM $ F.FoldM
    (comLayVerx vbos)
    (return $ imLayers 0)
    return

imLayers :: a -> IM.IntMap (PicShads a)
{-# INLINE imLayers #-}
imLayers x = IM.fromList $ zip [0..5] $ repeat ps
  where
    ps = pure x

pokeStride :: VertexType -> Int
{-# INLINE pokeStride #-}
pokeStride PolyV    = 7
pokeStride PolyzV{} = 8
pokeStride BezV{}   = 11
pokeStride TextV{}  = 9
pokeStride ArcV{}   = 10
pokeStride EllV     = 7

addLayCountVerx :: IM.IntMap (PicShads Int) -> Int -> VertexType -> IM.IntMap (PicShads Int)
{-# INLINE addLayCountVerx #-}
addLayCountVerx m lay vt = case vt of
    PolyV       -> m & ix lay . psPoly  +~ 1
    PolyzV _    -> m & ix lay . psPolyz +~ 1
    BezV _      -> m & ix lay . psBez   +~ 1
    TextV _     -> m & ix lay . psText  +~ 1
    ArcV _      -> m & ix lay . psArc   +~ 1
    EllV        -> m & ix lay . psEll   +~ 1
    --PolyV       -> IM.adjust (psPoly  +~ 1) lay m 
    --PolyzV _    -> IM.adjust (psPolyz +~ 1) lay m 
    --BezV _      -> IM.adjust (psBez   +~ 1) lay m 
    --TextV _     -> IM.adjust (psText  +~ 1) lay m 
    --ArcV _      -> IM.adjust (psArc   +~ 1) lay m 
    --EllV        -> IM.adjust (psEll   +~ 1) lay m 

poke224s :: Ptr Float -> [((Point2,Point2),Point4)] -> IO Int
poke224s ptr vals0 = go vals0 0
  where
    go [] n = return n
    go ((((V2 a b),(V2 c d)),(V4 e f g h)):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
        pokeElemOff ptr (off 6) g
        pokeElemOff ptr (off 7) h
        go vals (n+1)
      where
        off i = n*8 + i