root/lang/haskell/aobench/ao.hs @ 30638

module Main () where

import System (getArgs)
import System.IO (withBinaryFile, IOMode(..), hPutStrLn)
import System.Random.Mersenne
import Data.Maybe (catMaybes, listToMaybe)
import Data.List (sort)
import qualified Data.ByteString.Char8 as B

data Vec = Vec {
        x :: Double,
        y :: Double,
        z :: Double
        } deriving (Eq)

data Isect = Isect {
        isect_t :: Double,
        isect_p :: Vec,
        isect_n :: Vec
        } deriving (Eq)

instance Ord Isect where
        i1 < i2         = isect_t i1 < isect_t i2
        i1 > i2         = isect_t i1 > isect_t i2
        i1 <= i2        = isect_t i1 <= isect_t i2
        i1 >= i2        = isect_t i1 >= isect_t i2

data Sphere = Sphere {
        sphere_center :: Vec,
        sphere_radius :: Double
        }

data Plane = Plane {
        plane_p :: Vec,
        plane_n :: Vec
        }

data Ray = Ray {
        ray_org :: Vec,
        ray_dir :: Vec
        }

sq x = x * x

clamp lower upper = min upper . max lower

drand48 = getStdGen >>= random

justIf cond val
        | cond          = Just val
        | otherwise     = Nothing

zerovec = Vec { x = 0, y = 0, z = 0 }

v0 `vadd` v1 = Vec (x v0 + x v1) (y v0 + y v1) (z v0 + z v1)
v0 `vsub` v1 = Vec (x v0 - x v1) (y v0 - y v1) (z v0 - z v1)
v0 `vscale` s = Vec (x v0 * s) (y v0 * s) (z v0 * s)

v0 `vdot` v1 = (x v0 * x v1) + (y v0 * y v1) + (z v0 * z v1)

v0 `vcross` v1 = Vec cx cy cz
        where
                cx = y v0 * z v1 - z v0 * y v1
                cy = z v0 * x v1 - x v0 * z v1
                cz = x v0 * y v1 - y v0 * x v1

vlength c = sqrt $ c `vdot` c

vnormalize c
        | l > 1.0e-17   = c `vscale` (1.0 / l)
        | otherwise             = c
        where
                l = vlength c

ray_sphere_intersect ray sphere =
        justIf (d > 0.0 && t > 0.0) isect
        where
                rs = ray_org ray `vsub` sphere_center sphere
                b = rs `vdot` ray_dir ray
                c = rs `vdot` rs - sq (sphere_radius sphere)
                d = b * b - c
                t = -b - sqrt d
                hitpos = ray_org ray `vadd` (ray_dir ray `vscale` t)
                isect = Isect {
                        isect_t = t,
                        isect_p = hitpos,
                        isect_n = vnormalize $ hitpos `vsub` sphere_center sphere
                        }

ray_plane_intersect ray plane =
        justIf (abs v >= 1.0e-17 && t > 0) isect
        where
                d = -(plane_p plane `vdot` plane_n plane)
                v = ray_dir ray `vdot` plane_n plane
                t = -(ray_org ray `vdot` plane_n plane + d) / v
                isect = Isect {
                        isect_t = t,
                        isect_p = ray_org ray `vadd` (ray_dir ray `vscale` t),
                        isect_n = plane_n plane
                        }

orthoBasis n = (basis0, basis1, basis2)
        where
                basis2 = n
                basis1'
                        | x n < 0.6 && x n > -0.6       = Vec { x = 1, y = 0, z = 0 }
                        | y n < 0.6 && y n > -0.6       = Vec { x = 0, y = 1, z = 0 }
                        | z n < 0.6 && z n > -0.6       = Vec { x = 0, y = 0, z = 1 }
                        | otherwise                                     = Vec { x = 1, y = 0, z = 0 }
                basis0 = vnormalize $ basis1' `vcross` basis2
                basis1 = vnormalize $ basis2 `vcross` basis0

hemisphereDirectionRay p (basis0, basis1, basis2) r1 r2 =
        Ray { ray_org = p, ray_dir = r }
        where
                theta = sqrt r1
                phi = 2.0 * pi * r2
                v = Vec {
                        x = cos phi * theta,
                        y = sin phi * theta,
                        z = sqrt $ 1.0 - theta * theta
                        }
                -- local -> global
                r = Vec {
                        x = x v * x basis0 + y v * x basis1 + z v * x basis2,
                        y = x v * y basis0 + y v * y basis1 + z v * y basis2,
                        z = x v * z basis0 + y v * z basis1 + z v * z basis2
                        }

data Object = ObjSphere Sphere | ObjPlane Plane

