-- ghc -O2 -fllvm -funbox-strict-fields -with-rtsopts=-N4 -threaded -o ray ray.hs

import Data.Maybe
import Control.Parallel
import Control.Parallel.Strategies

---- precise comparison of Doubles

eps = 1e-8
x !>  y = x     > y+eps
x !<  y = x+eps < y
x !>= y = x+eps > y
x !<= y = x     < y+eps
x !== y = x !<= y && x !>= y

---- color (rgb values are between 0 and 1)

data Color = Color !Double !Double !Double

instance Show Color where
  show (Color r g b) = format r ++ " " ++ format g ++ " " ++ format b
    where format = show . round . (*255) . min 1

magnify (Color x y z)  k            = Color (k*x) (k*y) (k*z)
apply   (Color x y z) (Color a b c) = Color (x*a) (y*b) (z*c)
combine (Color x y z) (Color a b c) = Color (x+a) (y+b) (z+c)
combineAll = foldr1 combine

---- 3d vector (can represent a point)

data Vector = Vector !Double !Double !Double
            deriving Show

(Vector x y z) ^* k = Vector (x*k) (y*k) (z*k)
(Vector x y z) ^/ k = Vector (x/k) (y/k) (z/k)

(Vector x y z) ^+^ (Vector a b c) = Vector (x+a) (y+b) (z+c)
(Vector x y z) ^-^ (Vector a b c) = Vector (x-a) (y-b) (z-c)
(Vector x y z) ^*^ (Vector a b c) = Vector (y*c-b*z) (z*a-c*x) (x*b-a*y)
(Vector x y z) ^.^ (Vector a b c) = x*a + y*b + z*c
p              ^|^ q              = sqrt $ (p^-^q) ^.^ (p^-^q)

neg (Vector x y z) = Vector (-x) (-y) (-z)
norm v = sqrt $ v ^.^ v
normalize v = v ^/ norm v

---- ray directed from starting point

data Ray = Ray { rayFrom :: Vector
               , rayDir  :: Vector
               } deriving Show

rayPoints :: Vector -> Vector -> Ray
rayPoints a b = Ray a (normalize $ b^-^a)

---- light source

data Light = Light { lightPos   :: Vector
                   , lightColor :: Color
                   } deriving Show

---- objects in scenes

data Object = Sphere { sphCenter  :: Vector
                     , sphRadius  :: Double
                     , sphDiffuse :: Color
                     , sphReflect :: Color
                     , sphRough   :: Double
                     } deriving Show

---- intersection of a ray and an object

data Inter = Inter { interPoint   :: Vector
                   , interDist    :: Double
                   , interOuter   :: Bool
                   , interNormal  :: Vector
                   , interDiffuse :: Color
                   , interReflect :: Color
                   , interRough   :: Double
                   } deriving Show

findInter :: Ray -> Object -> Maybe Inter
findInter (Ray st dir) (Sphere c r diffuse reflect rough)
  | d0 !>= r  = Nothing
  | t2 !<= 0  = Nothing
  | otherwise = Just $ Inter point dist outer normal diffuse reflect rough
  where
    t0     = (dir ^.^ (c^-^st)) / (dir ^.^ dir)
    p0     = st ^+^ (dir^*t0)
    d0     = p0 ^|^ c
    delta  = sqrt $ r^2 - d0^2
    t1     = t0 - delta
    t2     = t0 + delta
    outer  = t1 !> 0
    time   = if outer then t1 else t2
    point  = st ^+^ (dir^*time)
    dist   = point ^|^ st
    normal = normalize $ point ^-^ c

test 0 r = return ()
test n r = do
  let (Just x) = obstacle r
  print $ interPoint x
  test (n-1) (Ray (interPoint x) (rayDir r))

---- light effects

ambiental :: Color
ambiental = Color 0.06 0.06 0.06

diffusion :: Light -> Inter -> Color
diffusion l i = lightColor l `apply` interDiffuse i `magnify` cosA `magnify` 0.5
  where dirL = normalize $ lightPos l ^-^ interPoint i
        cosA = dirL ^.^ interNormal i

reflection :: Ray -> Light -> Inter -> Color
reflection r l i = lightColor l `apply` interReflect i `magnify` powCosA
  where normal  = interNormal i
        dirR    = rayDir r
        dirL    = normalize $ lightPos l ^-^ interPoint i
        cosA    = normal ^* (normal ^.^ dirL) ^* 2 ^-^ dirL ^.^ neg dirR
        powCosA = cosA `max` 0 ** interRough i

