haskell-raytracer/app/Main.hs

module Main where

import Camera
import Control.Monad (mapM)
import Control.Monad.Random
import Data.Ix
import Data.Maybe
import Hittable
import IOUtils
import Linear.V3
import Linear.Vector
import Material
import Maths
import Ray

type Coord = (Int, Int)

type ImageDimensions = (Int, Int)

-- TODO
-- antialiasing

-- scatter rays based on the material of the object that was hit
scatter :: (RandomGen g) => Ray -> Hit -> Rand g (Ray)
scatter (Ray o d) (Hit t p n (Lambertian c)) = do
  r <- randV3
  let target = p ^+^ colorWithNormal 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` 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 -> Rand g (V3 Double)
getRayColor ws d r = do
  let h = getBestHit (mapMaybe (testHit r) ws)
  getRayColor' r h ws d

-- get color of pixel according to the closest hit found
getRayColor' :: (Shape s, RandomGen g) => Ray -> 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)

-- used for coloring the background, color based on direction vector
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

-- from a list of possible hits, get the closest possible hit found
getBestHit :: [Hit] -> Maybe Hit
getBestHit [] = Nothing
getBestHit (h : hs) = Just (getBestHit' hs h)

getBestHit' :: [Hit] -> Hit -> Hit
getBestHit' (h : hs) bestHit =
  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)

-- heart of loop, for each x y coordinate pixel, convert it into world space and then run the ray tracing algorithm
-- run [samples] times in order to increase fidelity
rayTrace :: (Shape s) => (Int, Int) -> ImageDimensions -> [s] -> Int -> Camera -> IO (V3 Double)
rayTrace (x, y) (w, h) world samples camera = do
  let maxDepth = 50
  l <- replicateM samples (evalRandIO (getRayColor world maxDepth $ buildRay (xf / wf) (yf / hf) camera))
  return $ foldl (^+^) (V3 0.0 0.0 0.0) l
  where
    xf = fromIntegral x
    yf = fromIntegral y
    wf = fromIntegral (w - 1)
    hf = fromIntegral (h - 1)

main :: IO ()
main = do
  let w = 400
  let h = 400

  let aspect = 1.0 --16.0 / 9.0
  let viewHeight = 2.0
  let viewWidth = aspect * viewHeight
  let focalLength = 1.0

  let camera = makeCamera (V3 0.0 0.0 0.0) viewWidth viewHeight focalLength

  let samples = 25
  let world =
        [ Circle (V3 0.0 0.0 (-1.0)) 0.5 (Glass (V3 1.0 1.0 1.0)),
          Circle (V3 1.0 0.0 (-1.0)) 0.5 (Lambertian (V3 0.7 0.1 0.1)),
          Circle (V3 0.0 (-100.5) (-1.0)) 100 (Lambertian (V3 0.7 0.6 0.5))
        ]
  let coords = [(x, y) | x <- reverse [0 .. w - 1], y <- [0 .. h - 1]]
  colors <- mapM (\x -> rayTrace x (w, h) world samples camera) coords
  let convertedColors = map v3ToColor colors
  -- print ppm header
  putStrLn "P3"
  putStrLn (show w ++ " " ++ show h)
  putStrLn "255"
  -- print each color converted to ppm format string
  mapM_ (colorToString samples) convertedColors
  mapM_ print convertedColors