packages feed

obj-0.1: src/Graphics/Formats/Obj/ObjModel.hs

{-# LANGUAGE NoMonomorphismRestriction #-}
{-# OPTIONS_GHC -Wall #-}
----------------------------------------------------------------------
-- |
-- Module      :  Graphics.Formats.Obj.ObjModel
-- Copyright   :  (c) Anygma BVBA and Thomas Davie 2008
-- License     :  BSD3
-- 
-- Maintainer  :  tom.davie@gmail.com
-- Stability   :  experimental
-- 
-- Describes an Obj Model
----------------------------------------------------------------------
module Graphics.Formats.Obj.ObjModel
       (geometryTests
       ,ObjModel()
       ,geometry,objFile
       ) where

import Graphics.Formats
import Graphics.Formats.Obj.Contents
import Graphics.Formats.Obj.ParserBits (anyOf)
import Graphics.Formats.Mtl.Contents

import Graphics.Rendering.OpenGL

import Test.QuickCheck
import Test.QuickCheck.Checkers
import Test.QuickCheck.Instances

import Data.List
import Data.Array
import Data.Function
import           Data.Map    (Map)
import qualified Data.Map as M

import Control.Monad
import Control.Applicative

data ObjModel = OM BufferSet [Object]
                   deriving (Show)

type BufferSet = (VertexBuffer, TexCoordBuffer, NormalBuffer)

type VertexBuffer   = Array Int (Vertex4   GLfloat)
type TexCoordBuffer = Array Int (TexCoord2 GLfloat)
type NormalBuffer   = Array Int (Normal3   GLfloat)

data Object = OFace     Material [VTriple]
            | OQuad     Material [VTriple]
            | OTriangle Material [VTriple]
            | OLine              [VDouble]
            | OPoint             [Int]
              deriving (Show,Eq,Ord)

instance EqProp ObjModel where
  m =-= m' = property (normalForm m == normalForm m')

normalForm :: ObjModel -> (BufferSet,[Object])
normalForm (OM (vb,tb,nb) os) =
  let lb = fst . bounds
      ub = snd . bounds
      indicies = (uncurry enumFromTo) . bounds
      mkArray a =
        let b = lb a
        in  array (0,ub a - lb a) [(i - b,a ! i) | i <- indicies a]
      stripMat (OFace     _ bs) = OFace emptyMat bs
      stripMat (OTriangle _ bs) = OFace emptyMat bs
      stripMat (OQuad     _ bs) = OFace emptyMat bs
      stripMat x                = x
  in  ((mkArray vb,mkArray tb,mkArray nb)
      ,sort $ offsetObjects (lb vb) (lb nb) (lb tb) (map stripMat os))

instance Eq ObjModel where
  (OM bs os) == (OM bs' os') = bs == bs' && os == os'

instance Arbitrary ObjModel where
  arbitrary = OM <$> ((>**<) arbitrary arbitrary arbitrary) <*> arbitrary
  coarbitrary (OM (v,t,n) os) =
    coarbitrary v . coarbitrary t . coarbitrary n . coarbitrary os

instance Arbitrary Object where
  arbitrary = oneof [OFace     <$> (return whiteMat)
                               <*> (nonEmpty 
                                     ((>**<) positive (maybeGen positive)
                                                      (maybeGen positive)))
                    ,OTriangle <$> (return whiteMat)
                               <*> (setLength 3
                                              ((>**<) positive
                                                      (maybeGen positive)
                                                      (maybeGen positive)))
                    ,OQuad     <$> (return whiteMat)
                               <*> (setLength 4
                                              ((>**<) positive
                                                      (maybeGen positive)
                                                      (maybeGen positive)))
                    ,OLine     <$> (nonEmpty
                                     (positive >*< (maybeGen positive)))
                    ,OPoint    <$> (nonEmpty positive)]
  coarbitrary (OFace     _ n) = coarbitrary n
  coarbitrary (OTriangle _ n) = coarbitrary n
  coarbitrary (OQuad     _ n) = coarbitrary n
  coarbitrary (OLine       n) = coarbitrary n
  coarbitrary (OPoint      n) = coarbitrary n

instance Renderable ObjModel where
  render (OM bs os) = 
    let (textured,unTextured) = partition isTextured os
        texturedGroups =   groupBy ((==)    `on` texID)
                         . sortBy  (compare `on` texID)
                         $ textured
    in do colorMaterial $= Just (FrontAndBack, AmbientAndDiffuse)
          mapM_ (renderOs bs) (unTextured:texturedGroups)

renderOs :: BufferSet -> [Object] -> IO ()
renderOs bs os =
  case os of
    []    -> return ()
    (o:_) -> do setTexturing o
                renderList bs (filter isTriangle os) Triangles
                renderList bs (filter isQuad     os) Quads
                mapM_ (renderObject bs) (filter isPolygon os)

material :: Object -> Maybe Material
material (OFace     m _) = Just m
material (OTriangle m _) = Just m
material (OQuad     m _) = Just m
material _               = Nothing

points :: Object -> Maybe [VTriple]
points (OFace     _ ps) = Just ps
points (OTriangle _ ps) = Just ps
points (OQuad     _ ps) = Just ps
points _                = Nothing

faceTexture :: Object -> Either String TextureObject
faceTexture = maybe (Left "") diffuseTex . material

faceColour :: Object -> IO ()
faceColour f = 
  either (const (setColour (ambientColour  m)
                           (diffuseColour  m)
                           (specularColour m)))
         (const (setColour (Color4 (1.0 :: GLfloat) 1.0 1.0 1.0)
                           (Color4 (1.0 :: GLfloat) 1.0 1.0 1.0)
                           (Color4 (1.0 :: GLfloat) 1.0 1.0 1.0)))
         (diffuseTex m)
  where
    m = maybe (error "Face has no material") id $ material f

texID :: Object -> GLuint
texID = either (const 0) unTexObj . faceTexture
        where
          unTexObj (TextureObject x) = x

setTexturing :: Object -> IO ()
setTexturing o =
  case faceTexture o of
    Left _  ->    texture        Texture2D $= Disabled
    Right t -> do texture        Texture2D $= Enabled
                  textureBinding Texture2D $= Just t

isTextured :: Object -> Bool
isTextured = either (const False) (const True) . faceTexture

renderList :: BufferSet -> [Object] -> PrimitiveMode -> IO ()
renderList  bs x f =
  do renderPrimitive f $ forM_ x (renderCalls bs)

renderCalls :: BufferSet -> Object -> IO ()
renderCalls bs f =
  do faceColour f
     maybe (return ()) (mapM_ (renderOperation3 bs)) $ points f

renderObject :: BufferSet -> Object -> IO ()
renderObject bs (OFace     _ ps) = 
  renderPrimitive Polygon   $ forM_ ps (renderOperation3 bs)
renderObject bs (OTriangle _ ps) =
  renderPrimitive Triangles $ forM_ ps (renderOperation3 bs)
renderObject bs (OQuad     _ ps) =
  renderPrimitive Quads     $ forM_ ps (renderOperation3 bs)
renderObject bs (OLine       ps) =
  renderPrimitive LineStrip $ forM_ ps (renderOperation2 bs)
renderObject bs (OPoint      ps) =
  renderPrimitive Points    $ forM_ ps (renderOperation1 bs)

renderOperation3 :: BufferSet -> VTriple -> IO ()
renderOperation3 (vs,ts,ns) (v,t,n) =
  let nop = maybe top ((>> top) . normal   . (ns !)) n
      top = maybe vop ((>> vop) . texCoord . (ts !)) t
      vop = vertex (vs ! v)
  in nop

renderOperation2 :: BufferSet -> VDouble -> IO ()
renderOperation2 bs (v,t) = renderOperation3 bs (v,Nothing,t)

renderOperation1 :: BufferSet -> Int -> IO ()
renderOperation1 bs v = renderOperation3 bs (v,Nothing,Nothing)
     
isTriangle :: Object -> Bool
isTriangle (OTriangle _ _) = True
isTriangle _               = False

isQuad :: Object -> Bool
isQuad (OQuad _ _ ) = True
isQuad _            = False

isPolygon :: Object -> Bool
isPolygon (OFace _ _) = True
isPolygon _           = False

setColour :: (ColorComponent a, ColorComponent b, ColorComponent c) =>
             Color4 a -> Color4 b -> Color4 c -> IO ()
setColour _ d _ = color d
{-  do materialAmbient  FrontAndBack $= a
       materialDiffuse  FrontAndBack $= d
       materialSpecular FrontAndBack $= s -}

{-renderNormals :: BufferSet -> VTriple -> IO ()
renderNormals (vs,_,ns) (v,_,n) =
  case n of
    Just n' -> vertex (vs ! v) >> vertex ((vs ! v) ..+^^ (ns ! n'))
    Nothing -> return ()

