module BabyRays.Sample where

import Core.Constants
import Core.Geometry
import BabyRays.Camera
import BabyRays.Types
import BabyRays.Scene

import Control.Monad.State
import System.Random
import Debug.Trace (trace)

sample :: Camera -> DFloat -> DFloat -> [ Point ]
sample camera width height =
   let (Camera camerapos viewdist lookingat viewup) = camera
       grid = [ (Point x y 0) | y <- [0..width-1], x <- [0..height-1] ]
       centering_offset = (Vector (- width / 2.0) (-height / 2.0)  0)
       pixel_offsets = map (.+^ centering_offset) grid
       viewdir = hat (lookingat .-. camerapos)
       screen_center = camerapos .+^ (viewdist *^ viewdir)
       viewright = viewdir `cross` viewup
       (Point x y z) .+. (Point a b c) = Point (x+a) (y+b) (z+c)
       transform (Point x y _) = (screen_center .+^ (x *^ viewright)) .+^ (y *^ (neg viewup))
   in map transform pixel_offsets

gridSample :: Camera                       -- | Camera object
              -> Image                     -- | Image triple
              -> (Int, Int)                -- | horizontal and vertical resolution
              -> Int                       -- | sqrt of samples per pixel
              -> [[(Pixel, [Ray])]]        -- | List of pixel rows
gridSample camera (sw, nw, se) (imHeight, imWidth) rate = pixels
  where height = norm (nw .-. sw)
        width  = norm (se .-. sw)

        unitY = hat (nw .-. sw)
        unitX = hat (se .-. sw)

        pxHeight = height / (fromIntegral imHeight)
        pxWidth  = width / (fromIntegral imWidth)

        pxBase (x,y) = let (x', y') = (pxWidth * fromIntegral x, pxHeight * fromIntegral y)
                       in sw .+^ ((x' *^ unitX) ^+^ (y' *^ unitY))

        mkPixel p@(x,y) = let base = pxBase p
                              horiz = linspace rate base (pxBase (x+1,y))
                              vert  = linspace rate base (pxBase (x,y+1))
                          in concat [map (.+^ (v .-. base)) horiz | v <- vert]

        mkRay = map (perspective camera) . mkPixel

        pixels = [[let p = (x,y) in (p, mkRay p) | x <- [0..(imWidth-1)]] | y <- [0..(imHeight-1)]]


jitterSample :: StdGen                     -- | PRNG Seed
              -> Camera                    -- | Camera object
              -> Image                     -- | Image triple
              -> (Int, Int)                -- | horizontal and vertical resolution
              -> Int                       -- | sqrt of samples per pixel
              -> [[(Pixel, [Ray])]]        -- | list of pixel rows
jitterSample gen camera (sw, nw, se) (imHeight, imWidth) rate = evalState pixels gen
  where height = norm (nw .-. sw)
        width  = norm (se .-. sw)

        unitY = hat (nw .-. sw)
        unitX = hat (se .-. sw)

        pxHeight = height / (fromIntegral imHeight)
        pxWidth  = width / (fromIntegral imWidth)

        randomX = state $ randomR (0, pxWidth)
        randomY = state $ randomR (0, pxHeight)

        mkRay :: Pixel -> State StdGen Ray
        mkRay p = do
          rx <- randomX
          ry <- randomY
          let base = pxBase p
          return . perspective camera $ base .+^ ((rx *^ unitX) ^+^ (ry *^ unitY))

        mkPixel :: Pixel -> State StdGen (Pixel, [Ray])
        mkPixel p = do
          rays <- mapM (const $ mkRay p) [1..rate*rate]
          return (p, rays)

        mkRow = mapM mkPixel
        mkMat = mapM mkRow

        pxBase (x,y) = let (x', y') = (pxWidth * fromIntegral x, pxHeight * fromIntegral y)
                       in sw .+^ ((x' *^ unitX) ^+^ (y' *^ unitY))

        pixels = mkMat [[(x,y) | x <- [0..(imWidth-1)]] | y <- [0..(imHeight-1)]]


-- jitterSample :: Int                          -- | PRNG Seed
--                 -> Camera                    -- | Camera object
--                 -> Image                     -- | Image triple
--                 -> (Int, Int)                -- | horizontal and vertical resolution
--                 -> Int                       -- | Samples per pixel
--                 -> [[(Pixel, [Ray])]]        -- | List of pixel rows
-- jitterSample seed camera (sw, nw, se) (imHeight, imWidth) rate = group imWidth pixels
--   where height = norm (nw .-. sw)
--         width  = norm (se .-. sw)

--         unitY = hat (nw .-. sw)
--         unitX = hat (se .-. sw)

--         (g1, g2) = split (mkStdGen seed)
--         randXs = randomRs (0, pxWidth) g1
--         randYs = randomRs (0, pxHeight) g2
--         rands  = group rate $ zipWith (\x y -> Vector x y 0) randXs randYs

--         pxHeight = height / (fromIntegral imHeight)
--         pxWidth  = width / (fromIntegral imWidth)

--         pxBase (x,y) = let (x', y') = (fromIntegral x, fromIntegral y)
--                        in sw .+^ ((x' *^ unitX) ^+^ (y' *^ unitY))

--         mkPixel (p, rs) = map (pxBase p .+^) rs

--         mkRay :: (Pixel, [Vector]) -> [Ray]
--         mkRay = map (parallel_projection camera) . mkPixel

--         pixels = [let p = (x,y) in (p, mkRay (p,r)) | (r, x) <- zip rands [0..imWidth], y <- [0..imHeight]]

--         group :: Int -> [a] -> [[a]]
--         group _ [] = []
--         group n l = take n l : group n (drop n l)

-- -- gridSample view (defaultPoint, defaultPoint .+^ (Vector 0 1 0), defaultPoint .+^ (Vector 1 0 0)) (640, 480)