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 Maths
import Ray
import Text.Printf

type Color = (Double, Double, Double)

type Coord = (Int, Int)

type ImageDimensions = (Int, Int)

type ViewPortDimensions = (Double, Double)

-- TODO
-- DONE add multiple samples per pixel
-- fix coordinate system
-- DONE fix random number generator

backgroundColor = V3 0.0 0.0 0.0

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)) ws depth =
  do
    r <- randV3
    let target = p ^+^ colorWithNormal (unitVector n) ^+^ r
    nextVal <- getRayColor ws (depth - 1) (Ray p (target ^-^ p))
    return (0.5 *^ 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

-- testHit needs to take into account other hits, not here
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
  -- weirdly, the bottom circle appears on top and bottom
  -- Circle (V3 0.0 (-10.05) 1.0) 5]
  let world = [Circle (V3 0.0 0.0 (-1.0)) 0.25, Circle (V3 0.0 (-10.25) (-1.0)) 10]
  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