findRayObjIntersect ray (ObjSphere sphere) = ray_sphere_intersect ray sphere
findRayObjIntersect ray (ObjPlane  plane)  = ray_plane_intersect  ray plane

ambient_occlusion objs nao_samples isect = do
        randomRays <- mapM chooseRandomRay samples
        let occludedNum = length $ filter id $ map isHitOccluder randomRays
        return $ vscale white $ calcOccludedRatio occludedNum
        where
                ntheta = nao_samples
                nphi = nao_samples
                eps = 0.0001
                p = isect_p isect `vadd` (isect_n isect `vscale` eps)
                basis = orthoBasis $ isect_n isect
                samples = [(j, i) | j <- [0 .. ntheta-1], i <- [0 .. nphi-1]]
                chooseRandomRay _ = do
                        r1 <- drand48
                        r2 <- drand48
                        return $ hemisphereDirectionRay p basis r1 r2
                isHitOccluder ray =
                        not $ null $ catMaybes $ map (findRayObjIntersect ray) objs
                calcOccludedRatio occludedSampleNum =
                        (ntheta * nphi - toEnum occludedSampleNum) / (ntheta * nphi)
                white = Vec { x = 1, y = 1, z = 1 }

screenRay nsubsamples w h x y u v =
        Ray { ray_org = zerovec, ray_dir = vnormalize (Vec { x = px, y = py, z = -1 }) }
        where
                px = (x + (u / nsubsamples) - (w / 2.0)) / (w / 2.0)
                py = -(y + (v / nsubsamples) - (h / 2.0)) / (h / 2.0)

findNearestIsect objs ray =
        listToMaybe $ sort $ catMaybes $ map (findRayObjIntersect ray) objs

data Scene = Scene {
        scene_objs :: [Object]
        }

raytrace shading objs ray = do
        shading $ findNearestIsect objs ray

oversampling trace makeRay nsubsamples = do
        mapM trace subrays >>= return . averageColor
        where
                subpixels = [(u, v) | v <- [0 .. nsubsamples-1], u <- [0 .. nsubsamples-1]]
                subrays = map (uncurry makeRay) subpixels
                averageColor = flip vscale (1.0 / sq nsubsamples) . foldl vadd zerovec

calcPixelColor trace w h nsubsamples x y = do
        oversampling trace makeRay nsubsamples
        where
                makeRay = screenRay nsubsamples w h x y

render w h renderPixel = do
        fimg <- mapM (uncurry renderPixel) screenPixels
        let img = map fpcol2icol fimg
        return img
        where
                screenPixels = [(x, y) | y <- [0 .. h-1], x <- [0 .. w-1]]
                fpcol2icol vec = (f2i (x vec), f2i (y vec), f2i (z vec))
                f2i = clamp 0 255 . floor . (* 255.99)

rendering scene w h nsubsamples nao_samples = do
        render w h $ calcPixelColor (raytrace shading objs) w h nsubsamples
        where
                shading Nothing = return zerovec
                shading (Just isect) = ambient_occlusion objs nao_samples isect
                objs = scene_objs scene

saveppm fname w h img =
        withBinaryFile fname WriteMode $ \fp -> do
                hPutStrLn fp "P6"
                hPutStrLn fp $ show (floor w) ++ " " ++ show (floor h)
                hPutStrLn fp "255"
                B.hPut fp $ B.pack $ map toEnum $ concatMap (\(r,g,b) -> [r,g,b]) img

init_scene =
        Scene {
                scene_objs = [
                        ObjSphere (Sphere { sphere_center = Vec { x = -2.0, y = 0.0, z = -3.5 }, sphere_radius = 0.5 }),
                        ObjSphere (Sphere { sphere_center = Vec { x = -0.5, y = 0.0, z = -3.0 }, sphere_radius = 0.5 }),
                        ObjSphere (Sphere { sphere_center = Vec { x =  1.0, y = 0.0, z = -2.2 }, sphere_radius = 0.5 }),
                        ObjPlane  (Plane  { plane_p = Vec { x = 0.0, y = -0.5, z = 0.0 }, plane_n = Vec { x = 0.0, y = 1.0, z = 0.0 } })
                        ]
                }

getParam idx defaultVal args
        | length args > idx     = args !! idx
        | otherwise                     = defaultVal

main = do
        args <- getArgs
        let width  = read $ getParam 0 "256" args
        let height = read $ getParam 1 "256" args
        let nsubsamples = read $ getParam 2 "2" args
        let nao_samples = read $ getParam 3 "8" args

        img <- rendering init_scene width height nsubsamples nao_samples
        saveppm "ao.ppm" width height img