{- ORMOLU_DISABLE -}
-- Implicit CAD. Copyright (C) 2011, Christopher Olah (chris@colah.ca)
-- Copyright 2016, Julia Longtin (julial@turinglace.com)
-- Released under the GNU AGPLV3+, see LICENSE

module Graphics.Implicit.Export.RayTrace( Color(Color), average, Camera(Camera), Light(Light), Scene(Scene), traceRay, cameraRay) where

import Prelude(Show, RealFrac, Maybe(Just, Nothing), Bool(False, True), (-), (.), ($), (*), (/), min, fromInteger, max, round, fromIntegral, unzip, fmap, length, sum, maximum, minimum, (>), (+), (<), (==), pred, flip, not, abs, floor, toRational, otherwise, pure)

-- Our number system, and the definition of a 3D object.
import Graphics.Implicit.Definitions (, Fastℕ, ℝ2, ℝ3, Obj3)

import Codec.Picture (Pixel8)

import Control.Monad (guard)

import Control.Arrow ((***))

import Linear
    ( V3(V3), cross, Metric(dot, norm), V2(V2), normalize, (*^) )

default (Fastℕ, )

-- Definitions

data Camera = Camera ℝ3 ℝ3 ℝ3 
    deriving Int -> Camera -> ShowS
[Camera] -> ShowS
Camera -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Camera] -> ShowS
$cshowList :: [Camera] -> ShowS
show :: Camera -> String
$cshow :: Camera -> String
showsPrec :: Int -> Camera -> ShowS
$cshowsPrec :: Int -> Camera -> ShowS
Show

-- | A ray. A point, and a normal pointing in the direction the ray is going.
data Ray    = Ray ℝ3 ℝ3
    deriving Int -> Ray -> ShowS
[Ray] -> ShowS
Ray -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Ray] -> ShowS
$cshowList :: [Ray] -> ShowS
show :: Ray -> String
$cshow :: Ray -> String
showsPrec :: Int -> Ray -> ShowS
$cshowsPrec :: Int -> Ray -> ShowS
Show

data Scene  = Scene Obj3 Color [Light] Color

-- | A light source. source point, and intensity.
data Light  = Light ℝ3 
    deriving Int -> Light -> ShowS
[Light] -> ShowS
Light -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Light] -> ShowS
$cshowList :: [Light] -> ShowS
show :: Light -> String
$cshow :: Light -> String
showsPrec :: Int -> Light -> ShowS
$cshowsPrec :: Int -> Light -> ShowS
Show

-- | A colour. Red Green Blue and Alpha components.
data Color  = Color Pixel8 Pixel8 Pixel8 Pixel8

-- Math

-- | The distance traveled by a line segment from the first point to the second point.
vectorDistance :: ℝ3 -> ℝ3 -> 
vectorDistance :: ℝ3 -> ℝ3 -> ℝ
vectorDistance ℝ3
a ℝ3
b = forall (f :: * -> *) a. (Metric f, Floating a) => f a -> a
norm (ℝ3
bforall a. Num a => a -> a -> a
-ℝ3
a)

-- | Multiply a colour by an intensity.
colorMult :: Pixel8 -> Color -> Color
Pixel8
s colorMult :: Pixel8 -> Color -> Color
`colorMult` (Color Pixel8
a Pixel8
b Pixel8
c Pixel8
d) = Pixel8 -> Pixel8 -> Pixel8 -> Pixel8 -> Color
Color (Pixel8
s Pixel8 -> Pixel8 -> Pixel8
`mult` Pixel8
a) (Pixel8
s Pixel8 -> Pixel8 -> Pixel8
`mult` Pixel8
b) (Pixel8
s Pixel8 -> Pixel8 -> Pixel8
`mult` Pixel8
c) Pixel8
d
    where
        bound :: RealFrac a => a -> a
        bound :: forall a. RealFrac a => a -> a
