{-# LANGUAGE OverloadedStrings #-}

module Graphics.Formats.STL.Types
       (
           STL(..),
           Triangle(..),
           Vector,
           triple,
       ) where

import           Control.Monad      (replicateM)
import qualified Data.ByteString    as BS
import           Data.Serialize     (Get, Put, PutM, Serialize (..),
                                     getFloat32le, getWord32le, putFloat32le,
                                     putWord32le, skip)
import           Data.Text          (Text)
import qualified Data.Text          as T
import           Data.Text.Encoding (encodeUtf8)
import           Data.Word          (Word16, Word32)

-- | A representation of an STL file, consisting of a (possibly empty)
-- object name, and a list of triangles.
data STL = STL { STL -> Text
name      :: Text
               , STL -> [Triangle]
triangles :: [Triangle]
               }

-- | A single triangle in STL is represented by a normal vector and
-- three vertices.
data Triangle = Triangle { Triangle -> Maybe Vector
normal   :: Maybe Vector
                         , Triangle -> (Vector, Vector, Vector)
vertices :: (Vector, Vector, Vector)
                         }

type Vector = (Float, Float, Float)

triple :: a -> a -> a -> (a, a, a)
triple :: a -> a -> a -> (a, a, a)
triple a
a a
b a
c = (a
a, a
b, a
c)

--------------------------------------------------------------------------------
-- Binary Output
--------------------------------------------------------------------------------

instance Serialize Triangle where
    get :: Get Triangle
get = Maybe Vector -> (Vector, Vector, Vector) -> Triangle
Triangle (Maybe Vector -> (Vector, Vector, Vector) -> Triangle)
-> Get (Maybe Vector) -> Get ((Vector, Vector, Vector) -> Triangle)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get (Maybe Vector)
getNormal Get ((Vector, Vector, Vector) -> Triangle)
-> Get (Vector, Vector, Vector) -> Get Triangle
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get (Vector, Vector, Vector)
t Get Triangle -> Get () -> Get Triangle
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Int -> Get ()
skip Int
2 where
      t :: Get (Vector, Vector, Vector)
t = (,,) (Vector -> Vector -> Vector -> (Vector, Vector, Vector))
-> Get Vector -> Get (Vector -> Vector -> (Vector, Vector, Vector))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Vector
getVector Get (Vector -> Vector -> (Vector, Vector, Vector))
-> Get Vector -> Get (Vector -> (Vector, Vector, Vector))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Vector
getVector Get (Vector -> (Vector, Vector, Vector))
-> Get Vector -> Get (Vector, Vector, Vector)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Vector
getVector
    put :: Putter Triangle
put (Triangle Maybe Vector
n (Vector
a, Vector
b, Vector
c)) = Maybe Vector -> PutM ()
maybeNormal Maybe Vector
n PutM () -> PutM () -> PutM ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Vector -> PutM ()
v3 Vector
a PutM () -> PutM () -> PutM ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Vector -> PutM ()
v3 Vector
b PutM () -> PutM () -> PutM ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Vector -> PutM ()
v3 Vector
c PutM () -> PutM () -> PutM ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Putter Word16
forall t. Serialize t => Putter t
put (Word16
0x00 :: Word16)

instance Serialize STL where
    get :: Get STL
get = do
        ()
_  <- Get ()
getHeader
        Word32
ct <- Get Word32
getWord32le
        Text -> [Triangle] -> STL
STL Text
"" ([Triangle] -> STL) -> Get [Triangle] -> Get STL
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get Triangle -> Get [Triangle]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
ct) Get Triangle
forall t. Serialize t => Get t
get
    put :: Putter STL
put (STL Text
n [Triangle]
tris) = Putter ByteString
forall t. Serialize t => Putter t
put (Text -> ByteString
header Text
n) PutM () -> PutM () -> PutM ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Putter Word32
putWord32le Word32
ct PutM () -> PutM () -> PutM ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Putter Triangle -> [Triangle] -> PutM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Putter Triangle
forall t. Serialize t => Putter t
put [Triangle]
tris where
      ct :: Word32
      ct :: Word32
ct = Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word32) -> ([Triangle] -> Int) -> [Triangle] -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Triangle] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Triangle] -> Word32) -> [Triangle] -> Word32
forall a b. (a -> b) -> a -> b
$ [Triangle]
tris  -- here's the space leak

-- | header is always exactly 80 characters long
header :: T.Text -> BS.ByteString
header :: Text -> ByteString
header Text
n = [ByteString] -> ByteString
BS.concat [ByteString
lib, ByteString
truncatedName, ByteString
padding] where
  lib :: ByteString
lib = Text -> ByteString
encodeUtf8 Text
"http://hackage.haskell.org/package/STL "
  truncatedName :: ByteString
truncatedName = Int -> ByteString -> ByteString
BS.take (Int
72 Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
BS.length ByteString
lib) (ByteString -> ByteString)
-> (Text -> ByteString) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
encodeUtf8 (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ Text
n
  padding :: ByteString
padding = Int -> Word8 -> ByteString
BS.replicate (Int
72 Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
BS.length ByteString
truncatedName Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
BS.length ByteString
lib) Word8
0x20
-- header _ = BS.replicate 72 0x20 -- cereal adds 8 bytes giving the length of the BS

putFloat :: Float -> Put
putFloat :: Float -> PutM ()
putFloat = Float -> PutM ()
putFloat32le

v3 :: Vector -> PutM ()
v3 :: Vector -> PutM ()
v3 (Float
x,Float
y,Float
z) = Float -> PutM ()
putFloat Float
x PutM () -> PutM () -> PutM ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Float -> PutM ()
putFloat Float
y PutM () -> PutM () -> PutM ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Float -> PutM ()
putFloat Float
z

maybeNormal :: Maybe Vector -> PutM ()
maybeNormal :: Maybe Vector -> PutM ()
maybeNormal Maybe Vector
n = case Maybe Vector
n of
    Maybe Vector
Nothing -> Vector -> PutM ()
v3 (Float
0,Float
0,Float
0)
    Just Vector
n' -> Vector -> PutM ()
v3 Vector
n'

getHeader :: Get ()
getHeader :: Get ()
getHeader = Int -> Get ()
skip Int
80

getFloat :: Get Float
getFloat :: Get Float
getFloat = Get Float
getFloat32le

getVector :: Get Vector
getVector :: Get Vector
getVector = (,,) (Float -> Float -> Float -> Vector)
-> Get Float -> Get (Float -> Float -> Vector)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Float
getFloat Get (Float -> Float -> Vector)
-> Get Float -> Get (Float -> Vector)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Float
getFloat Get (Float -> Vector) -> Get Float -> Get Vector
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Float
getFloat

getNormal :: Get (Maybe Vector)
getNormal :: Get (Maybe Vector)
getNormal = do
    Vector
v <- Get Vector
getVector
    Maybe Vector -> Get (Maybe Vector)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Vector -> Get (Maybe Vector))
-> Maybe Vector -> Get (Maybe Vector)
forall a b. (a -> b) -> a -> b
$ case Vector
v of
        (Float
0,Float
0,Float
0) -> Maybe Vector
forall a. Maybe a
Nothing
        Vector
n'      -> Vector -> Maybe Vector
forall a. a -> Maybe a
Just Vector
n'