import Data.List
import System.Random
import Criterion.Main
import Control.Monad.Par
import Control.Exception (evaluate)

parDivConq :: NFData b
           => (a -> Bool) -- test if problem is small enough
           -> (a -> [a])  -- or (a -> (a,a)), to split the problem into smaller ones
           -> (a -> b)    -- solver for sequential
           -> ([b] -> b)  -- or ((b,b) -> b), merge sub-solutions ...
           -> a           -- input
           -> Par b
parDivConq smallEnough splitInputs seqSolver merge input =
  if smallEnough input then return (seqSolver input)
  else do
    let ins = splitInputs input
    outs <- parMapM (parDivConq smallEnough splitInputs seqSolver merge) ins
    return $ merge outs

divConq :: NFData b
        => (a -> Bool) -- test if problem is small enough
        -> (a -> [a])  -- or (a -> (a,a)), to split the problem into smaller ones
        -> (a -> b)    -- solver for sequential
        -> ([b] -> b)  -- or ((b,b) -> b), merge sub-solutions ...
        -> a           -- input
        -> b           -- output
divConq smallEnough splitInputs seqSolver merge input =
  runPar $ parDivConq smallEnough splitInputs seqSolver merge input

mergeTwo :: Ord a => [a] -> [a] -> [a]
mergeTwo xs [] = xs
mergeTwo [] ys = ys
mergeTwo (x:xs) (y:ys) =
  if y < x then y : mergeTwo (x:xs) ys
  else x : mergeTwo xs (y:ys)

splitInTwo :: [a] -> ([a], [a])
splitInTwo [] = error "Cannot split two lists"
splitInTwo elements =
  let
    size = length elements
    middle  = size `div` 2
  in
    splitAt middle elements

seqFusionSort :: Ord a => [a] -> [a]
seqFusionSort [] = []
seqFusionSort [x] = [x]
seqFusionSort xs =
  let
    (as, bs) = splitInTwo xs
    sortedA = seqFusionSort as
    sortedB = seqFusionSort bs
  in
    mergeTwo sortedA sortedB

parFusionSort :: (NFData a, Ord a) => Int -> [a] -> [a]
parFusionSort sizeThreshold = divConq smallEnough splitInputs seqFusionSort (foldl mergeTwo [])
  where
    smallEnough :: [a] -> Bool
    smallEnough elements = length elements < sizeThreshold

    splitInputs :: [a] -> [[a]]
    splitInputs elements =
      let (xs, ys) = splitInTwo elements in
        [xs, ys]

main :: IO ()
main = do
  let n = 10^(7 :: Int)
  let elements = (take n (randoms (mkStdGen 211570155)) :: [Int] )
  putStrLn "Begin sort"
  _ <- evaluate $ parFusionSort (10^(6 :: Int)) elements
  putStrLn "Finished sorting."