bound = forall a. Ord a => a -> a -> a
max a
0 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => a -> a -> a
min a
255
        mult :: Pixel8 -> Pixel8 -> Pixel8
        mult :: Pixel8 -> Pixel8 -> Pixel8
mult Pixel8
x Pixel8
y = forall a b. (RealFrac a, Integral b) => a -> b
round forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. RealFrac a => a -> a
bound forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Real a => a -> Rational
toRational forall a b. (a -> b) -> a -> b
$ Pixel8
x forall a. Num a => a -> a -> a
* Pixel8
y

-- | Average a set of colours.
average :: [Color] -> Color
average :: [Color] -> Color
average [Color]
l =
    let
        (([ℝ]
rs, [ℝ]
gs), ([ℝ]
bs, [ℝ]
as)) = (forall a b. [(a, b)] -> ([a], [b])
unzip forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** forall a b. [(a, b)] -> ([a], [b])
unzip) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
            (\(Color Pixel8
r Pixel8
g Pixel8
b Pixel8
a) -> ((forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
r, forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
g), (forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
b, forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
a)))
            [Color]
l :: (([], []), ([], []))
        n :: 
        n :: ℝ
n = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
length [Color]
l
        (r', g', b', a') = (forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ℝ]
rsforall a. Fractional a => a -> a -> a
/n, forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ℝ]
gsforall a. Fractional a => a -> a -> a
/n, forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ℝ]
bsforall a. Fractional a => a -> a -> a
/n, forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ℝ]
asforall a. Fractional a => a -> a -> a
/n)
    in Pixel8 -> Pixel8 -> Pixel8 -> Pixel8 -> Color
Color
        (forall a. Num a => Integer -> a
fromInteger forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ r') (forall a. Num a => Integer -> a
fromInteger forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ g') (forall a. Num a => Integer -> a
fromInteger forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ b') (forall a. Num a => Integer -> a
fromInteger forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ a')

-- Ray Utilities

cameraRay :: Camera -> ℝ2 -> Ray
cameraRay :: Camera -> ℝ2 -> Ray
cameraRay (Camera ℝ3
p ℝ3
vx ℝ3
vy f) (V2 x y) =
    let
        v :: ℝ3
v  = ℝ3
vx forall a. Num a => V3 a -> V3 a -> V3 a
`cross` ℝ3
vy
        p' :: ℝ3
p' = ℝ3
p forall a. Num a => a -> a -> a
+ fforall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ℝ3
v forall a. Num a => a -> a -> a
+ xforall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ℝ3
vx forall a. Num a => a -> a -> a
+ yforall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ℝ3
vy
        n :: ℝ3
n  = forall a (f :: * -> *).
(Floating a, Metric f, Epsilon a) =>
f a -> f a
normalize (ℝ3
p' forall a. Num a => a -> a -> a
- ℝ3
p)
    in
        ℝ3 -> ℝ3 -> Ray
Ray ℝ3
p' ℝ3
n

-- | Create a ray from two points.
rayFromTo :: ℝ3 -> ℝ3 -> Ray
rayFromTo :: ℝ3 -> ℝ3 -> Ray
rayFromTo ℝ3
p1 ℝ3
p2 = ℝ3 -> ℝ3 -> Ray
Ray ℝ3
p1 (forall a (f :: * -> *).
(Floating a, Metric f, Epsilon a) =>
f a -> f a
normalize forall a b. (a -> b) -> a -> b
$ ℝ3
p2 forall a. Num a => a -> a -> a
- ℝ3
p1)

rayBounds :: Ray -> (ℝ3, ℝ3) -> ℝ2
rayBounds :: Ray -> (ℝ3, ℝ3) -> ℝ2
rayBounds Ray
ray (ℝ3, ℝ3)
box =
    let
        Ray (V3 cPx cPy cPz) (V3 cVx cVy cVz) = Ray
ray
        (V3 x1 y1 z1, V3 x2 y2 z2) = (ℝ3, ℝ3)
box
        xbounds :: [ℝ]
xbounds = [(x1 forall a. Num a => a -> a -> a
- cPx)forall a. Fractional a => a -> a -> a
/cVx, (x2forall a. Num a => a -> a -> a
-cPx)forall a. Fractional a => a -> a -> a
/cVx]
        ybounds :: [ℝ]
ybounds = [(y1forall a. Num a => a -> a -> a
-cPy)forall a. Fractional a => a -> a -> a
/cVy, (y2forall a. Num a => a -> a -> a
-cPy)forall a. Fractional a => a -> a -> a
/cVy]
        zbounds :: [ℝ]
zbounds = [(z1forall a. Num a => a -> a -> a
-cPz)forall a. Fractional a => a -> a -> a
/cVz, (z2forall a. Num a => a -> a -> a
-cPz)forall a. Fractional a => a -> a -> a
/cVz]
        lower :: ℝ
lower   = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [ℝ]
xbounds, forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [ℝ]
ybounds, forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [ℝ]
zbounds]
        upper :: ℝ
