Ray.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE FlexibleInstances #-}
module Graphics.Ray 
  ( -- * Camera
    CameraSettings(..), defaultCameraSettings
    -- * Ray Tracing
  , ToRandom(toRandom), raytrace
    -- * Image IO
  , readImage, writeImage, writeImageSqrt
    -- * Re-exports
  , module Graphics.Ray.Core
  , module Graphics.Ray.Geometry
  , module Graphics.Ray.Material
  , module Graphics.Ray.Texture
  , module Graphics.Ray.Noise
  ) where

import Graphics.Ray.Core
import Graphics.Ray.Geometry
import Graphics.Ray.Material
import Graphics.Ray.Texture
import Graphics.Ray.Noise

import Linear (V2(V2), V3(V3), (*^), (^*), normalize, cross, (^/), zero, dot)
import System.Random (StdGen, random, splitGen)
import Data.Massiv.Array (B, D, S, U, Ix2((:.)), (!))
import qualified Data.Massiv.Array as A
import qualified Data.Massiv.Array.IO as I
import Graphics.Pixel.ColorSpace (SRGB, Linearity(Linear, NonLinear))
import qualified Graphics.Pixel.ColorSpace as C
import Control.Monad.State (State, state, evalState)
import Control.Monad (replicateM)
import Data.Functor.Identity (Identity, runIdentity)
import Data.Maybe (listToMaybe)

data CameraSettings = CameraSettings
  { cs_center :: Point3 
    -- ^ Camera position 
  , cs_lookAt :: Point3 
    -- ^ Point for the camera to look at 
  , cs_up :: Vec3 
    -- ^ Camera \"up\" vector 
  , cs_vfov :: Double 
    -- ^ Vertical field of view (in radians) 
  , cs_aspectRatio :: Double 
    -- ^ Width-to-height ratio of image 
  , cs_imageWidth :: Int 
    -- ^ Image width in pixels 
  , cs_samplesPerPixel :: Int 
    -- ^ Number of top-level rays created per pixel 
  , cs_maxRecursionDepth :: Int 
    -- ^ Number of times a ray can reflect before recursion stops 
  , cs_background :: Ray -> Color 
    -- ^ Background color (which can depend on direction) 
  , cs_defocusAngle :: Double 
    -- ^ If this is positive, the image will be somewhat blurry in the foreground and background, 
    -- with only a single plane in focus (like an image produced by a real camera)
  , cs_focusDist :: Double 
    -- ^ Distance from the camera to the plane of focus (only matters if defocus angle is nonzero) 
  , cs_redirectTargets :: [(Double, Point3, Vec3, Vec3)] 
    -- ^ List of parallelograms to send scattered rays toward (if the material allows it) along with 
    -- probabilities, which should add up to less than 1. Sending more rays toward lights in 
    -- scenes with small light sources can lead to faster convergence (i.e. less noise).
  }

-- | By default, the camera is positioned at the origin looking in the negative z direction, with the positive y direction being upward
-- (and the positive x direction being rightward). The remaining attributes are as follows:
-- 
-- @
-- cs_vfov = pi / 2
-- cs_aspectRatio = 1.0
-- cs_imageWidth = 100
-- cs_samplesPerPixel = 10
-- cs_maxRecursionDepth = 10
-- cs_background = const (V3 1 1 1)
-- cs_defocusAngle = 0.0
-- cs_focusDist = 10.0
-- cs_redirectTargets = []
-- @
defaultCameraSettings :: CameraSettings
defaultCameraSettings = CameraSettings
  { cs_center = V3 0 0 0
  , cs_lookAt = V3 0 0 (-1)
  , cs_up = V3 0 1 0
  , cs_vfov = pi / 2
  , cs_aspectRatio = 1.0
  , cs_imageWidth = 100
  , cs_samplesPerPixel = 10
  , cs_maxRecursionDepth = 10
  , cs_background = const (V3 1 1 1)
  , cs_defocusAngle = 0.0
  , cs_focusDist = 10.0
  , cs_redirectTargets = []
  }

