haskell-raytracer/app/Main.hs

module Main where

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

type Coord = (Int, Int)

type ImageDimensions = (Int, Int)

type ViewPortDimensions = (Double, Double)

-- 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 (unitVector n) ^+^ r
  return $ Ray p (target ^-^ p)
scatter (Ray o d) (Hit t p n (Metal c)) = do
  let reflected = reflect (unitVector d) (unitVector 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 (unitVector n)
      return $ Ray p reflected
    else do
      let refracted = refract ud (unitVector 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 -> ViewPortDimensions -> [s] -> Int -> IO (V3 Double)
rayTrace (x, y) (w, h) (vw, vh) world samples = do
  let maxDepth = 50
  l <- replicateM samples (evalRandIO (getRayColor world maxDepth $ (buildRay (xf * (vw / wf)) (yf * (vh / hf)))))
  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 samples = 25
  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 = [(x, y) | x <- reverse $ [-1 * (w `div` 2) .. (w `div` 2) - 1], y <- [-1 * (h `div` 2) .. (h `div` 2) - 1]]
  colors <- mapM (\x -> rayTrace x (w, h) (1.0, 1.0) world samples) 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