Ticket #2284: ObjView.hs

-- Scott Dillard, sedillard@gmail.com

{-# OPTIONS -fglasgow-exts #-}

module Main where

import System.IO
import System.Exit
import Foreign
import Foreign.C.Types
import Control.Monad
import Control.Parallel

import Data.List
import Data.Maybe
import Data.Char
import Data.Array

import qualified Data.Sequence as Seq
import qualified Data.Foldable as Fold
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy.Char8 as C

import Data.Map (Map)
import qualified Data.Map as Map

import Data.Set (Set)
import qualified Data.Set as Set

import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap

import Graphics.UI.GLUT hiding (normalize) --why is this exported by glut?
import Graphics.Rendering.OpenGL.GL hiding (normalize,rotate)
import qualified Graphics.Rendering.OpenGL.GL as GL

import Debug.Trace


main = 
  do

  --initialize glut
  initialDisplayCapabilities $= [ With DisplayRGB, With DisplayDouble, With DisplayDepth ]
  (name,args) <- getArgsAndInitialize
  when (null args) (error "Give me an obj file to load")

  --load mesh
  (vertList,faceList) <- loadObj (head args)
  putStrLn $ "Loaded " ++ (head args) ++ ", " ++ show(length vertList) ++ " vertices, " ++ show(length faceList) ++ " faces." 

  vertList <- return $  map vecFromList vertList

  putStrLn "use x,y,z keys to rotate" 


  let
    nverts = length vertList
    vertArray = listArray (0, nverts-1) vertList
    imesh = trace "BUILDING MESH" $ memoMesh (nextMapFromLists faceList)
    faces = findFaces imesh
    normals = computeNormals (vertArray!) faces
    normArray = listArray (0, nverts-1) $ IntMap.elems normals
    lo = minimum vertList
    hi = maximum vertList
    center = lo `vadd` (0.5 `vmult` (hi`vsub`lo))
    scaleFactor = 1 / vfold max (hi`vsub`lo) 

  let
    reshape size@(Size w h) = 
      do
      viewport $= ( Position 0 0 , size )

      matrixMode $= Projection
      loadIdentity
      let ratio = (fromIntegral h) / (fromIntegral w)
      ortho (-1) 1 (-ratio) ratio (-2) 2

      matrixMode $= Modelview 0
      loadIdentity
      scale scaleFactor scaleFactor scaleFactor
      case center of (Vec3 x y z) -> translate $ Vector3 x y z

      -- init opengl
      lighting $= Enabled
      light (Light 0) $= Enabled
      depthFunc $= Just Less
      clearColor $= Color4 0 0 0 0
      
    display = 
      do
      clear [ColorBuffer, DepthBuffer]
      color $ Color3 (0.5) (0.5::Float) 0
      colorMaterial $= (Just (FrontAndBack,AmbientAndDiffuse))
      lighting $= Enabled
      GL.normalize $= Enabled
      forM_ faces $ \f ->
        renderPrimitive Polygon $
          forM_ (face f) $ \(Edge i _ _) -> do
            case normArray!i of (Vec3 x y z) -> normal $ Normal3 x y z
            case vertArray!i of (Vec3 x y z) -> vertex $ Vertex3 x y z
      checkGLErrors "Display"

    keyMouse key keyState mods (Position x y) = 
      do
      case key of 
        Char 'q' -> exitWith ExitSuccess
        Char 'x' -> GL.rotate 1 (Vector3 1 0 (0::Float))
        Char 'y' -> GL.rotate 1 (Vector3 0 1 (0::Float))
        Char 'z' -> GL.rotate 1 (Vector3 0 0 (1::Float))
        _ -> return ()

  --create a window and bind callbacks
  createWindow name
  displayCallback $= (display >> swapBuffers) 
  keyboardMouseCallback $= Just (\a b c d -> do keyMouse a b c d; postRedisplay Nothing)
  reshapeCallback $= Just reshape

  mainLoop





--the normal at a vertex is the sum of the normals of the faces incident on it
computeNormals :: (Int -> Vec3) -> [Edge Int] -> IntMap Vec3
computeNormals vf faces = IntMap.map normalize $ normals (Seq.fromList faces)
  where
    normal e = 
      let inds@(i:j:k:_) = map edgeOrg (face e)
          [vi,vj,vk] = map vf [i,j,k]
          --n = normalize $ cross (vj`vsub`vi) (vk`vsub`vi) --uniform
          n = cross (vj`vsub`vi) (vk`vsub`vi) --area weighted
      in  zip inds (repeat n)
    normals s =
      case Seq.length s of
        n | n < 100  -> foldr (\(i,n) m -> IntMap.insertWith vadd i n m) 
                               IntMap.empty (concatMap normal $ Fold.toList s)
          | otherwise -> let (a,b) = Seq.splitAt (n`div`2) s 
                             a' = normals a
                             b' = normals b
                         in  par a' (seq b' (IntMap.unionWith vadd a' b' ))
   