class ToRandom m where
  toRandom :: m a -> State StdGen a

instance ToRandom Identity where
  toRandom :: Identity a -> State StdGen a
  toRandom = pure . runIdentity

instance ToRandom (State StdGen) where
  toRandom :: State StdGen a -> State StdGen a
  toRandom = id

-- [private]
data RedirectTarget = RedirectTarget
  { rt_origin :: Point3
  , rt_U :: Vec3
  , rt_V :: Vec3
  , rt_cross :: Vec3
  , rt_hit :: Ray -> Maybe Double
  }

-- | Produce an image from the given camera settings, world, and seed.
raytrace :: ToRandom m => CameraSettings -> Geometry m Material -> StdGen -> A.Matrix D Color
raytrace (CameraSettings {..}) (Geometry _ hitWorld) seed = let
  imageHeight = round (fromIntegral cs_imageWidth / cs_aspectRatio)
  viewportHeight = cs_focusDist * tan (cs_vfov / 2) * 2
  viewportWidth = viewportHeight * fromIntegral cs_imageWidth / fromIntegral imageHeight

  w = normalize (cs_center - cs_lookAt)
  u = normalize (cross cs_up w) 
  v = cross w u

  across = viewportWidth *^ u
  down = -(viewportHeight *^ v)

  topLeft = cs_center - w ^* cs_focusDist - across ^/ 2 - down ^/ 2
  pixelU = across ^/ fromIntegral cs_imageWidth
  pixelV = down ^/ fromIntegral imageHeight

  redirectTargets = flip map cs_redirectTargets $ \(_, q, s0, s1) -> RedirectTarget
    { rt_origin = q
    , rt_U = s0
    , rt_V = s1
    , rt_cross = cross s0 s1
    , rt_hit = 
        let Geometry _ hit = parallelogram q s0 s1 in
        \ray -> fmap (hr_t . fst) (runIdentity (hit undefined ray (0, infinity)))
    }

  probs = map (\(p, _, _, _) -> p) cs_redirectTargets
  remProb = 1 - sum probs
  thresholds = zip redirectTargets (scanl1 (+) probs)
  getTarget r = fmap fst (listToMaybe (dropWhile ((r >=) . snd) thresholds))

  defocusRadius = cs_focusDist * tan (cs_defocusAngle / 2)
  defocusDiskU = u ^* defocusRadius
  defocusDiskV = v ^* defocusRadius
  
  sampleDefocusDisk :: State StdGen Point3
  sampleDefocusDisk = do
    V2 x y <- randomInUnitDisk
    pure (cs_center + x *^ defocusDiskU + y *^ defocusDiskV)

  samplePixel :: Int -> Int -> State StdGen Point3
  samplePixel i j = do
    x <- state random
    y <- state random
    pure (topLeft + (fromIntegral i + x) *^ pixelU + (fromIntegral j + y) *^ pixelV)

  getRay :: Int -> Int -> State StdGen Ray
  getRay i j = do
    origin <- sampleDefocusDisk
    target <- samplePixel i j
    pure (Ray origin (normalize (target - origin)))

  rayColor :: Int -> Double -> Ray -> State StdGen Color
  rayColor depth time ray@(Ray _ rayDir)
    | depth <= 0 = pure zero
    | otherwise =
    toRandom (hitWorld time ray (0.0001, infinity)) >>= \case
      Nothing -> pure (cs_background ray)
      Just (hit, Material mat) -> do
        let (emitted, genRes) = mat rayDir hit
        res <- genRes
        (emitted +) <$> case res of
          Absorb -> pure zero
          Scatter attenuation dir -> (attenuation *) <$> rayColor (depth - 1) time (Ray (hr_point hit) dir)

          HemisphereF matF -> do
            choice <- getTarget <$> state random
            dir <- case choice of
              Nothing -> do
                uu <- randomUnitVector
                pure (normalize (hr_normal hit + uu))
              Just RedirectTarget {..} -> do
                (i, j) <- state random
                let lightPt = rt_origin + i *^ rt_U + j *^ rt_V
                pure (normalize (lightPt - hr_point hit))
            let pdf1 = dot dir (hr_normal hit) / pi
            if pdf1 <= 0 then pure zero else do
              let pdfs = flip map redirectTargets $ \RedirectTarget {..} ->
                    case rt_hit (Ray (hr_point hit) dir) of
                      Nothing -> 0
                      Just t -> t * t / abs (dot rt_cross dir)
              -- The pdf from which 'dir' was generated
              let pdf = remProb * pdf1 + sum (zipWith (*) probs pdfs)
              c <- rayColor (depth - 1) time (Ray (hr_point hit) dir)
              pure (matF dir * c ^* (pdf1 / pdf))

          SphereF matF -> do
            choice <- getTarget <$> state random
            dir <- case choice of
              Nothing -> randomUnitVector
              Just RedirectTarget {..} -> do
                (i, j) <- state random
                let lightPt = rt_origin + i *^ rt_U + j *^ rt_V
                pure (normalize (lightPt - hr_point hit))
            let pdf1 = 0.25 / pi
            let pdfs = flip map redirectTargets $ \RedirectTarget {..} ->
                  case rt_hit (Ray (hr_point hit) dir) of
                    Nothing -> 0
                    Just t -> t * t / abs (dot rt_cross dir)
            -- The pdf from which 'dir' was generated
            let pdf = remProb * pdf1 + sum (zipWith (*) probs pdfs) 
            c <- rayColor (depth - 1) time (Ray (hr_point hit) dir)
            pure (matF dir * c ^* (pdf1 / pdf))
  
  pixelColor :: Int -> Int -> State StdGen Color
  pixelColor i j = do
    colors <- replicateM cs_samplesPerPixel $ do
      time <- state random
      ray <- getRay i j 
      rayColor cs_maxRecursionDepth time ray
    pure (sum colors ^/ fromIntegral cs_samplesPerPixel)

  -- array of random seeds for each pixel (constructed using splitGen)
  seeds :: A.Matrix B StdGen
  (_, seeds) = A.randomArrayS seed (A.Sz (imageHeight :. cs_imageWidth)) splitGen

  in A.makeArray A.Par (A.Sz (imageHeight :. cs_imageWidth)) (\ix@(j :. i) -> evalState (pixelColor i j) (seeds ! ix))

