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

module Ray.Trace where

import Common.Type
import Scene.Light
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 :: Scene -> Exp Position -> Exp Direction
normEstimate scene p =
  let del = 0.0001

      e = sceneEstimator scene
      gx1 = e (p - constant (V3 del 0 0))
      gx2 = e (p + constant (V3 del 0 0))
      gy1 = e (p - constant (V3 0 del 0))
      gy2 = e (p + constant (V3 0 del 0))
      gz1 = e (p - constant (V3 0 0 del))
      gz2 = 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 :: Scene -> 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 < sceneMaxiter scene && r > 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 = (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')

simpleShading :: Scene -> Exp Direction -> Exp Float -> Exp Colour
simpleShading scene d noise =
  let o = sceneEye scene + noise *^ d
      (r :: Exp Float, p :: Exp Position) = unlift $ traceRay scene o d
      n = normEstimate scene p
   in r < sceneEps scene ?
      (applyLights (sceneLights scene) (sceneEye scene) p n, sceneColour scene)