checkGLErrors where_ = 
  do
  errs <- get errors
  when (not (null errs)) $ (do putStrLn (where_++": "); (mapM_ print errs))
  

  
  

--half edge mesh
data Edge v = Edge { edgeOrg :: v , edgeSym :: Edge v , edgeNext :: Edge v }

instance Eq v => Eq (Edge v) where
  (Edge i (Edge j _ _) _) == (Edge x (Edge y _ _) _) = (i,j) == (x,y)

instance Ord v => Ord (Edge v) where
  (Edge i (Edge j _ _) _) < (Edge x (Edge y _ _) _) = (i,j) < (x,y)

instance Show v => Show (Edge v) where
  show (Edge i _ _) = show i

--given a map from nextMapFromLists, produce a memoized Edge mesh
memoMesh :: Map (Int,Int) (Int,Int) -> Edge Int
memoMesh nexts = head $ Map.elems ties
  where
    ties = Map.mapWithKey (\ij _ -> make ij) nexts
    lookup ij = fromJust $ Map.lookup ij ties
    make ij@(i,j) = Edge i (lookup (j,i)) 
                           (lookup . fromJust $ Map.lookup ij nexts)

faceRing e = e : faceRing (edgeNext e)
orgRing  e = e : orgRing  (edgeNext . edgeSym $ e)
face e = e : takeWhile (/= e) (faceRing $ edgeNext e)
edgeVerts (Edge i (Edge j _ _) _) = (i,j)
edgeDest (Edge _ (Edge j _ _) _) = j


--given an Edge mesh, make a list of 1 Edge per face
findFaces e = faces' [e] Set.empty
  where
    faces' [] seen = []
    faces' (e:es) seen = 
      let ij = edgeVerts e in
      case Set.member ij seen of
        True -> faces' es seen
        False ->
          let f = face e in
          e : faces' (map edgeSym f ++ es) 
                     (foldr Set.insert seen (map edgeVerts f))





--given a list of faces, where each face is a list of vertex indicies,
--construct a map from an edge (Int,Int) to the next edge around the left face
--of that edge, in ccw order

nextMapFromLists :: [[Int]] -> Map (Int,Int) (Int,Int)
nextMapFromLists faces = make (Seq.fromList faces)
  where
    pairs xs = xs `zip` (tail . cycle $ xs)
    make s =
      case Seq.length s of
        n | n<100 -> foldr (\(e,f) m -> Map.insert e f m) 
                           Map.empty (Fold.concatMap (pairs . pairs) s)
          | otherwise -> 
              let (a,b) = Seq.splitAt (n`div`2) s
                  a' = make a
                  b' = make b
              in  par a' $ seq b' $ Map.union (make a) (make b)




--load an Obj model file (old alias/wavefront format, blender and wings3d can
--read/write it. This is barebones, but works with wings3d's exported files)

loadObj :: String -> IO ([[Double]],[[Int]])
loadObj filename =
  do
  txt <- C.hGetContents =<< openFile filename ReadMode
  let
    (verts,faces) = foldl' doLine ([],[]) $ C.lines txt
    doLine :: ([[Double]],[[Int]]) -> C.ByteString -> ([[Double]],[[Int]])
    doLine (vls,fls) line =
      case C.head line of 
        'v' -> 
          case C.head $ C.tail line of 
            ' ' -> (map readDouble (C.words ((C.tail . C.tail)line)) : vls, fls)
            _ -> (vls,fls)
        'f' -> 
          (vls, (map (pred . (fst . fromJust . C.readInt) . (C.takeWhile isDigit)) ((C.words . C.tail) line)):fls )
        _ -> (vls,fls)
  return (reverse verts,faces)

readDouble :: C.ByteString -> Double
readDouble ls = unsafePerformIO $ B.useAsCString s $ \cstr ->
    realToFrac `fmap` c_strtod cstr nullPtr
  where
    s = B.concat . C.toChunks $ ls

foreign import ccall unsafe "static stdlib.h strtod" c_strtod
    :: Ptr CChar -> Ptr (Ptr CChar) -> IO CDouble






data Vec3 = Vec3 !Double !Double !Double deriving (Eq,Ord,Show)
cross (Vec3 ux uy uz) (Vec3 vx vy vz) = Vec3 (uy*vz-uz*vy) (uz*vx-ux*vz) (ux*vy-uy*vx) 
normalize (Vec3 x y z) = case sqrt(x*x+y*y+z*z) of n -> Vec3 (x/n) (y/n) (z/n)
vfold f (Vec3 x y z) = x`f`y`f`z
vadd (Vec3 x y z) (Vec3 a b c) = Vec3 (x+a) (y+b) (z+c)
vsub (Vec3 x y z) (Vec3 a b c) = Vec3 (x-a) (y-b) (z-c)
vmult s (Vec3 a b c) = Vec3 (s*a) (s*b) (s*c)
vecFromList [x,y,z] = Vec3 x y z