{-# LANGUAGE PatternSignatures   #-}

import IO
import Control.Concurrent
import Control.Concurrent.STM

import MultiSetRewrite.Base
import MultiSetRewrite.RuleSyntax
import MultiSetRewrite.RuleCompiler
import MultiSetRewrite.StoreRepresentation

import CHRSolver

{- Mergesort

  Leq(x,a) \ Leq(x,b) <==> a<b | Leq(a,b)
  Merge(n,a), Merge(n,b) <==> a<b | \Leq(a,b),\Merge(n+1,a)

-}

(.<.) x y = do
   v_x <- readVar x
   v_y <- readVar y
   return (v_x < v_y)

ruleMSort :: CHR -> ActiveConstraint -> IO ()
ruleMSort (chr@(chrD,chrR)) activeCons = do
  x <- newVar :: IO (VAR Int)
  a <- newVar :: IO (VAR Int)
  b <- newVar :: IO (VAR Int)
  n <- newVar :: IO (VAR Int)
  let prog = compileRulePattern
               [ ([cons "Leq" (x,a)] .\. [cons "Leq" (x,b)])
                  `when` (a .<. b)  .->. do v_a <- readVar a
                                            v_b <- readVar b
                                            stackGoal chr (cons "Leq" (v_a,v_b))

               , ([cons "Merge" (n,a), cons "Merge" (n,b)])
                  `when` (a .<. b)  .->. do v_a <- readVar a
                                            v_b <- readVar b
                                            v_n <- readVar n
                                            stackGoal chr (cons "Leq" (v_a,v_b))
                                            queueGoal chr (cons "Merge" (v_n + 1, v_a))

                 -- collecting the result

               , [cons "Start" (), cons "Merge" (n,a)] 
                     .->. do v_a <- readVar a
                             putStr $ show v_a
                             stackGoal chr (cons "Next" v_a)
               
               , [cons "Next" a, cons "Leq" (a,b)]
                     .->. do v_b <- readVar b
                             putStr $ show v_b
                             stackGoal chr (cons "Next" v_b)
               ]
  res <- executeRules (store chrD) activeCons prog
  case res of
         Just action -> action
         Nothing -> return ()



-- the input MUST be a list of size 2^k
msort :: [Int] -> IO ()
msort xs = do
  let threads = 4
  initStore <- newStore valHashOpMsg
  let initG  = [ cons "Merge" (1::Int,x) | x <- xs]
  initGoals <- atomically $ newTVar initG
  inActive <- atomically $ newTVar 0

  let chr = (Data {store = initStore,
                   goals = initGoals,
                   numberOfSolverThreads = threads,
                   inActiveThreads = inActive},
             Rules { rules = ruleMSort } )

  mapM_ (\x -> forkIO (solverThread chr)) [1..threads]

  suspendUntilRulesHaveExhaustivelyFired chr

  -- collect results

  atomically $ do v <- readTVar (inActiveThreads (fst chr))
                  writeTVar (inActiveThreads (fst chr)) (v-1)
  
  stackGoal chr (cons "Start" ())

  forkIO (solverThread chr)

  suspendUntilRulesHaveExhaustivelyFired chr

  putStr "Done"