upper   = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [ℝ]
xbounds, forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [ℝ]
ybounds, forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [ℝ]
zbounds]
    in
        forall a. a -> a -> V2 a
V2 lower upper

-- Intersection
-- FIXME: magic numbers.
intersection :: Ray -> ((,), ) ->  -> Obj3 -> Maybe ℝ3
intersection :: Ray -> ((ℝ, ℝ), ℝ) -> ℝ -> (ℝ3 -> ℝ) -> Maybe ℝ3
intersection r :: Ray
r@(Ray ℝ3
p ℝ3
v) ((a, aval),b) res ℝ3 -> ℝ
obj =
    let
        step :: ℝ
step | avalforall a. Fractional a => a -> a -> a
/4 forall a. Ord a => a -> a -> Bool
> res = res
             | avalforall a. Fractional a => a -> a -> a
/2 forall a. Ord a => a -> a -> Bool
> res = resforall a. Fractional a => a -> a -> a
/2
             | Bool
otherwise = resforall a. Fractional a => a -> a -> a
/10
        a' :: ℝ
a'  = a forall a. Num a => a -> a -> a
+ step
        a'val :: ℝ
a'val = ℝ3 -> ℝ
obj (ℝ3
p forall a. Num a => a -> a -> a
+ a'forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ℝ3
v)
    in if a'val forall a. Ord a => a -> a -> Bool
< 0
    then
        let a'' :: ℝ
a'' = ℝ2 -> (ℝ -> ℝ) -> ℝ
refine (forall a. a -> a -> V2 a
V2 a a') (\s -> ℝ3 -> ℝ
obj (ℝ3
p forall a. Num a => a -> a -> a
+ sforall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ℝ3
v))
        in forall a. a -> Maybe a
Just (ℝ3
p forall a. Num a => a -> a -> a
+ a''forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ℝ3
v)
    else if a' forall a. Ord a => a -> a -> Bool
< b
    then Ray -> ((ℝ, ℝ), ℝ) -> ℝ -> (ℝ3 -> ℝ) -> Maybe ℝ3
intersection Ray
r ((a',a'val), b) res ℝ3 -> ℝ
obj
    else forall a. Maybe a
Nothing

refine :: ℝ2 -> ( -> ) -> 
refine :: ℝ2 -> (ℝ -> ℝ) -> ℝ
refine (V2 a b) ℝ -> ℝ
obj =
    let
        (aval, bval) = (ℝ -> ℝ
obj a, ℝ -> ℝ
obj b)
    in if bval forall a. Ord a => a -> a -> Bool
< aval
    then Fastℕ -> ℝ2 -> ℝ2 -> (ℝ -> ℝ) -> ℝ
refine' Fastℕ
10 (forall a. a -> a -> V2 a
V2 a b) (forall a. a -> a -> V2 a
V2 aval bval) ℝ -> ℝ
obj
    else Fastℕ -> ℝ2 -> ℝ2 -> (ℝ -> ℝ) -> ℝ
