{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns        #-}

module Ray.Trace where

import Common.Type
import Scene.Object
import Scene.World

import Data.Array.Accelerate as A
import Data.Array.Accelerate.Data.Colour.Names
import Data.Array.Accelerate.Data.Colour.RGB as RGB
import Data.Array.Accelerate.Linear.Metric
import Data.Array.Accelerate.Linear.V3
import Data.Array.Accelerate.Linear.Vector
import Graphics.Gloss.Accelerate.Data.Point

import qualified Prelude as P


castViewRays :: Int -> Int -> Float -> Acc (Array DIM2 Direction)
castViewRays sizeX sizeY fov =
  let sizeX' = P.fromIntegral sizeX
      sizeY' = P.fromIntegral sizeY
      aspect = sizeX' / sizeY'
      fov' = constant (P.sin (fov / 2))
      fovX = fov' * aspect
      fovY = fov'
   in A.generate
        (constant (Z :. sizeY :. sizeX))
        (\ix ->
           let (x, y) = xyOfPoint $ pointOfIndex sizeX sizeY ix
            in normalize $ lift (V3 (x * fovX) ((-y) * fovY) 1))

normEstimate :: Estimator a => Scene a -> Exp Position -> Exp Direction
normEstimate scene p =
  let del = 0.0001

      e = sceneEstimator scene
      gx1 = estimate e (p - constant (V3 del 0 0))
      gx2 = estimate e (p + constant (V3 del 0 0))
      gy1 = estimate e (p - constant (V3 0 del 0))
      gy2 = estimate e (p + constant (V3 0 del 0))
      gz1 = estimate e (p - constant (V3 0 0 del))
      gz2 = estimate e (p + constant (V3 0 0 del))

      gradX = (gx2 - gx1) / (constant del)
      gradY = (gy2 - gy1) / (constant del)
      gradZ = (gz2 - gz1) / (constant del)
   in normalize $ lift (V3 gradX gradY gradZ)

traceRay :: Estimator a => Scene a -> Exp Position -> Exp Direction -> Exp (Float, Position)
traceRay scene p d =
  let test :: Exp (Int, Float, Position) -> Exp Bool
      test (unlift -> (i :: Exp Int, r :: Exp Float, _ :: Exp Position)) =
        i < constant (sceneMaxiter scene) && r > constant (sceneEps scene) && r < 5

      step :: Exp (Int, Float, Position) -> Exp (Int, Float, Position)
      step (unlift -> (i :: Exp Int, _ :: Exp Float, p :: Exp Position)) =
        let r = estimate (sceneEstimator scene) p
         in lift (i + 1, r, p + r *^ d)

      (_ :: Exp Int, r' :: Exp Float, p' :: Exp Position) = unlift $ while test step (lift (0 :: Exp Int, 1 :: Exp Float, p))
   in lift (r', p')

phong :: Light -> Exp Position -> Exp Position -> Exp Direction -> Exp Colour
phong light eye p n =
  let ambient = 1
      diffuse = constant (lightColour light)
      specularExponent = 10 :: Exp Int
      specularity = 0.5

      lightP = constant (lightPosition light)

      l = normalize (lightP - p)
      e = normalize (eye - p)
      mag = n `dot` l
      r = 2 * mag *^ n - l
   in (ambient *^ lightGrey +
       diffuse ^* (max mag 0) ^+
       (specularity * (max (e `dot` r) 0) ^ specularExponent)) ^/
      (ambient + 1 + specularity)

simpleShading :: Estimator a => Scene a -> Exp Direction -> Exp Float -> Exp Colour
simpleShading scene d noise =
  let o = constant (sceneEye scene) + noise *^ d
      (r :: Exp Float, p :: Exp Position) = unlift $ traceRay scene o d
      n = normEstimate scene p

      calc :: Light -> Exp Colour
      calc light = let l = normalize (constant (lightPosition light) - p)
                       o' = p + (constant (10 * (sceneEps scene))) *^ n
                       (r' :: Exp Float, p' :: Exp Position) = unlift $ traceRay scene o' l
--                       shadow = (r' < constant (sceneEps scene) ? (constant 0.4, constant 1))
                    in (r < constant (sceneEps scene) ?
                        ( phong light (constant (sceneEye scene)) p n
                        , constant (sceneColour scene)))
   in (P.foldr (+) (constant 0) (P.fmap calc (sceneLights scene)))/2