(..+^^) :: Vertex4 GLfloat -> Normal3 GLfloat -> Vertex4 GLfloat
(..+^^) (Vertex4 x y z w) (Normal3 i j k) = Vertex4 (x+i) (y+j) (z+k) w-}

geometry :: ObjFile -> MtlFile -> ObjModel
geometry (OF f) mtls =
  OM bs (unsmoothedObjects ++ smoothedObjects)
  where
    vertexBuffer   = listArray (1,length vertexList    ) vertexList
    normalBuffer   = listArray (1,length fullNormalList) fullNormalList
    texCoordBuffer = listArray (1,length texCoordList  ) texCoordList
    vertexList     = map vToVertex    . filter isVertex   $ f
    normalList     = map vnToNormal   . filter isNormal   $ f
    fullNormalList = normalList ++ concat newNormals
    texCoordList   = map vtToTexCoord . filter isTexCoord $ f
    bs             = (vertexBuffer,texCoordBuffer,normalBuffer)
    unsmoothedObjects = maybe [] id (M.lookup Nothing objects)
    (smoothedObjects,newNormals,_) =
      foldr (smoothGroup vertexBuffer)
            ([],[],length normalList + 1)
            smoothingGroups
    smoothingGroups = M.elems . M.filterWithKey (\k _ -> k /= Nothing)
                              $ objects
    objects = fst6 $ foldl (addObj mtls)
                           (M.empty,whiteMat,Nothing,1,1,1)
                           (filter (anyOf [isNormal,isTexCoord,isVertex
                                          ,isObject,isUseMtl,isSmoothG]) f)