-- | Read an image file, converting each pixel to linear RGB color space.
readImage :: FilePath -> IO (A.Matrix U Color)
readImage path = A.compute . A.map fromPixel <$> (I.readImageAuto path :: IO (A.Matrix S (C.Pixel (SRGB 'Linear) Double)))
  where
    fromPixel :: C.Pixel (SRGB 'Linear) Double -> Color
    fromPixel (C.Pixel (C.ColorSRGB r g b)) = V3 r g b

-- | Write an array of linear RGB colors to an image file.
writeImage :: (A.Source r Color) => FilePath -> A.Matrix r Color -> IO ()
writeImage path m = I.writeImageAuto path (A.map toPixel m)
  where
    toPixel :: Color -> C.Pixel (SRGB 'Linear) Double
    toPixel (V3 r g b) = C.Pixel (C.ColorSRGB r g b)

-- | Write an array to an image file, using a slightly incorrect color space conversion function.
-- This function exists for testing purposes.
writeImageSqrt :: (A.Source r Color) => FilePath -> A.Matrix r Color -> IO ()
writeImageSqrt path m = I.writeImageAuto path (A.map toPixel m)
  where
    toPixel :: Color -> C.Pixel (SRGB 'NonLinear) Double
    toPixel (V3 r g b) = C.Pixel (C.ColorSRGB (sqrt r) (sqrt g) (sqrt b))