You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
haskell-raytracer/app/Main.hs

133 lines
4.0 KiB

module Main where
import Control.Monad (mapM)
import Control.Monad.Random
import Data.Ix
import Data.Maybe
import Debug.Trace
import Hittable
import Linear.V2
import Linear.V3
import Linear.Vector
import Material
import Maths
import Ray
import Text.Printf
type Color = (Double, Double, Double)
type Coord = (Int, Int)
type ImageDimensions = (Int, Int)
type ViewPortDimensions = (Double, Double)
-- TODO
-- antialiasing
backgroundColor = V3 0.0 0.0 0.0
scatter :: (RandomGen g) => Ray Double -> Hit -> Rand g (Ray Double)
scatter (Ray o d) (Hit t p n (Lambertian c)) = do
r <- randV3
let target = p ^+^ colorWithNormal (unitVector n) ^+^ r
return $ Ray p (target ^-^ p)
scatter (Ray o d) (Hit t p n (Metal c)) = do
let reflected = reflect (unitVector d) n
return $ Ray p reflected
scatter (Ray o d) (Hit t p n (Glass c)) = do
let refraction_ratio = if isFrontFace (Ray o d) n then 1.0 / 1.5 else 1.0
let ud = unitVector d
let ct = min (ud `dotP` (unitVector n)) 1.0
let st = sqrt (1.0 - ct * ct)
rVal <- rand
if refraction_ratio * st > 1.0 || reflectance ct refraction_ratio < rVal
then do
let reflected = reflect ud n
return $ Ray p reflected
else do
let refracted = refract ud n refraction_ratio
return $ Ray p refracted
getRayColor :: (Shape s, RandomGen g) => [s] -> Int -> Ray Double -> Rand g (V3 Double)
getRayColor ws d r = do
let h = getBestHit (mapMaybe (testHit r) ws)
v <- getRayColor' r h ws d
return v
getRayColor' :: (Shape s, RandomGen g) => Ray Double -> Maybe Hit -> [s] -> Int -> Rand g (V3 Double)
getRayColor' (Ray o d) Nothing _ _ = do return $ colorGradient (unitVector d)
getRayColor' (Ray o d) _ _ 0 = do return $ colorGradient (unitVector d)
getRayColor' (Ray o d) (Just (Hit t p n m)) ws depth = do
let a = attenuation m
ray <- scatter (Ray o d) (Hit t p n m)
nextVal <- getRayColor ws (depth - 1) ray
return (a * nextVal)
colorGradient :: V3 Double -> V3 Double
colorGradient (V3 x y z) = do
let t = 0.5 * (y + 1.0)
(1.0 - t) *^ V3 1.0 1.0 1.0 + t *^ V3 0.5 0.7 1.0
getBestHit [] = Nothing
getBestHit (h : hs) = Just (getBestHit' hs h)
getBestHit' :: [Hit] -> Hit -> Hit
getBestHit' (h : hs) bestHit =
-- always pick closest hit to camera
if root h < root bestHit
then getBestHit' hs h
else getBestHit' hs bestHit
getBestHit' [] hit = hit
colorWithNormal :: V3 Double -> V3 Double
colorWithNormal v = 0.5 *^ (v ^+^ V3 1.0 1.0 1.0)
rayTrace :: (Shape s) => (Int, Int) -> ImageDimensions -> ViewPortDimensions -> [s] -> IO (V3 Double)
rayTrace (x, y) (w, h) (vw, vh) world = do
evalRandIO (getRayColor world 50 $ buildRay (xf * (vw / wf)) (yf * (vh / hf)))
where
xf = fromIntegral x
yf = fromIntegral y
wf = fromIntegral (w - 1)
hf = fromIntegral (h - 1)
clamp :: Double -> Double -> Double -> Double
clamp x mi ma
| x < mi = mi
| x > ma = ma
| otherwise = x
colorToString :: Color -> IO ()
colorToString (x, y, z) = tupToString (xr, yr, zr)
where
-- scale 1
scale = 1.0
xr = truncate ((clamp (sqrt (x * scale)) 0.0 0.999) * 256.0)
yr = truncate ((clamp (sqrt (y * scale)) 0.0 0.999) * 256.0)
zr = truncate ((clamp (sqrt (z * scale)) 0.0 0.999) * 256.0)
tupToString :: (Show a) => (a, a, a) -> IO ()
tupToString (x, y, z) = putStrLn (show x ++ " " ++ show y ++ " " ++ show z)
v3ToColor :: V3 Double -> Color
v3ToColor (V3 x y z) = (x, y, z)
main :: IO ()
main = do
let w = 400
let h = 400
let world =
[ Circle (V3 (0.0) 0.0 (-1.0)) 0.25 (Glass (V3 1.0 1.0 1.0)),
Circle (V3 (-0.5) 0.0 (-1.0)) 0.25 (Lambertian (V3 0.7 0.1 0.1)),
Circle (V3 0.0 (-10.25) (-1.0)) 10 (Lambertian (V3 0.7 0.6 0.5))
]
let coords = reverse $ range ((-1 * (w `div` 2), -1 * (h `div` 2)), ((w `div` 2) - 1, (h `div` 2) - 1)) -- "canvas"
colors <- mapM (\x -> rayTrace x (w, h) (1.0, 1.0) world) coords
let convertedColors = map v3ToColor colors
putStrLn "P3"
putStrLn (show w ++ " " ++ show h)
putStrLn "255"
mapM_ colorToString convertedColors
mapM_ print convertedColors