refraction :: Int -> Ray -> Inter -> Color
refraction iter r i
  = traceRay (iter-1) rayO `magnify` 0.85 `combine`
    traceRay (iter-1) rayI `magnify` 0.85
  where
    pt       = interPoint i
    normal   = interNormal i
    dirR     = rayDir r
    outer    = interOuter i
    n1       = if outer then 1.0 else 1.5
    n2       = if outer then 1.5 else 1.0
    normal'  = if outer then normal else neg normal
    alpha    = acos $ neg dirR ^.^ normal'
    arm      = normalize $ normal' ^* cos alpha ^+^ dirR
    sinB     = sin alpha * (n1/n2)
    validB   = sinB > -1 && sinB < 1
    beta     = if validB then asin sinB else alpha
    normal'' = if validB then normal' else neg normal'

    dirO     = normalize $ normal ^* (normal ^.^ neg dirR) ^* 2 ^+^ dirR
    rayO     = Ray pt dirO
    dirI     = normalize $ arm ^* tan beta ^-^ normal''
    rayI     = Ray pt dirI

dropSingle :: Ray -> Light -> Inter -> Color
dropSingle r l i | isNothing liM               = Color 0 0 0
                 | interDist li !< (posL^|^pt) = Color 0 0 0
                 | otherwise                   = result
  where posL   = lightPos l
        pt     = interPoint i
        liM    = obstacle $ rayPoints posL pt
        li     = fromJust liM
        result = diffusion l i `combine` reflection r l i

dropLight :: Int -> Ray -> Inter -> Color
dropLight iter r i = combineAll [dropSingle r l i | l <- lights]
                     `combine` refraction iter r i
                     `combine` ambiental

---- rendering

obstacle :: Ray -> Maybe Inter
obstacle ray = foldl (\a -> closer a . findInter ray) Nothing objs
  where closer Nothing  b        = b
        closer a        Nothing  = a
        closer (Just x) (Just y) | interDist x < interDist y = Just x
                                 | otherwise                 = Just y

traceRay :: Int -> Ray -> Color
traceRay 0    ray                  = Color 0 0 0
traceRay iter ray | isNothing obst = Color 0 0 0
                  | otherwise      = dropLight iter ray i
  where obst = obstacle ray
        i    = fromJust obst

getPixel :: Int -> Int -> Int -> Int -> Color
getPixel w h x y = traceRay 7 $ rayPoints eye pt
  where w' = fromIntegral w
        h' = fromIntegral h
        x' = fromIntegral x
        y' = fromIntegral y
        pt = corner ^+^ (xAxis ^* (x' / (w'-1) * viewW))
                    ^-^ (yAxis ^* (y' / (h'-1) * viewH))

render :: Int -> Int -> [Color]
render w h = runEval $ parBuffer 10000 rseq im
  where im = [getPixel w h x y | y <- [0..h-1], x <- [0..w-1]]

---- writing to file

formatPGM :: Int -> Int -> String
formatPGM w h = format $ render w h
  where format ps = "P3\n" ++ show w ++ " " ++ show h ++ "\n" ++ "255\n" ++
                   concatMap (\c -> show c ++ "\n") ps

writePGM :: Int -> Int -> IO ()
writePGM w h = writeFile file $ formatPGM w h

main = writePGM 700 700

---- constants

file    = "scene.ppm"

viewH   = 20
viewW   = 20
eye     = Vector (-20) 10 0
view    = Vector  0 0  0
viewUp  = Vector  0 0 10
viewDir = normalize $ view ^-^ eye
xAxis   = normalize $ viewDir ^*^ yAxis
yAxis   = normalize $ viewUp ^-^ (viewDir ^* (viewDir^.^viewUp))
zAxis   = normalize $ neg viewDir
corner  = view ^-^ (xAxis ^* (viewW/2)) ^+^ (yAxis ^* (viewH/2))

lights  = [ Light (Vector 10   5    5 ) (Color   0 0.0 0.9)
          , Light (Vector  2   5  (-1)) (Color 0.7 0.4 0.4)
          , Light (Vector 10 (-5)   5 ) (Color   0 0.5   0)
          ]
objs    = [ Sphere (Vector (-2) (-4)   4 ) 5 (Color 1 1 1) (Color 0.5 0.5 0.5) 10
          , Sphere (Vector (-2)   4  (-5)) 5 (Color 1 1 1) (Color 0.5 0.5 0.5) 10
          ] ++
          [ Sphere (Vector (-8) (i'/9*25-12.5) (j'/9*25-12.5)) 0.5
                   (Color 1 1 1) (Color 0.5 0.5 0.5) 4
          | i <- [0..9], j <- [0..9]
          , let i' = fromIntegral i, let j' = fromIntegral j ]