module Main where

import Graphics.Imlib
import Data.Array.IArray

type Point = (Double, Double, Double)
type Vector = (Double, Double, Double)
type Color = (Double, Double, Double)

data Ray = Ray Point Vector
data Light = Light Point Color
data Surface = Surface {color :: Color
                       ,reflectivity :: Double
                       ,smoothness :: Double
                       ,opacity :: Double
                       ,refractiveIx :: Double}
data Object = Plane {point :: Point
                    ,normal :: Vector
                    ,surface :: Surface}
            | Sphere {point :: Point
                     ,radius :: Double
                     ,surface :: Surface}
data Hit = NoHit
         | Hit Double Point Vector Surface

type View = Array (Int, Int) Ray
type Scene = ([Light], [Object])

black = (0, 0, 0) :: Color
white = (1, 1, 1) :: Color
red = (1, 0, 0) :: Color
green = (0, 1, 0) :: Color
blue = (0, 0, 1) :: Color
yellow = (1, 1, 0) :: Color
blueGreen = (0, 1, 1) :: Color
violet = (1, 0, 1) :: Color

instance Num a => Num (a, a, a) where
    (x1, y1, z1) + (x2, y2, z2) = (x1 + x2, y1 + y2, z1 + z2)
    (x1, y1, z1) - (x2, y2, z2) = (x1 - x2, y1 - y2, z1 - z2)
    (x1, y1, z1) * (x2, y2, z2) = (y1 * z2 - z1 * y2
                                  ,z1 * x2 - x1 * z2
                                  ,x1 * y2 - x2 * y1)
    negate (x, y, z) = (-x, -y, -z)

k `scl` (x, y, z) = (k * x, k * y, k * z)

(x1, y1, z1) `dot` (x2, y2, z2) = x1 * x2 + y1 * y2 + z1 * z2

mag (x, y, z) = sqrt (x^2 + y^2 + z^2)

a `dist` b = mag (b - a)

sq v = v `dot` v

a `dotScl` b = (a `dot` b) / (mag a * mag b)

v `sclTo` k = (k / mag v) `scl` v

rotateBy :: Vector -> (Double, Double) -> Vector
v@(x, y, z) `rotateBy` (dtheta, dphi) =
    let r = sqrt (x^2 + y^2)
        theta = atan (y / x)
        phi = atan (z / r)
    in (r * cos (theta + dtheta)
       ,r * sin (theta + dtheta)
       ,mag v * sin (phi + dphi))

normalAt :: Object -> Point -> Vector
(Sphere c _ _) `normalAt` p = p - c
(Plane _ n _) `normalAt` _ = n

newRay :: Point -> Vector -> Ray
newRay s d = Ray s (d `sclTo` 1)

rayTo :: Point -> Point -> Ray
a `rayTo` b = Ray a ((b - a) `sclTo` 1)

ptAt :: Ray -> Double -> Point
Ray s d `ptAt` l = s + (d `sclTo` l)

reflect :: Vector -> Vector -> Vector
reflect v n = v - (2 * (v `dot` n) / sq n) `scl` n

refract :: Vector -> Vector -> Double -> Vector
refract v n i = negate n `sclTo` i + v

hitDist :: Ray -> Object -> Double
(Ray s d) `hitDist` (Sphere p r _) =
    let v = s - p
--      a = sq d --because d has magnitude 1
        b = 2 `scl` v `dot` d
        c = v `dot` v - r^2
        dis = sqrt (b^2 - 4 {-* a-} * c)
    in min ((-b + dis) / (2 {-* a-})) ((-b - dis) / (2 {-* a-}))
(Ray s d) `hitDist` (Plane p n _) = (n `dot` (p - s)) / (n `dot` d)

intersect :: Ray -> Object -> Hit
ray `intersect` obj =
    let t = ray `hitDist` obj
        p = ray `ptAt` t
    in if isNaN t || isInfinite t || t <= 1e-11
       then NoHit
       else Hit t p (obj `normalAt` p) (surface obj)

intersects :: Ray -> [Object] -> Hit
intersects ray objs = foldl1 minHit $ map (intersect ray) objs
    where minHit a@(Hit x _ _ _) b@(Hit y _ _ _) = if x < y
                                                   then a
                                                   else b
          minHit NoHit hit = hit
          minHit hit NoHit = hit

lighted :: Hit -> Scene -> Color
lighted _ ([], _) = black
lighted hit@(Hit _ p n _) (((Light q color) : rest), objs) =
    let d = q - p
    in lighted hit (rest, objs) +
           (case intersects (newRay p d) objs of
              NoHit -> (d `dotScl` n / sq d) `scl` color
              Hit _ _ _ _ -> black)

reflected :: Ray -> Hit -> Scene -> Color
reflected (Ray _ v) (Hit _ p n sfc) scene = 
    reflectivity sfc `scl`
    rayTrace scene (newRay p (reflect v n))

refracted :: Ray -> Hit -> Scene -> Color
refracted (Ray _ v) (Hit _ p n sfc) scene =
    opacity sfc `scl`
    rayTrace scene (newRay p (refract v n (refractiveIx sfc)))

onSurface :: Color -> Surface -> Color
(r, g, b) `onSurface` (Surface (x, y, z) _ _ _ _) = (x * r, y * g, z * b)

rayTrace :: Scene -> Ray -> Color
rayTrace scene ray =
    case intersects ray (snd scene) of
      NoHit -> black
      hit@(Hit _ _ _ sfc) -> (lighted hit scene +
                              refracted ray hit scene +
                              reflected ray hit scene)
                             `onSurface` sfc

newView :: Ray -> (Double, Double) -> (Int, Int) -> View
newView (Ray s d) (theta, phi) (x, y) =
    let stepAngles angle num =
            let num' = fromIntegral num
            in zip [0 .. num] $ map (* angle) $
               map (/ num') [- num' / 2 .. num' / 2]
    in array ((0, 0), (x, y)) $ concat $
       map (\(j, dphi) ->
            map (\(i, dtheta) ->
                 ((i, j), newRay s (d `rotateBy` (dtheta, dphi)))) $
            stepAngles theta x) $ stepAngles phi y

viewTrace :: Scene -> View -> Array (Int, Int) Color
viewTrace scene = amap (rayTrace scene)

--and now, actually using the raytracer:

defaultSurface = Surface {color = white
                         ,reflectivity = 0.5
                         ,smoothness = 0.5
                         ,opacity = 0.5
                         ,refractiveIx = 0.5}

scene = ([Light (0, 0, 5) white]
        ,[Sphere (1, 1, 3) 1 defaultSurface
         ,Plane (0, 0, 0) (0, 0, 1) defaultSurface])

view = newView ((-5, -5, 5) `rayTo` (0, 0, 0)) (pi / 2, pi / 2) (512, 512)

main :: IO ()
main = (createImageUsingArray $ amap toImlibColor $ viewTrace scene view)
       >>= contextSetImage
       >> saveImage "out.png"
    where toImlibColor :: Color -> ImlibColor
          toImlibColor (r, g, b) = ImlibColor 0 (s255 r) (s255 g) (s255 b)
              where s255 s = round (max 0 (min 1 s) * 255)

instance Show Hit where
    show NoHit = "NoHit"
    show (Hit a b c _) = show a ++ show b ++ show c