module Observe where

import Data.IORef
import System.IO.Unsafe
import List (intersperse,nub,sort)
import Char (chr)
import Control.Exception(catch)
import Prelude hiding (catch)

{-- for older Hugs versions
-- unfortunately Ctrl-C does not work
-- you have to add modifyIORef (at the end of this file) as well!
import IOExts
import List (intersperse,nub)
-}

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

data EvalTree = Cons Int String [EvalTreeRef] 
	      | Fun Int [(Int,EvalTreeRef,EvalTreeRef)]
	      | Uneval 
	      | Demand Int

type EvalTreeRef = IORef EvalTree

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

showEval :: Int -> Int -> IO String -> IO String
showEval n m act = if n < m then return "_" else act

showEvalTreeRef :: Int -> EvalTreeRef -> IO String
showEvalTreeRef n r = do
  val <- readIORef r
  xs <- showEvalTree n val
  return xs

showEvalTree :: Int -> EvalTree -> IO String
showEvalTree _ Uneval = return "_"
showEvalTree n (Demand m) = showEval n m $ return "!"
showEvalTree n (Cons m  cons []) = showEval n m $ return cons
showEvalTree n (Cons m "(,)" ts) = showEval n m $ do
  [aa,ab] <- mapM (showEvalTreeRef n) ts
  return ("(" ++ aa ++ "," ++ ab ++ ")")
showEvalTree n (Cons m "(,,)" ts) = showEval n m $ do
  [aa,ab,ac] <- mapM (showEvalTreeRef n) ts
  return ("(" ++ aa ++ "," ++ ab ++ "," ++ ac ++ ")")
showEvalTree n (Cons m "(:)" ts) = showEval n m $ do
  [aa,ab] <- mapM (showEvalTreeRef n) ts
  return ("(" ++ aa ++ ":" ++ ab ++ ")")
showEvalTree n (Cons m cons ts) = showEval n m $ do
  args <- mapM (showEvalTreeRef n) ts
  return ("(" ++ concat (intersperse " " (cons:args)) ++ ")")
showEvalTree n (Fun m appls) = showEval n m $ do
  resStrs <- mapM showApp (reverse (filter (\ (k,_,_) -> k <= n) appls))
  return (concat (intersperse "\n" resStrs))
 where showApp (_,rArg,rRes) = do
         arg <- showEvalTreeRef n rArg
         res <- showEvalTreeRef n rRes
	 return ("{" ++ arg ++ " -> " ++ res ++ "}")

moments :: EvalTreeRef -> IO [Int]
moments r = readIORef r >>= evalMoments

evalMoments :: EvalTree -> IO [Int]
evalMoments Uneval = return []
evalMoments (Demand n) = return [n]
evalMoments (Cons n _ rs) = collect moments n rs
evalMoments (Fun n apps) = collect appMoments n apps
 where
  appMoments (m,arg,res) = collect moments m [arg,res]

collect :: (a -> IO [Int]) -> Int -> [a] -> IO [Int]
collect moms n xs = mapM moms xs >>= return . (n:) . concat

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

type Obs a = a -> EvalTreeRef -> a 

class Observe a where
  obs :: Obs a

observer :: Observe a => Obs a 
observer x r = unsafePerformIO $ do
  n <- getCount
  writeIORef r (Demand n) -- n is not accessed by current implementation
  return (obs x r)

oPrim :: Show a => Obs a
oPrim x r = unsafePerformIO $ do
      seq x (return ())
      -- if show x == show x then (return ()) else (return ()) 
      mkEvalTreeCons (show x) r 0
      return x 

o0 :: a -> String -> EvalTreeRef -> a
o0 cons consName r = unsafePerformIO $ do
      mkEvalTreeCons consName r 0
      return cons

o1 :: Observe a => (a -> b) -> String -> a -> EvalTreeRef -> b
o1 cons consName vA r = unsafePerformIO $ do
      [aRef] <- mkEvalTreeCons consName r 1
      return (cons (observer vA aRef))

o2 :: (Observe a,Observe b) => (a -> b -> c) -> String -> a -> b ->
      EvalTreeRef -> c
o2 cons consName vA vB r = unsafePerformIO $ do
      [aRef,bRef] <- mkEvalTreeCons consName r 2
      return (cons (observer vA aRef)
                   (observer vB bRef)) 

