{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Data.List (zipWith4)
import Control.Parallel.Eden
import Control.Parallel.Eden.Grace

---------------------------------------------------------------------
-- Master Worker core  through Graph spec

mwGrace:: forall t r. (Trans t, Trans r) =>
       Int -> Int        -- #workers, prefetch
       -> ([t] -> [r])   -- worker function
       -> [t] -> [[r]]   -- what to do
mwGrace np prefetch wf tasks = fst $ (start $ build (0, master) (number workers) edges)
    where
        master :: Lister ([[(Int, r)]] -> ([[r]], [(Int, t)]))
        master = lister (\xs -> (map (map snd) xs, zip (initReqs ++ map fst (merge xs)) tasks)) [np]

        initReqs = concat (replicate prefetch [0..np-1])

         -- worker specifications
        worker :: Int -> [t] -> [(Int, r)]
        worker n ts = zip [n, n..] $ wf ts

        workers :: [Function]
        workers = toFL [worker i | i <- [0..np-1]]

        -- edge definitions
        edges :: [Edge Int Int]
        edges = zipWith4 E [1..np] [0,0..] [1,1..] nothings
                ++ zipWith4 E [0,0..] [1..np] [1,1..] [Just (toWorkerSelect i) | i <- [0..np-1]]

        toWorkerSelect :: Int -> ([r], [(Int, t)]) -> [t]
        toWorkerSelect n (_, xs) = map snd $ filter ((==n) . fst) xs

---------------------------------------------------------------------
-- deterministic and nondeterministic interface to the master-worker-core

mwNonDetGrace, mwDetGrace :: (Trans t, Trans r) =>
       Int -> Int        -- #workers, prefetch
       -> ([t] -> [r])   -- worker function
       -> [t] -> [r]     -- what to do

mwNonDetGrace np prefetch wf tasks = merge $ mwGrace np prefetch wf tasks

mwDetGrace np prefetch wf tasks = map snd ( mergeByTags (mwGrace np prefetch wf' tasks'))
   where wf' ls  = let (lsI,lsC) = unzip ls
                       lsC'      = wf lsC
                   in zip lsI lsC'
         tasks'  = zip [1..] tasks 

---------------------------------------------------------------------
-- helper functions

mergeByTags :: [[(Int,r)]] -> [(Int,r)]
mergeByTags [] = []
mergeByTags [wOut] = wOut
mergeByTags [w1,w2] = merge2ByTag w1 w2
mergeByTags wOuts = merge2ByTag
                     (mergeHalf wOuts)
                     (mergeHalf (tail wOuts))
   where mergeHalf = mergeByTags . (takeEach 2)
         takeEach n [] = []
         takeEach n xs@(x:_) = x:takeEach n (drop n xs)

merge2ByTag [] w2 = w2
merge2ByTag w1 [] = w1
merge2ByTag w1@(r1@(i,_):w1s) w2@(r2@(j,_):w2s)
                       | i < j = r1: merge2ByTag w1s w2
                       | i > j = r2: merge2ByTag w1 w2s
                       | otherwise = error "found tags i == j"

---------------------------------------------------------------------


--A simple example call could be:

test :: [Int]
test = mwDetGrace (noPe-1) 10 (map (\x -> x*x)) [1..1000]

main = print test