refine' Fastℕ
10 (forall a. a -> a -> V2 a
V2 b a) (forall a. a -> a -> V2 a
V2 aval bval) ℝ -> ℝ
obj

refine' :: Fastℕ -> ℝ2 -> ℝ2 -> ( -> ) -> 
refine' :: Fastℕ -> ℝ2 -> ℝ2 -> (ℝ -> ℝ) -> ℝ
refine' Fastℕ
0 (V2 a _) ℝ2
_ ℝ -> ℝ
_ = a
refine' Fastℕ
n (V2 a b) (V2 aval bval) ℝ -> ℝ
obj =
    let
        mid :: ℝ
mid = (aforall a. Num a => a -> a -> a
+b)forall a. Fractional a => a -> a -> a
/2
        midval :: ℝ
midval = ℝ -> ℝ
obj mid
    in
        if midval forall a. Eq a => a -> a -> Bool
== 0
        then mid
        else if midval forall a. Ord a => a -> a -> Bool
< 0
        then Fastℕ -> ℝ2 -> ℝ2 -> (ℝ -> ℝ) -> ℝ
refine' (forall a. Enum a => a -> a
pred Fastℕ
n) (forall a. a -> a -> V2 a
V2 a mid) (forall a. a -> a -> V2 a
V2 aval midval) ℝ -> ℝ
obj
        else Fastℕ -> ℝ2 -> ℝ2 -> (ℝ -> ℝ) -> ℝ
refine' (forall a. Enum a => a -> a
pred Fastℕ
n) (forall a. a -> a -> V2 a
V2 mid b) (forall a. a -> a -> V2 a
V2 midval bval) ℝ -> ℝ
obj

intersects :: Ray -> ((, ), ) ->  -> Obj3 -> Bool
intersects :: Ray -> ((ℝ, ℝ), ℝ) -> ℝ -> (ℝ3 -> ℝ) -> Bool
intersects Ray
a ((ℝ, ℝ), ℝ)
b c ℝ3 -> ℝ
d = case Ray -> ((ℝ, ℝ), ℝ) -> ℝ -> (ℝ3 -> ℝ) -> Maybe ℝ3
intersection Ray
a ((ℝ, ℝ), ℝ)
b c ℝ3 -> ℝ
d of
    Maybe ℝ3
Nothing -> Bool
False
    Just ℝ3
_  -> Bool
True

-- Trace
-- FIXME: magic numbers.
traceRay :: Ray ->  -> (ℝ3, ℝ3) -> Scene -> Color
traceRay :: Ray -> ℝ -> (ℝ3, ℝ3) -> Scene -> Color
traceRay ray :: Ray
ray@(Ray ℝ3
cameraP ℝ3
cameraV) step (ℝ3, ℝ3)
box (Scene ℝ3 -> ℝ
obj Color
objColor [Light]
lights Color
defaultColor) =
    let
        (V2 a b) = Ray -> (ℝ3, ℝ3) -> ℝ2
rayBounds Ray
ray (ℝ3, ℝ3)
box
    in case Ray -> ((ℝ, ℝ), ℝ) -> ℝ -> (ℝ3 -> ℝ) -> Maybe ℝ3
intersection Ray
ray ((a, ℝ3 -> ℝ
obj (ℝ3
cameraP forall a. Num a => a -> a -> a
+ aforall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ℝ3
cameraV)), b) step ℝ3 -> ℝ
obj of
        Just ℝ3
p  -> forall a b c. (a -> b -> c) -> b -> a -> c
flip Pixel8 -> Color -> Color
colorMult Color
objColor forall a b. (a -> b) -> a -> b
$ forall a b. (RealFrac a, Integral b) => a -> b
floor (forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall a b. (a -> b) -> a -> b
$ 0.2 forall a. a -> [a] -> [a]
: do
            Light ℝ3
lightPos lightIntensity <- [Light]
lights
            let
                ray' :: Ray
ray'@(Ray ℝ3
_ ℝ3
v) = ℝ3 -> ℝ3 -> Ray
rayFromTo ℝ3
p ℝ3
lightPos
                v' :: ℝ3