fst6 :: (a,b,c,d,e,f) -> a
fst6 (x,_,_,_,_,_) = x

addObj :: MtlFile
       -> (Map (Maybe Int) [Object],Material,Maybe Int,Int,Int,Int)
       -> Statement
       -> (Map (Maybe Int) [Object],Material,Maybe Int,Int,Int,Int)
addObj (MF mtls) (os,_,sg,vc,nc,tc) (UseMtl m) =
  maybe (error ("Material not found: " ++ show m))
        (\mtl' -> (os,mtl',sg,vc,nc,tc))
        (M.lookup m mtls)
addObj _   (os,cm,sg,vc,nc,tc) (P ps) =
  (M.insertWith (++) sg [OPoint (map (mkAbs vc) ps)] os, cm,sg,vc,nc,tc)
addObj _   (os,cm,sg,vc,nc,tc) (L ps) =
  (M.insertWith (++) sg [OLine (absoluteRefs2 vc tc ps)] os, cm,sg,vc,nc,tc)
addObj _   (os,cm,sg,vc,nc,tc) (F ps) =
  case length ps of
    3 -> (M.insertWith (++) sg [OTriangle cm (absoluteRefs3 vc nc tc ps)] os
         ,cm,sg,vc,nc,tc)
    4 -> (M.insertWith (++) sg [OQuad     cm (absoluteRefs3 vc nc tc ps)] os
         ,cm,sg,vc,nc,tc)
    _ -> (M.insertWith (++) sg [OFace     cm (absoluteRefs3 vc nc tc ps)] os
         ,cm,sg,vc,nc,tc)
addObj _   (os,cm,sg,vc,nc,tc) (V _ _ _ _) = (os,cm,sg,vc+1,nc  ,tc  )
addObj _   (os,cm,sg,vc,nc,tc) (VN _ _ _)  = (os,cm,sg,vc  ,nc+1,tc  )
addObj _   (os,cm,sg,vc,nc,tc) (VT _ _ _)  = (os,cm,sg,vc  ,nc  ,tc+1)
addObj _   (os,cm,_ ,vc,nc,tc) (SG s)      = (os,cm,s ,vc  ,nc  ,tc  )
addObj _   (os,cm,sg,vc,nc,tc) _           = (os,cm,sg,vc  ,nc  ,tc  )

smoothGroup :: VertexBuffer -> [Object]
                            -> ([Object],[[Normal3 GLfloat]],Int)
                            -> ([Object],[[Normal3 GLfloat]],Int)
