{-# LANGUAGE PatternSignatures #-}

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

import Join.Base
import Join.Join
import Join.JoinPrettyPrinter

import Data.Time
import System.Time

{- UnionFind

union    @ Union(a,b) = do { x <- newVar
                           ; y <- newVar
                           ; Find(a,x) 
                           ; Find(b,y)
                           ; waitSync x
                           ; waitSync y
                           ; Link(x,y) }
                           
findNode @ Edge(a,b) \ Find(a,x) = Find(b,x)

findRoot @ Root(a) \ Find(a,x) = do { Found(a,x)
                                    ; x .=. () }
                                    
found    @ Edge(a,b) \ Found(a,x) = Found(b,x)

linkeq   @ Link(x,y),Found(a,x),Found(a,y) = return ()

link     @ Link(x,y),Found(a,x),Found(b,y),Root(a),Root(b) = do { Edge(b,a) 
                                                                ; Root(a) }

-}

union a b = method "Union" (a,b)
edge a b = method "Edge" (a,b)
find a x = method "Find" (a,x)
root a   = method "Root" a
found a x = method "Found" (a,x)
link x y = method "Link" (x,y)

unionfindRules ucount join activeMethod = do
  x <- newVar :: IO (VAR (Sync Bool))
  y <- newVar :: IO (VAR (Sync Bool))
  a <- newVar :: IO (VAR Int)
  b <- newVar :: IO (VAR Int)
  let prog = 
       translateJoinDefinitions
       [ 
         [union a b] .->. do { v_a <- readVar a
                             ; v_b <- readVar b
                             ; x'  <- newSync :: IO (Sync Bool)
                             ; y'  <- newSync :: IO (Sync Bool)
                             ; call join "Find" (v_a,x') 
                             ; call join "Find" (v_b,y')
                             ; waitSync x'
                             ; waitSync y'
                             ; call join "Link" (x',y') }
                             
       , ([edge a b] .\. [find a x]) .->. do { v_b <- readVar b
                                             ; v_x <- readVar x
                                             ; call join "Find" (v_b,v_x)
                                             ; return () }
                                             
       , ([root a] .\. [find a x]) .->. do { v_a <- readVar a
                                           ; v_x <- readVar x
                                           ; call join "Found" (v_a,v_x)
                                           ; x .=. True }
                                           
       , ([edge a b] .\. [found a x]) .->. do { v_b <- readVar b
                                              ; v_x <- readVar x
                                              ; call join "Found" (v_b,v_x)
                                              ; return () }
                                              
       , [link x y,found a x,found a y] .->. (atomically $ inc ucount)
       
       , [link x y,found a x,found b y,root a,root b] .->. do { v_a <- readVar a
                                                              ; v_b <- readVar b
                                                              ; call join "Edge" (v_b,v_a)
                                                              ; call join "Root" v_a 
                                                              ; atomically $ inc ucount }
       ]
              
  res <- runJoinOnGoal (store join) activeMethod prog
  case res of
    Just action -> action
    Nothing -> return ()  
  where
    inc tv = do { i <- readTVar tv
                ; writeTVar tv (i+1) }
    
-- Auxiliary code for tree construction

-- mkTrees join <Num of trees> <Size of trees> = <Arbitral leaves of each tree>
mkTrees :: Join -> Int -> Int -> IO [Int]
mkTrees join num size = 
   mkTrees' num 0
   where
     mkTrees' 0 _ = return []
     mkTrees' num curr = do
       { mkTree join [curr..(curr+size)]
       ; ls <- mkTrees' (num-1) (curr+size+1) 
       ; return $ (curr:ls) }

mkTree :: Join -> [Int] -> IO ()
mkTree join [] = return ()
mkTree join [s] = callPassive join "Root" s
mkTree join [s1,s2] = do { callPassive join "Edge" (s1,s2)
                         ; callPassive join "Root" s2 }
mkTree join (s1:s2:s3:ss) = do { callPassive join "Edge" (s1,s3)
                               ; callPassive join "Edge" (s2,s3) 
                               ; mkTree join (s3:ss) } 

unionPar :: Join -> [Int] -> IO ()
unionPar join (i1:i2:is) = do
  { forkIO $ call join "Union" (i1,i2)
  ; unionPar join (i2:is) }
unionPar _ _ = return ()
                                             
-- sample code

main :: IO ()
main = test1

test1 = do
    -- we count the number of successful execution of link
  { cnt <- atomically $ newTVar 0
  ; let no :: Int = 3

  ; jStore <- newJoinStore 
  ; let join = Join {store = jStore, rules = unionfindRules cnt}
  
  ; ls <- mkTrees join 4 3 -- 4 200
  
  ; s <- prettyIt jStore
  ; putStrLn s
  
  ; t1 <- getCurrentTime
  ; putStrLn "Started Clock"
  
  ; unionPar join ls

  ; atomically $ do { i <- readTVar cnt
                    ; if i >= no then return ()
                                 else retry }

  ; t2 <- getCurrentTime
  ; putStrLn "End Clock"
                               
  ; s <- prettyIt jStore
  ; putStrLn s
  
  ; let diff = diffUTCTime t2 t1
  ; putStrLn $ "\nTime Taken (Sec):\n" ++ (show diff) }