module Optimus.Util where
	import Flite.Fresh
	import Flite.Syntax
	import Flite.Traversals hiding (freshen)
	import Optimus.Uniplate
	import Flite.Let
	import Flite.Matching
	
	import Data.Char
	import Data.List
	import Control.Monad
	import Data.Generics.Uniplate
	import qualified Data.Map as Map
	import Data.Maybe
	
	primitives = ["(+)", "(-)", "(==)", "(/=)", "(<=)", "emit", "emitInt"]
	
	freshen :: Exp -> Id -> Fresh Exp
	freshen e v = fresh >>= return . fix e . Var
		where
			fix (Case x as) w = Case x [ if (Var v `elem` universe p) then (subst w v p, subst w v y) else (p, y) | (p, y) <- as]
			fix (Let bs x) w@(Var w_) | v `elem` map fst bs = Let [ (if v == v_ then w_ else v_, subst w v y) | (v_, y) <- bs ] (subst w v x)
									  | otherwise			 = Let bs x
			fix _ _ = error "Not designed to freshen this."
	
	freshenMany :: Exp -> [Id] -> Fresh Exp
	freshenMany = foldM freshen
	
	freshenDescend :: Exp -> Id -> Fresh Exp
	freshenDescend e@(Case _ _) v = freshen e v >>= descendM (flip freshenDescend v)
	freshenDescend e@(Let _ _) v = freshen e v >>= descendM (flip freshenDescend v)
	freshenDescend e v = descendM (flip freshenDescend v) e
	
	freshenMD :: Exp -> [Id] -> Fresh Exp
	freshenMD = foldM freshenDescend
	
	autoFresh :: Fresh a -> a
	autoFresh x = snd $ runFresh x "v" 0
	
	freshProg :: (Prog -> Fresh Prog) -> Prog -> Prog
	freshProg f x = snd $ runFresh (f x) "v" $ (+) 1 $ maximum $ -1 : [ read rest | 'v':rest <- allNames x, (not . null) rest, all isDigit rest ]
	
	allNames :: Prog -> [Id]
	allNames p = nub $ concat [ [f] ++ concatMap patVars a ++ allnames r | Func f a r <- p ]
		where
			allnames (Var v) = [v]
			allnames (Fun f) = [f]
			allnames e@(Case _ as) = concatMap (patVars . fst) as ++ extract allnames e
			allnames e@(Let bs _) = map fst bs ++ extract allnames e
			allnames e = extract allnames e

	-- allNames p = concat [ [f] ++ (concatMap patVars a) ++ [ v | Var v <- universe r ] ++ (concat [ patVars p | Case _ as <- universe r, (p, _) <- as ]) | Func f a r <- p ]
		
	alphaExp :: Exp -> Fresh Exp
	alphaExp e = (alpha . Func "NOVAR" []) e >>= return . funcRhs
	
	autoAlphaExp :: Exp -> Exp
	autoAlphaExp e = snd $ runFresh (alphaExp e) "v" $ (+) 1 $ maximum $ -1 : [ read rest | 'v':rest <- freeVars e, (not . null) rest, all isDigit rest ]
	
	insertMany :: Ord k => Map.Map k a -> [(k, a)] -> Map.Map k a
	insertMany = foldr (uncurry Map.insert)
	
	alpha :: Decl -> Fresh Decl
	alpha (Func i a r) = alphaM
		where
			alphaM = do
				l <- assocFresh ([ v | Var v <- concatMap universe a ] ++ [i]) >>= return . Map.fromList
				r' <- freshSubst l r
				a' <- mapM (freshSubst l) a
				return $ Func (fromJust $ Map.lookup i l) a' r'
			assocFresh :: [Id] -> Fresh [(Id, Id)]
			assocFresh = mapM (\v -> fresh >>= \v' -> return (v, v'))
			freshSubst :: Map.Map Id Id -> Exp -> Fresh Exp
			freshSubst l (Var v) 	| isJust v' = return . Var . fromJust $ v'
									| otherwise	= return . Var $ v
				where
					v' = Map.lookup v l
			freshSubst l (Let bs x) = let (vs, ys) = unzip bs in do
				m <- assocFresh vs
				descendM (freshSubst $ insertMany l m) (Let (zip (map snd m) ys) x)
			freshSubst l (Case x as) = do
				x' <- freshSubst l x
				as' <- sequence [ assocFresh [v | Var v <- universe p] >>= \m -> let fS' = freshSubst $ insertMany l m in liftM2 (,) (fS' p) (fS' y) | (p, y) <- as]
				return (Case x' as')
			freshSubst l e = descendM (freshSubst l) e
			
	desugar :: Prog -> Fresh Prog
	desugar p = return p >>= desugarCase >>= desugarEqn >>= inlineLinearLet >>= inlineSimpleLet
	
	funcReuse :: Prog -> Prog
	funcReuse p = if null reused then p else error $ "Reused the name of a function as a local variable. " ++ show reused
		where
			reused = funcs p `intersect` (nub . concatMap (\d -> (concatMap patVars . funcArgs) d ++ (boundVars . funcRhs) d) $ p)
	
	byFuncName :: Prog -> Map.Map Id Decl
	byFuncName []					= Map.empty
	byFuncName (d@(Func f _ _):ds)	= Map.insert f d (byFuncName ds)
	
	boundVars :: Exp -> [Id]
	boundVars = nub . bv
		where
			bv e@(Case _ as) = concatMap (patVars . fst) as ++ extract bv e
			bv e@(Let bs _) = map fst bs ++ extract bv e
			bv e = extract bv e