v' = forall a (f :: * -> *).
(Floating a, Metric f, Epsilon a) =>
f a -> f a
normalize ℝ3
v
            forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ Ray -> ((ℝ, ℝ), ℝ) -> ℝ -> (ℝ3 -> ℝ) -> Bool
intersects Ray
ray' ((0, ℝ3 -> ℝ
obj ℝ3
p),20) step ℝ3 -> ℝ
obj
            let
                pval :: ℝ
pval = ℝ3 -> ℝ
obj ℝ3
p
                dirDeriv :: ℝ3 -> 
                dirDeriv :: ℝ3 -> ℝ
dirDeriv ℝ3
v'' = (ℝ3 -> ℝ
obj (ℝ3
p forall a. Num a => a -> a -> a
+ stepforall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ℝ3
v'') forall a. Num a => a -> a -> a
- pval)forall a. Fractional a => a -> a -> a
/step
                deriv :: ℝ3
deriv = forall a. a -> a -> a -> V3 a
V3 (ℝ3 -> ℝ
dirDeriv (forall a. a -> a -> a -> V3 a
V3 1 0 0)) (ℝ3 -> ℝ
dirDeriv (forall a. a -> a -> a -> V3 a
V3 0 1 0)) (ℝ3 -> ℝ
dirDeriv (forall a. a -> a -> a -> V3 a
V3 0 0 1))
                normal :: ℝ3
normal = forall a (f :: * -> *).
(Floating a, Metric f, Epsilon a) =>
f a -> f a
normalize ℝ3
deriv
                unitV :: ℝ3
unitV = forall a (f :: * -> *).
(Floating a, Metric f, Epsilon a) =>
f a -> f a
normalize ℝ3
v'
                -- proj :: InnerSpace v => v -> v -> v
                proj :: f a -> f a -> f a
proj f a
a' f a
b' = (f a
a' forall (f :: * -> *) a. (Metric f, Num a) => f a -> f a -> a
`dot` f a
b')forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^f a
b'
                dist :: ℝ
dist  = ℝ3 -> ℝ3 -> ℝ
vectorDistance ℝ3
p ℝ3
lightPos
                illumination :: ℝ
illumination = forall a. Ord a => a -> a -> a
max 0 (ℝ3
normal forall (f :: * -> *) a. (Metric f, Num a) => f a -> f a -> a
`dot` ℝ3
unitV) forall a. Num a => a -> a -> a
* lightIntensity forall a. Num a => a -> a -> a
* (25 forall a. Fractional a => a -> a -> a
/dist)
                rV :: ℝ3
rV =
                    let
                        normalComponent :: ℝ3
normalComponent = forall {f :: * -> *} {a}. (Num a, Metric f) => f a -> f a -> f a
proj ℝ3
v' ℝ3
normal
                        parComponent :: ℝ3
parComponent    = ℝ3
v' forall a. Num a => a -> a -> a
- ℝ3
normalComponent
                    in
                        ℝ3
normalComponent forall a. Num a => a -> a -> a
- ℝ3
parComponent
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ illuminationforall a. Num a => a -> a -> a
*(3 forall a. Num a => a -> a -> a
+ 0.3forall a. Num a => a -> a -> a
*forall a. Num a => a -> a
abs(ℝ3
rV forall (f :: * -> *) a. (Metric f, Num a) => f a -> f a -> a
`dot` ℝ3
cameraV)forall a. Num a => a -> a -> a
*forall a. Num a => a -> a
abs(ℝ3
rV forall (f :: * -> *) a. (Metric f, Num a) => f a -> f a -> a
`dot` ℝ3
cameraV))
            )
        Maybe ℝ3
Nothing   -> Color
defaultColor