smoothGroup vb g (os,ns,nns) =
  (gos ++ os,gns : ns, nns + (length gns))
  where
    gos = map (applyNormals normalMap) g
    normalMap = M.fromList $ zip vs [nns..]
    gns = map (makeNormal vb g) vs
    vs = concatMap objVerticies g

applyNormals :: Map Int Int -> Object -> Object
applyNormals m (OFace     mat vs) = OFace     mat (appNorms m vs)
applyNormals m (OTriangle mat vs) = OTriangle mat (appNorms m vs)
applyNormals m (OQuad     mat vs) = OQuad     mat (appNorms m vs)
applyNormals _ x                  = x

appNorms :: Map Int Int -> [VTriple] -> [VTriple]
appNorms m vs = map (\(v,t,n) -> case n of
                                   Just _  -> (v,t,n           )
                                   Nothing -> (v,t,M.lookup v m))
                    vs

makeNormal :: VertexBuffer -> [Object] -> Int -> Normal3 GLfloat
makeNormal vb os =
 uncurry3 Normal3 . averageVec
                   . map (uncurry crossProduct . createVectors
                                               . lookupVerticies vb)
                   . (findVertexNeighbors os)
  where
    findVertexNeighbors :: [Object] -> Int -> [(Int,Int,Int)]
    findVertexNeighbors objs v = foldr (findVertexPair v) [] objs
    
    findVertexPair :: Int -> Object -> [(Int,Int,Int)] -> [(Int,Int,Int)]
    findVertexPair v (OFace     _ vtripples) x = findVP v vtripples x
    findVertexPair v (OTriangle _ vtripples) x = findVP v vtripples x
    findVertexPair v (OQuad     _ vtripples) x = findVP v vtripples x
    findVertexPair _ _                       x = x

    findVP v vtripples ns =
      maybe ns (:ns) (find3 v (map trippleVertex vtripples))
      where
        find3  :: Int -> [Int] -> Maybe (Int,Int,Int)
        find3  x ys = find3' x (length ys) $ cycle ys
        find3' :: Int -> Int -> [Int] -> Maybe (Int,Int,Int)
        find3' _ 0 _  = Nothing
        find3' x r (l:c:n:ys)
          | x == c    = Just (l,c,n)
          | otherwise = find3' x (r-1) (c:n:ys)
        find3' _ _ _  =
          error "find3' called incorrectly.  Input list not infinite."
    
    lookupVerticies :: VertexBuffer
                    -> (Int,Int,Int)
                    -> (Vertex4 GLfloat,Vertex4 GLfloat,Vertex4 GLfloat)
    lookupVerticies buff (a,b,c) = (buff ! a, buff ! b, buff ! c)
    
    createVectors :: (Vertex4 GLfloat,Vertex4 GLfloat,Vertex4 GLfloat)
                  -> ((GLfloat,GLfloat,GLfloat),(GLfloat,GLfloat,GLfloat))
    createVectors (a,b,c) = (a .-. b,c .-. b)
    
    crossProduct :: Num a => (a,a,a) -> (a,a,a) -> (a,a,a)
    crossProduct (x,y,z) (x',y',z') =
      (y * z' - z * y', z * x' - x * z', x * y' - y * x')
    
    averageVec :: Floating a => [(a,a,a)] -> (a,a,a)
    averageVec [] = error "Average vectors: no vectors to average."
    averageVec xs = normalise ((sumVec (map normalise xs)) ^/ (fromIntegral $ length xs))
    
    normalise :: Floating a => (a,a,a) -> (a,a,a)
    normalise x = x ^/ (mag x)
    
    sumVec = foldr1 (^+^)
    
    mag (x,y,z) = sqrt (x * x + y * y + z * z)
    
    (^+^) (x,y,z) (x',y',z') = (x + x', y + y', z + z')
    (.-.) (Vertex4 x y z w) (Vertex4 x' y' z' w') = 
      (x * w - x' * w', y * w - y' * w', z * w - z' * w')
    
    v ^/ s = v ^* (1 / s)
    (x,y,z) ^* s = (x * s, y * s, z * s)

objVerticies :: Object -> [Int]
objVerticies (OFace     _ vs) = map trippleVertex vs
objVerticies (OTriangle _ vs) = map trippleVertex vs
objVerticies (OQuad     _ vs) = map trippleVertex vs
objVerticies _ = []

mkAbs :: Int -> Int -> Int
mkAbs c x = if x < 0 then c+x else x

absoluteRefs2 :: Int -> Int -> [(Int,Maybe Int)] -> [(Int,Maybe Int)]
absoluteRefs2 c c' =
  uncurry zip . (\(x,y) -> (map (mkAbs c) x
                           ,map (liftM (mkAbs c')) y)) . unzip

absoluteRefs3 :: Int -> Int -> Int -> [(Int,Maybe Int,Maybe Int)]
                                   -> [(Int,Maybe Int,Maybe Int)]
absoluteRefs3 c c' c'' =
  uncurry3 zip3 . (\(x,y,z) -> (map (mkAbs c) x
                               ,map (liftM (mkAbs c')) y
                               ,map (liftM (mkAbs c'')) z)) . unzip3

uncurry3 :: (a -> b -> c -> d) -> (a,b,c) -> d
uncurry3 f (x,y,z) = f x y z

trippleVertex :: (Int,Maybe Int,Maybe Int) -> Int
trippleVertex (v,_,_) = v

offsetObjects :: Int -> Int -> Int -> [Object] -> [Object]
offsetObjects vo no to = map (offsetObj vo no to)
offsetObj :: Int -> Int -> Int -> Object -> Object
offsetObj vo no to (OFace    m ts) =
  OFace m     $ map (\(x,y,z) -> (x-vo
                                 ,y >>= return . ((flip (-)) no)
                                 ,z >>= return . ((flip (-)) to))) ts
offsetObj vo no to (OTriangle m ts) =
  OTriangle m $ map (\(x,y,z) -> (x-vo
                                 ,y >>= return . ((flip (-)) no)
                                 ,z >>= return . ((flip (-)) to))) ts
offsetObj vo no to (OQuad     m ts) =
  OQuad m     $ map (\(x,y,z) -> (x-vo
                                 ,y >>= return . ((flip (-)) no)
                                 ,z >>= return . ((flip (-)) to))) ts
offsetObj vo _  to (OLine       ts) =
  OLine $ map (\(x,z) -> (x-vo
                         ,z >>= return . ((flip (-)) to))) ts
offsetObj vo _  _  (OPoint      ts) =
  OPoint $ map ((flip (-)) vo) ts

objFile :: ObjModel -> ObjFile
objFile (OM (vb,tb,nb) os) =
  OF $ concat [writeBuffer vertexToV    vb
              ,writeBuffer texCoordToVT tb
              ,writeBuffer normalToVN   nb
              ,map writeObject (offsetObjects (-1) (-1) (-1) os)]
  where
    writeBuffer :: (a -> Statement) -> Array Int a -> [Statement]
    writeBuffer f b =
      map f $ elems b
    writeObject :: Object -> Statement
    writeObject (OFace     _ fs) = F fs
    writeObject (OTriangle _ fs) = F fs
    writeObject (OQuad     _ fs) = F fs
    writeObject (OLine ls)       = L ls
    writeObject (OPoint ps)      = P ps

vnToNormal :: Statement -> Normal3 GLfloat
vnToNormal (VN i j k)   = Normal3 i j k
vnToNormal _            = error "Obj statement was not a normal."

vtToTexCoord :: Statement -> TexCoord2 GLfloat
vtToTexCoord (VT u v _) = TexCoord2 u v
vtToTexCoord _          = error "Obj statement was not a texture coordinate."

vToVertex :: Statement -> Vertex4 GLfloat
vToVertex (V x y z w)   = Vertex4 x y z w
vToVertex _             = error "Obj statement was not a vertex."

vertexToV :: Vertex4 GLfloat -> Statement
vertexToV (Vertex4 x y z w) = V x y z w

normalToVN :: Normal3 GLfloat -> Statement
normalToVN (Normal3 i j k) = VN i j k

texCoordToVT :: TexCoord2 GLfloat -> Statement
texCoordToVT (TexCoord2 u v) = VT u v 0.0

prop_geomUnGeom :: ObjModel -> Property
prop_geomUnGeom x =
  ((((flip geometry) (MF M.empty)) . objFile $ x) =-= x)

geometryTests :: IO ()
geometryTests = do putStr "prop_geomUnGeom:   "
                   quickCheck prop_geomUnGeom