haskell-raytracer/app/Main.hs

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

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 (Metal (V3 0.8 0.8 0.8)),
          Circle (V3 0.0 (-10.25) (-1.0)) 10 (Lambertian (V3 0.7 0.3 0.3))
        ]
  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