module Optimus.CallGraph where
	import Flite.Pretty
	import Flite.Syntax
	import Flite.Traversals
	import Optimus.Uniplate
	import Optimus.Util
	
	import Control.Monad
	import Control.Monad.Fix
	import Control.Monad.State
	import qualified Data.Map as Map
	import qualified Data.Set as Set
	import Data.Generics.Uniplate
	import Data.Graph.Inductive.Graph
	import Data.Graph.Inductive.PatriciaTree
	import Data.GraphViz hiding (Int)
	import Data.List
	import Data.Maybe
	import Debug.Trace
	
	callGraph :: Prog -> Gr Id Id
	callGraph p = mkGraph (zip [1..] nodes) edges
		where
			nodes = nub $ [ f | Func f _ _ <- p ]
			lookupNodes = flip lookup (zip nodes [1..])
			edges = nub . catMaybes $ [ liftM2 (\x y -> (x, y, "")) (lookupNodes f) (lookupNodes g) | Func f _ r <- p, Fun g <- universe r ]
	
	conGraph :: Prog -> Gr Id [Id]
	conGraph p = mkGraph (zip [1..] nodes) edges
		where
			nodes = nub $ [ f | Func f _ _ <- p ]
			lookupNodes = flip lookup (zip nodes [1..])
			edges = nub . catMaybes $ [ liftM2 (\x y -> (x, y, consUsed r)) (lookupNodes f) (lookupNodes g) | Func f _ r <- p, Fun g <- universe r ]
			consUsed e = nub $ [ c | Con c <- universe e ] ++ [ c | Case _ as <- universe e, (p', _) <- as, Con c <- universe p' ]
	
	produceDot :: (Graph gr, Show b) => gr Id b -> String
	produceDot g = printDotGraph $ graphToDot True g [] labelNode labelEdge
		where
			labelNode (_, l) = [Label (StrLabel l), Color [ColorName $ if words l !! 0 `elem` primitives then "black" else "red"]]
			labelEdge (_, _, l) = if show l == "\"\"" then [] else [Label (StrLabel $ show l)]
			
	-----------------------
	
	arityGraph :: Prog -> Gr Id Id
	arityGraph p = mkGraph (zip [1..] nodes) edges
		where
			nodes = nub $ [ f | Func f _ _ <- p ]
			lookupNodes = flip lookup (zip nodes [1..])
			arities = [ (f, length a) | Func f a _ <- p ] ++ map (flip (,) 2) primitives
			arity f = fromMaybe (error "Missing function definition") $ lookup f arities
			edges = nub . catMaybes $ [ liftM2 (\x y -> (x, y, "")) (lookupNodes f) (lookupNodes g) | Func f _ r <- p, App (Fun g) ys <- universe r, arity f <= length ys ]
	
	hofGraph :: Prog -> Gr Id Id
	hofGraph p = mkGraph (zip [1..] nodes) edges
		where
			nodes = nub $ [ f | Func f _ _ <- p ]
						 ++ [ concat $ intersperse " " $ g : [ if hof e then showArg e else "_" | e <- ys ]
							| Func f _ r <- p, App (Fun g) ys <- universe r, arity f <= length ys, any hof ys ]
			lookupNodes = flip lookup (zip nodes [1..])
			arities = [ (f, length a) | Func f a _ <- p ] ++ map (flip (,) 2) primitives
			arity f = fromMaybe (error "Missing function definition") $ lookup f arities
			edges = nub . catMaybes $ [ if any hof ys then Nothing else Just ()
										>> liftM2 (\x y -> (x, y, "")) (lookupNodes f) (lookupNodes g)
										| Func f _ r <- p, App (Fun g) ys <- universe r, arity f <= length ys ]
			hof (App (Fun f) ys) = arity f > length ys
			hof (Fun f) = arity f > 0
			hof _ = False
	
	data TravArcType = Call | Composition
	
	travGraph :: Prog -> Gr Id TravArcType
	travGraph p = mkGraph (zip [1..] nodes) edges
		where
			edgesRaw = trav p
			nodes = nub . uncurry (++) . (\(_, x, y) -> (x, y)) . unzip3 $ edgesRaw
			lookupNodes = fromJust . flip lookup (zip nodes [1..])
			edges = [ (lookupNodes f, lookupNodes g, t) | (t@(Call), f, g) <- edgesRaw ]
	
	produceDotTrav :: Gr Id TravArcType -> String
	produceDotTrav g = printDotGraph $ graphToDot True g [] labelNode labelEdge
		where
			labelNode (_, l) = [Label (StrLabel l)] ++ if words l !! 0 `elem` primitives then [Style [SItem Dotted []]] else []
			labelEdge (_, _, t) = case t of { Call -> [] ; Composition -> [Style [SItem Dashed []], ArrowHead emptyArr] }
	
	trav :: Prog -> [(TravArcType, Id, Id)]
	trav p = fst $ execState (t [] "main" (body "main")) ([], [])
		where
			m = byFuncName (p ++ [ Func f [Var "_", Var "_"] Bottom | f <- primitives ])
			body f = fromMaybe (error "Missing function definition") $ Map.lookup f m >>= return . funcRhs
			arity = length . args
			args f = fromMaybe (error "Missing function definition") $ Map.lookup f m >>= return . concatMap patVars . funcArgs
			unsat (App (Fun f) ys) = arity f > length ys
			unsat (Fun f) = arity f > 0
			unsat _ = False
			t :: [(Id, Exp)] -> Id -> Exp -> State ([(TravArcType, Id, Id)], [Id]) Exp
			t env f (App (App x ys) zs) = t env f (App x $ ys ++ zs)
			t env f (Var v) = return $ fromMaybe (Con "") $ lookup v env
			t env f (App (Fun g) ys) | arity g <= length ys = do
										(oldEdges, oldNames) <- get
										put (oldEdges, [])
										bs <- sequence [ t env f e >>= \e' -> return $ if unsat e' then Just (v, e') else Nothing | (v, e) <- zip (args g) ys ]
										g' <- return . concat . intersperse " " $ g : (map (\x -> fromMaybe "_" (return . show . snd =<< x)) $ bs)
										(newEdges, newNames) <- get
										put ((Call, f, g'): map ((,,) Composition g') newNames ++ newEdges, [])
										e' <- if g' `notElem` [ x | (Call, _, x) <- newEdges]
											then t (catMaybes bs) g' (body g)
											else return Bottom
										(newerEdges, _) <- get
										put (newerEdges, g':oldNames)
										return e'
			t env f (App x ys) = do
							x' <- t env f x
							if x' == x || x' == Bottom then mapM (t env f) ys >>= return . App x' else t env f (App x' ys)
			t env f (Let bs x) = do
							bs' <- sequence [ liftM ((,) v) (t env f e) | (v, e) <- bs ]
							t (bs ++ env) f x
			t env f (Case x as) = do
							x <- t env f x
							ys <- sequence [ t (map (flip (,) Bottom) (patVars p) ++ env) f y | (p, y) <- as ]
							return $ head (ys ++ [x])
			t env f e = return e
	
	selectFuncs n p = nub $ selectFuncs2 n p
	
	selectFuncs1 :: Prog -> [Id]
	selectFuncs1 p = (evalState (traverse "main") Set.empty)
		where
			m = byFuncName p
			body f = fromMaybe Bottom (Map.lookup f m >>= return . funcRhs)
			traverse :: Id -> State (Set.Set Id) [Id]
			traverse f = let
				cs = filter (flip notElem $ f : primitives) . calls . body $ f
				csSet = Set.fromList cs in do
					visited <- get
					put $ csSet `Set.union` visited
					rest <- mapM traverse . Set.toList $ csSet `Set.difference` visited
					return $ if length cs > 1 then concat rest ++ [f] else concat rest
					
	selectFuncs2 :: Int -> Prog -> [Id]
	selectFuncs2 n p = take n [ f | (f, n) <- selected, n >= limit ]
		where
			selected :: [(Id, Int)]
			selected = map (\f -> maybe (f, -1) ((,) f) $ lookup f call_count) $ selectFuncs1 p
			call_count :: [(Id, Int)]
			call_count = map (\xs@(x:_) -> (x, length xs)) . group . sort . concatMap (calls . funcRhs) $ p
			limit :: Int
			limit | length selected > n = (reverse . sort . map snd $ selected) !! n
				  | otherwise			= 0