import Data.Maybe

---- 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

scaleC (Color x y z) k = Color (k*x) (k*y) (k*z)
multC (Color x y z) (Color a b c) = Color (x*a) (y*b) (z*c)
addC (Color x y z) (Color a b c) = Color (x+a) (y+b) (z+c)
addAllC = foldr1 addC

---- 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
    time   = if outer then t1 else t2
    point  = st ^+^ (dir^*time)
    dist   = point ^|^ st
    outer  = t1 !> 0
    normal = normalize $ point ^-^ c

---- light effects

ambiental :: Color
ambiental = Color 0.06 0.06 0.06

diffusion :: Light -> Inter -> Color
diffusion l i = (lightColor l) `multC` (interDiffuse i) `scaleC` cosA
  where dirL = normalize $ (lightPos l) ^-^ (interPoint i)
        cosA = dirL ^.^ (interNormal i)

reflection :: Ray -> Light -> Inter -> Color
reflection r l i = (lightColor l) `multC` (interReflect i) `scaleC` 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)

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 = addAllC [diffusion l i, reflection r l i]

dropLight :: Ray -> Inter -> Color
dropLight r i = addAllC [dropSingle r l i | l <- lights] `addC` 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 :: Ray -> Color
traceRay ray | isNothing obst = Color 0 0 0
             | otherwise      = dropLight ray i
  where obst = obstacle ray
        i    = fromJust obst

getPixel :: Int -> Int -> Int -> Int -> Color
getPixel w h x y = traceRay $ 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 = [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 400 400

---- constants

file    = "slika.ppm"

viewH   = 20
viewW   = 20
eye     = Vector 10 0  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.4     0    0)
          , Light (Vector 10 (-5)   5 ) (Color   0  0.31 0.31)
          , Light (Vector  2   5  (-1)) (Color 0.31 0.31    0)
          ]
objs    = [ Sphere (Vector (-2) (-2)   2 ) 5 (Color 1 1 1) (Color 0.5 0.5 0.5) 10
          , Sphere (Vector (-2)   2  (-3)) 5 (Color 1 1 1) (Color 0.5 0.5 0.5) 10
          ] ++
          [ Sphere (Vector (-5) (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 ]