-- | The Grid module just contains the Grid type and two constructors
--   for it. We hide the main Grid constructor because we don't want
--   to allow instantiation of a grid with h <= 0.
module Grid
where

import qualified Data.Array.Repa as R
import Test.QuickCheck (Arbitrary(..), Gen, Positive(..))

import Cube (Cube(Cube), find_containing_tetrahedra)
import FunctionValues
import Point (Point)
import ScaleFactor
import Tetrahedron (polynomial)
import Values (Values3D, dims, empty3d, zoom_shape)


-- | Our problem is defined on a Grid. The grid size is given by the
--   positive number h. The function values are the values of the
--   function at the grid points, which are distance h from one
--   another in each direction (x,y,z).
data Grid = Grid { h :: Double, -- MUST BE GREATER THAN ZERO!
                   function_values :: Values3D }
          deriving (Eq, Show)


instance Arbitrary Grid where
    arbitrary = do
      (Positive h') <- arbitrary :: Gen (Positive Double)
      fvs <- arbitrary :: Gen Values3D
      return (make_grid h' fvs)


-- | The constructor that we want people to use. If we're passed a
--   non-positive grid size, we throw an error.
make_grid :: Double -> Values3D -> Grid
make_grid grid_size values
    | grid_size <= 0 = error "grid size must be positive"
    | otherwise = Grid grid_size values


-- | Creates an empty grid with grid size 1.
empty_grid :: Grid
empty_grid = Grid 1 empty3d


-- | Returns a three-dimensional list of cubes centered on the grid
--   points of g with the appropriate 'FunctionValues'.
cubes :: Grid -> [[[Cube]]]
cubes g
    | xsize == 0 || ysize == 0 || zsize == 0 = [[[]]]
    | otherwise =
        [[[ Cube (h g) i j k (make_values fvs i j k) | i <- [0..xsize]]
              | j <- [0..ysize]]
              | k <- [0..zsize]]
    where
      fvs = function_values g
      (xsize, ysize, zsize) = dims fvs


-- | Takes a grid and a position as an argument and returns the cube
--   centered on that position. If there is no cube there (i.e. the
--   position is outside of the grid), it will throw an error.
cube_at :: Grid -> Int -> Int -> Int -> Cube
cube_at g i j k
    | i < 0 = error "i < 0 in cube_at"
    | j < 0 = error "j < 0 in cube_at"
    | k < 0 = error "k < 0 in cube_at"
    | otherwise =
        let zsize = length (cubes g) in
          if k >= zsize then
            error "k >= xsize in cube_at"
          else
            let ysize = length ((cubes g) !! k) in
              if j >= ysize then
                error "j >= ysize in cube_at"
              else
                let xsize = length (((cubes g) !! k) !! j) in
                  if i >= xsize then
                    error "i >= xsize in cube_at"
                  else
                    (((cubes g) !! k) !! j) !! i


--   The first cube along any axis covers (-h/2, h/2). The second
--   covers (h/2, 3h/2).  The third, (3h/2, 5h/2), and so on.
--
--   We translate the (x,y,z) coordinates forward by 'h/2' so that the
--   first covers (0, h), the second covers (h, 2h), etc. This makes
--   it easy to figure out which cube contains the given point.
calculate_containing_cube_coordinate :: Grid -> Double -> Int
calculate_containing_cube_coordinate g coord
    -- Don't use a cube on the boundary if we can help it. This
    -- returns cube #1 if we would have returned cube #0 and cube #1
    -- exists.
    | coord == offset && (xsize > 0 && ysize > 0 && zsize > 0) = 1
    | otherwise = (ceiling ( (coord + offset) / cube_width )) - 1
    where
      (xsize, ysize, zsize) = dims (function_values g)
      cube_width = (h g)
      offset = cube_width / 2


-- | Takes a 'Grid', and returns a 'Cube' containing the given 'Point'.
--   Since our grid is rectangular, we can figure this out without having
--   to check every cube.
find_containing_cube :: Grid -> Point -> Cube
find_containing_cube g p =
    cube_at g i j k
    where
      (x, y, z) = p
      i = calculate_containing_cube_coordinate g x
      j = calculate_containing_cube_coordinate g y
      k = calculate_containing_cube_coordinate g z


{-# INLINE zoom_lookup #-}
zoom_lookup :: Grid -> ScaleFactor -> a -> (R.DIM3 -> Double)
zoom_lookup g scale_factor _ = zoom_result g scale_factor


{-# INLINE zoom_result #-}
zoom_result :: Grid -> ScaleFactor -> R.DIM3 -> Double
zoom_result g (sfx, sfy, sfz) (R.Z R.:. i R.:. j R.:. k) =
  f p
  where
    i' = (fromIntegral i) / (fromIntegral sfx)
    j' = (fromIntegral j) / (fromIntegral sfy)
    k' = (fromIntegral k) / (fromIntegral sfz)
    p  = (i', j', k') :: Point
    c = find_containing_cube g p
    t = head (find_containing_tetrahedra c p)
    f = polynomial t


zoom :: Grid -> ScaleFactor -> Values3D
zoom g scale_factor
    | xsize == 0 || ysize == 0 || zsize == 0 = empty3d
    | otherwise =
        R.force $ R.traverse arr transExtent (zoom_lookup g scale_factor)
          where
            arr = function_values g
            (xsize, ysize, zsize) = dims arr
            transExtent = zoom_shape scale_factor