o3 :: (Observe a,Observe b,Observe c)
   => (a -> b -> c -> d) -> String -> a -> b -> c ->
      EvalTreeRef -> d
o3 cons consName vA vB vC r = unsafePerformIO $ do
      [aRef,bRef,cRef] <- mkEvalTreeCons consName r 3
      return (cons (observer vA aRef)
                   (observer vB bRef)
                   (observer vC cRef)) 

mkEvalTreeCons :: String -> EvalTreeRef -> Int ->
                  IO [EvalTreeRef]
mkEvalTreeCons consName r n = do
  refs <- mapM (const (newIORef Uneval)) [1..n]
  n <- getCount
  writeIORef r (Cons n consName refs)
  return refs

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

instance (Observe a,Observe b) => Observe (a,b) where
  obs (a,b) = o2 (,) "(,)" a b

instance (Observe a,Observe b) => Observe (a -> b) where
  obs f r x = unsafePerformIO $ do
     n <- getCount
     (m,applRefs) <- readIOFunRef n r
     [rArg,rRes] <- mapM (const (newIORef Uneval)) [1..2]
     writeIORef r (Fun m ((n,rArg,rRes):applRefs))
     return (observer (f (observer x rArg)) rRes)
   where readIOFunRef
           :: Int -> EvalTreeRef -> IO (Int,[(Int,EvalTreeRef,EvalTreeRef)])
	 readIOFunRef n r = do
           v <- readIORef r
	   case v of
             Fun m applRefs -> return (m,applRefs)
	     _            -> do
                 writeIORef r (Fun n [])
		 return (n,[])

instance Observe a => Observe [a] where
  obs (x:xs) = o2 (:) "(:)" x xs
  obs []     = o0 [] "[]"

instance Observe Int where
  obs = oPrim

instance Observe Char where
  obs = oPrim

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

global :: IORef [(String, IO [Int], Int -> IO ())]
global = unsafePerformIO $ newIORef []

counter :: IORef Int
counter = unsafePerformIO $ newIORef 1

getCount :: IO Int
getCount = do
  n <- readIORef counter
  writeIORef counter (n+1)
  return n

observe :: Observe a => String -> a -> a
observe label x = unsafePerformIO $ do
  r <- newIORef Uneval
  modifyIORef global
   (( label
    , moments r
    , (\n -> putStrLn (label++"\n"++replicate (length label) '-')
             >> showEvalTreeRef n r >>= putStrLn)):)
  return (observer x r)

run :: IO () -> IO ()
run io = do
  writeIORef global []
  writeIORef counter 1
  -- catch only works with ghc or hugs Version, newer than Decemper 2001
  -- catch has to be imported from module Control.Ecxeption
  catch io (\e -> putStrLn ("Runtime Error\n"++show e))
  --io
  printObs

printObs = do
  putStrLn ">>> Observations <<<"
  obs <- readIORef global
  n <- readIORef counter
  sequence (map (\ (_,_,ob) -> ob n) obs)
  return ()

runI :: IO () -> IO ()
runI io = do
  writeIORef global []
  writeIORef counter 1
  -- catch only works with ghc or hugs Version, newer than Decemper 2001
  -- catch has to be imported from module Control.Ecxeption
  catch io (\e -> putStrLn ("Runtime Error\n"++show e))
  --io
  interactiveObs

interactiveObs = do
  putStrLn ">>> Observations <<<"
  obs <- readIORef global
  queryLoop obs

queryLoop obs = do
  putStr "label (or 'quit'): "
  label <- getLine
  if label == "quit" then return ()
   else do
    stepObservations $ filter (\ (l,_,_) -> l==label) obs
    queryLoop obs

stepObservations :: [(String, IO [Int], Int -> IO ())] -> IO ()
stepObservations msobs
  = sequence moms >>= mapM_ step . (0:) . nub . sort . concat
 where
  (moms,os) = foldr (\ (_,ms,ob) (mss,obs)->(ms:mss,ob:obs)) ([],[]) msobs
  step n = do
    putStr (chr 27:"[2J"++chr 27:"[H") -- clear screen and move curser home
    sequence (map ($n) os)
    getLine

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

{- for Hugs:
modifyIORef :: IORef a -> (a -> a) -> IO ()
modifyIORef ref f =
  readIORef ref >>= (writeIORef ref) . f
-}