{-
  Copyright (C) 2010 by IPwn Studios
  Released under GNU General Public License v3
-}

{-# LANGUAGE ScopedTypeVariables #-}
module RasterizeAcc where

import Control.Arrow
import Control.Parallel
import Control.Parallel.Strategies
import Data.Array.Accelerate
import Data.Array.Accelerate.Interpreter
import Data.Bits ((.&.))
import Prelude   hiding (replicate, zip, unzip, map, filter, max, min, not,
    zipWith, round, uncurry, scanl, fst, snd, tail, even)
import qualified Prelude
import Debug.Trace


type Area = ((Int, Int), (Int, Int))
type Value = ((Int, Int), Float)
type Insertion = ((Int, Int), Float)
type Facet = (Int, Int, Int)

aseq :: (NFData a, Elem a, Ix dim) => String -> Acc (Array dim a) -> Acc (Array dim a)
aseq descr = id -- use . trace descr . hseq . run

aseq1 :: (NFData a, Elem a, Ix dim, Show a) => String -> Acc (Array dim a) -> Acc (Array dim a)
aseq1 descr aarr =
    let arr = run aarr
    in  use $ {-trace (descr ++ " " ++ show (arrayShape arr) ++ " " ++ show arr)-} arr

hseq :: (NFData a, Elem a, Ix dim) => Array dim a -> Array dim a
hseq arr =
    let l = toList arr `using` seqList rwhnf
    in  l `seq` fromList (arrayShape arr) l

rasterize :: Exp Area
          -> Acc (Vector Value)
          -> Acc (Vector Facet) 
          -> Acc (Array DIM2 Float)
rasterize area values facets =
    let -- Triangle vertices sorted by X co-ordinate
        tris :: Acc (Vector (Value, Value, Value))
        tris = flip map facets $ \tri ->
            let (iA, iB, iC) = untuple tri
            in  sort3Tuple lessThanX $ tuple (values ! iA, values ! iB, values ! iC)

        noOfTris = shape tris :: Exp Int

        -- Vector of (triangle index, (start column, end column))
        columnRanges :: Acc (Vector (Int, Int))
        columnRanges = aseq1 "columnRanges" $ map rangeOfColumns tris

        -- Vector of (triangle index, integer x co-ordinate of column)
        columnIndices :: Acc (Vector (Int, Int))
        columnIndices = aseq1 "columnIndices" $ flatten columnRanges

        columns :: Acc (Vector ((Int, Float, Float), (Int, Float, Float)))
        columns = aseq1 "columns" $ flip map columnIndices $ \pair ->
            let (triIx, xInt) = untuple pair
            in  rowRange (tris ! triIx) xInt

        pointRanges :: Acc (Vector (Int, Int))
        pointRanges = aseq1 "pointRanges" $ flip map columns $ \col ->
            let ((yStart, _, _), (yEnd, _, _)) = (untuple *** untuple) $ untuple col
                    :: ((Exp Int, Exp Float, Exp Float), (Exp Int, Exp Float, Exp Float))
            in  tuple (yStart, yEnd)

        -- Vector of (column index, integer y co-ordinate of column)
        pointIndices :: Acc (Vector (Int, Int))
        pointIndices = aseq1 "pointIndices" $ flatten pointRanges

        insertions :: Acc (Vector Insertion)
        insertions = aseq1 "insertions" $ flip map pointIndices $ \pi ->
            let (colIx, yInt) = untuple pi
                (triIx, xInt) = untuple $ columnIndices ! colIx
                    :: (Exp Int, Exp Int)
                ((_, yBot, vBot), (_, yTop, vTop)) =
                    (untuple *** untuple) $ untuple $ columns ! colIx
                    :: ((Exp Int, Exp Float, Exp Float), (Exp Int, Exp Float, Exp Float))
                -- Displace odd columns because the terrain mesh is hexagonal
                hexify yInt = even xInt ? (intToFloat yInt, intToFloat yInt + 0.5)
                v = interpolate (yBot, vBot) (yTop, vTop) (hexify yInt)
            in  tuple (tuple (xInt, yInt), v)
    in  toArray area insertions
  where
    ((x0,y0),(x1,y1)) = (untuple *** untuple) $ untuple area  :: ((Exp Int, Exp Int), (Exp Int, Exp Int))

    -- Give the start and end of the Y value of the specified row in both
    -- floating point and integer co-ordinates.  Output is a pair of
    -- (integer y, float y, value)
    rowRange :: Exp (Value, Value, Value) -> Exp Int -> Exp ((Int, Float, Float), (Int, Float, Float))
    rowRange tri xInt =
        let (a, b, c) = untuple tri
            ((axInt, ayInt), vA) = (untuple *** id) (untuple a) :: ((Exp Int, Exp Int), Exp Float)
            ((bxInt, byInt), vB) = (untuple *** id) (untuple b) :: ((Exp Int, Exp Int), Exp Float)
            ((cxInt, cyInt), vC) = (untuple *** id) (untuple c) :: ((Exp Int, Exp Int), Exp Float)
            ax = intToFloat axInt
            ay = intToFloat ayInt
            bx = intToFloat bxInt
            by = intToFloat byInt
            cx = intToFloat cxInt
            cy = intToFloat cyInt
            x = intToFloat xInt
            yAC = interpolate (ax, ay) (cx, cy) x :: Exp Float
            vAC = interpolate (ax, vA) (cx, vC) x :: Exp Float
            (yOther, vOther) = untuple (
                    x <* bx ? (tuple (interpolate (ax,ay) (bx,by) x, interpolate (ax,vA) (bx,vB) x),
                               tuple (interpolate (bx,by) (cx,cy) x, interpolate (bx,vB) (cx,vC) x))
                ) :: (Exp Float, Exp Float)
            -- Order AC and Other so that y is increasing
            (yBot, vBot, yTop, vTop) = untuple (
                    yAC <=* yOther ? (tuple (yAC, vAC, yOther, vOther),
                                      tuple (yOther, vOther, yAC, vAC))
                ) :: (Exp Float, Exp Float, Exp Float, Exp Float)
            yStart = max y0 (roundFloatToInt yBot)
            yEnd   = min (y1+1) (roundFloatToInt yTop)
        in  tuple (tuple (yStart, yBot, vBot), tuple (yEnd, yTop, vTop))

    rangeOfColumns :: Exp (Value, Value, Value) -> Exp (Int, Int)
    rangeOfColumns tri =
        let (a, _, c) = untuple tri :: (Exp Value, Exp Value, Exp Value)
            (xyA, _) = untuple a    :: (Exp (Int, Int), Exp Float)
            (xA, _)  = untuple xyA  :: (Exp Int, Exp Int)
            (xyC, _) = untuple c    :: (Exp (Int, Int), Exp Float)
            (xC, _)  = untuple xyC  :: (Exp Int, Exp Int)
        in  tuple (x0 `max` xA, (x1 + 1) `min` xC)

-- | Flatten a list of ranges with an exclusive upper value into a list of
-- indices into the concatenated ranges in the form (input ix, value in range). e.g.
--
-- > [(5,8),(10,15),(3,4)]
--
-- > [(0,5),(0,6),(0,7),(1,10),(1,11),(1,12),(1,13),(1,14),(2,3)]
flatten :: Acc (Vector (Int, Int))
        -> Acc (Vector (Int, Int))
flatten ranges0 =
    let -- ensure the ranges are not decreasing
        ranges = flip map ranges0 $ \ab ->
            let (a, b) = untuple ab
            in  tuple (a, a `max` b)
        noOfRanges = shape ranges :: Exp Int

        rangeWidths :: Acc (Vector Int)
        rangeWidths = aseq1 "f.rangeWidths" $ map (uncurry $ flip (-)) ranges

        -- The start position of each column in a yet-to-be-created array of columns
        (columnOffsets, noOfCols) = second (! constant ()) $ scanl (+) 0 $ rangeWidths
            :: (Acc (Vector Int), Exp Int)

        -- Input indices for each of the output elements, e.g. for the example
        -- above, [0,0,0,1,1,1,1,1,2]
        rangeIndices :: Acc (Vector Int)
        rangeIndices = aseq1 "f.rangeIndices" $
            let ones = replicate noOfRanges (unit 1) :: Acc (Vector Int)
                zeroes = replicate (noOfCols + 1) (unit 0) :: Acc (Vector Int)
                steps = permute (+) zeroes (columnOffsets !) ones
            in  tail $ Prelude.fst $ scanl (+) (-1) steps

        columnIndices :: Acc (Vector Int)
        columnIndices = aseq1 "f.columnIndices" $
            --   [10 - 8 + 1, 3 - 15 + 1] 
            --   [3, -11]
            -- 5 
            let ones = use $ run $ replicate (noOfCols + 1) (unit 1) :: Acc (Vector Int)
                -- ([5,10,3], [8,15,4]) 
                (rangeStarts, rangeEnds) = unzip ranges
                -- [3, -11]
                deltas = use $ run $ map (1+) $ zipWith (-) (tail rangeStarts) rangeEnds
                deltas_m1 = map (\x -> x - 1) deltas
                -- [1, 1, 1, 3, 1, 1, 1, 1,-11, 1]
                changes = use $ run $ permute (+) ones (\ix -> columnOffsets ! (ix + 1)) deltas_m1
            --   [5, 6, 7,10,11,12,13, 14, 3]
            in  tail $ Prelude.fst $ scanl (+) ((rangeStarts ! 0) - 1) changes

    in  zip rangeIndices columnIndices

tail :: Elem a => Acc (Vector a) -> Acc (Vector a)
tail xs = backpermute (shape xs - 1) (1+) xs

even :: Exp Int -> Exp Bool
even x = (x .&. 1) ==* 0

interpolate :: (Exp Float, Exp Float) -> (Exp Float, Exp Float) -> Exp Float -> Exp Float
interpolate (t0, x0) (t1, x1) t = (t - t0) * (x1 - x0) / (t1 - t0) + x0

-- | Give an array of values [0..length xs-1], e.g. [8,2,3,4] -> [0,1,2,3]
indicesOf :: Elem a => Acc (Vector a) -> Acc (Vector Int)
indicesOf = Prelude.fst . scanl (+) 0 . map (const 1)

lessThanX :: Exp Value -> Exp Value -> Exp Bool
lessThanX a b =
    let (xyA, _) = untuple a   :: (Exp (Int, Int), Exp Float)
        (xA, _)  = untuple xyA :: (Exp Int, Exp Int)
        (xyB, _) = untuple b   :: (Exp (Int, Int), Exp Float)
        (xB, _)  = untuple xyB :: (Exp Int, Exp Int)
    in  xA <* xB

sort3Tuple :: forall a . (Elem a) =>
              (Exp a -> Exp a -> Exp Bool)   -- ^ Less-than function
           -> Exp (a, a, a)                  -- ^ Tuple
           -> Exp (a, a, a)
sort3Tuple lessThan = shuttle1 . untuple
  where
    -- Shuttle sort
    shuttle1 :: (Exp a, Exp a, Exp a) -> Exp (a, a, a)
    shuttle1 (a, b, c) = a `lessThan` b ? (shuttle2 (a, b, c), shuttle2 (b, a, c))

    shuttle2 :: (Exp a, Exp a, Exp a) -> Exp (a, a, a)
    shuttle2 (a, b, c) = b `lessThan` c ? (tuple (a, b, c), shuttle3 (a, c, b))

    shuttle3 :: (Exp a, Exp a, Exp a) -> Exp (a, a, a)
    shuttle3 (a, b, c) = a `lessThan` b ? (tuple (a, b, c), tuple (b, a, c))

prop_sort3Tuple :: (Int, Int, Int) -> Bool
prop_sort3Tuple triple =
    let (a, b, c) = head $ toList $ run (map (sort3Tuple (<*)) $ unit $ constant triple)
    in  a <= b && b <= c

toArray :: Exp Area
        -> Acc (Vector Insertion)
        -> Acc (Array DIM2 Float)
toArray area insertions =
    let (bl, tr) = untuple area
        (x0, y0) = untuple bl
        (x1, y1) = untuple tr
        w = x1 - x0 + 1
        h = y1 - y0 + 1 
        z = zeroArray (tuple (h, w))
        (ixs, values) = unzip insertions
        perm :: Exp Int -> Exp (Int, Int)
        perm ix =
            let (x, y) = untuple (ixs ! ix)
            in
                x >=* x0 &&* x <=* x1 &&*
                y >=* y0 &&* y <=* y1
                    ? (tuple (y - y0, x - x0), ignore)
    in  permute const z perm values

zeroArray :: Exp (Int, Int) -> Acc (Array DIM2 Float)
zeroArray dim = replicate dim (unit 0)