module Optimus.Simplify where -- (simplify) where
	import Flite.Syntax
	import Flite.Traversals hiding (freshen)
	import Flite.Descend
	import Flite.Fresh
	import Optimus.Util
	import Data.List
	import Data.Maybe
	import Control.Monad
	import Optimus.Trace
	
	-- Issues are going to be with additional Apps in places
	
	-- |Not traced as too common and simply.
	stripApp :: Exp -> Exp
	stripApp e@(App (Fun _) []) = e
	stripApp (App x []) = stripApp x
	stripApp e = descend stripApp e
	
	-- |Not traced as too common and simply.
	appOfApp :: Exp -> Exp
	appOfApp (App (App x xs) ys) = appOfApp $ App x (xs ++ ys)
	appOfApp e = descend appOfApp e
			
	caseOfCons :: Exp -> Fresh Exp
	caseOfCons (Case (App (Con c) xs) ps) | isJust y' = fromJust y' >>= (caseOfCons . t_S "caseOfCons")
		where
			y' = listToMaybe [ mapM (const fresh) vs' >>=
							   \ws -> return $ Let (zip ws xs) (substMany y $ zipWith (\w v -> (Var w, v)) ws vs')
							| ( App (Con c') vs , y ) <- ps, c == c', length xs == length vs, let vs' = concatMap patVars vs ]
	caseOfCons (Case (Con c) ps) = caseOfCons $ Case (App (Con c) []) ps
	caseOfCons e = descendM caseOfCons e
			
	appToCase :: Exp -> Fresh Exp
	appToCase (App e@(Case _ _) zs) = do
		Case x ps <- freshenMany e (concatMap freeVars zs)
		appToCase $ t_S "appToCase" $ Case x [ (p, App y zs) | (p, y) <- ps ]
	appToCase e = descendM appToCase e
			
	appToLet :: Exp -> Fresh Exp
	appToLet (App e@(Let _ _) zs) = do
		Let bs y <- freshenMany e (concatMap freeVars zs)
		appToLet $ t_S "appToLet" $ Let bs (App y zs)
	appToLet e = descendM appToLet e
			
	caseOfLet :: Exp -> Fresh Exp
	caseOfLet (Case e@(Let _ _) as) = do
		Let bs y <- freshenMany e $ concat [ freeVarsExcept (patVars p) x | (p, x) <- as ]
		caseOfLet $ t_S "caseOfLet" $ Let bs (Case y as)
	caseOfLet e = descendM caseOfLet e
			
	caseOfCase :: Exp -> Fresh Exp
	caseOfCase (Case e@(Case _ _) as') = do
		Case x as <- freshenMany e $ concat [ freeVarsExcept (freeVars p) y | (p, y) <- as' ]
		caseOfCase $ t_S "caseOfCase" $ Case x [ (p, Case y as') | (p, y) <- as]
	caseOfCase e = descendM caseOfCase e
			
	substituteVar :: Exp -> Fresh Exp
	substituteVar e@(Case (Var v) as_) | v `elem` (freeVars $ Case Bottom as_) = do
		Case _ as <- freshen e v
		as' <- sequence [ liftM ((,) p) $ freshenMD y $ patVars p | (p, y) <- as ]
		--freshenMany e $ v : concat [ patVars p | Case _ as' <- extract universe e, (p, _) <- as' ]
		substituteVar $ t_S "substituteVar" $ Case (Var v) [ (p, subst p v y) | (p, y) <- as' ]
	substituteVar e = descendM substituteVar e
			
	letInCase :: Exp -> Fresh Exp
	letInCase (Let bs e@(Case y _)) | (not . null) safe_bs = do
			Case _ as <- freshenMany e (v:freeVars x)
			letInCase $ t_S "letInCase" $ Let rem_bs (Case y [ (p, Let [(v, x)] z) | (p, z) <- as ])
		where
			yFree = freeVars y
			bsFrees = [ (v, freeVars x) | (v, x) <- bs ]
			safe_bs = [ (v, x) | (v, x) <- bs, v `notElem` (yFree ++ concat [ ws | (w, ws) <- bsFrees, v /= w ]) ]
			(v, x) = head safe_bs
			rem_bs = [ (v', x') | (v', x') <- bs, v /= v' ]
	letInCase e = descendM letInCase e
			
	inlineLet :: Exp -> Fresh Exp
	inlineLet (Let bs y) | (not . null) safe_bs = do
			y' <- freshenMD y $ freeVars x_i
			inlineLet $ t_S "inlineLet" $ subst x_i v_i (Let rem_bs y')
		where
			(_, xs) = unzip bs
			(v_i, x_i) = head safe_bs
			rem_bs = [ b | b@(v_j, _) <- bs, v_j /= v_i ]
			safe_bs = 	[ b |
						  b@(v_j, x_j) <- bs
						, v_j `notElem` freeVars x_j
						, (sum . map (varRefs v_j)) (y:xs) <= 1 || isOk x_j ]
			isOk (Var _) = True
			isOk (Con _) = True
			isOk (Int _) = True
			isOk (Fun _) = True
			isOk (App e []) = isOk e
			isOk _ = False
	inlineLet e = descendM inlineLet e
			
	splitLet :: Exp -> Fresh Exp
	splitLet (Let bs y) | (not . null) safe_bs
		= do
			new_bs <- mapM build args
			splitLet $ t_S "splitLet" $ subst (App (Con c) (map Var (fst . unzip $ new_bs))) v_i (Let (rem_bs ++ new_bs) y)
		where
			(v_i, App (Con c) args) = head safe_bs
			rem_bs = [ b | b@(v_j, _) <- bs, v_j /= v_i ]
			safe_bs = [ b | b@(v_j, x_j@(App (Con _) _)) <- bs, v_j `notElem` freeVars x_j ]
			build x_j = fresh >>= \vF -> return (vF, x_j)
	splitLet e = descendM splitLet e
	
	removeLet :: Exp -> Exp
	removeLet (Let [] x) = descend removeLet x
	removeLet e = descend removeLet e
			
	letInLet :: Exp -> Fresh Exp
	letInLet (Let bs1 e@(Let _ _)) = do
			Let bs2 x <- freshenMany e (concatMap (freeVars . snd) bs1 ++ map fst bs1)
			letInLet $ t_S "letInLet" $ Let (bs1 ++ bs2) x
	letInLet e = descendM letInLet e
			
	simplify :: Exp -> Fresh Exp
	simplify p = do
			p' <- (return . stripApp . appOfApp . removeLet) =<< appToCase =<< appToLet =<< caseOfLet =<< caseOfCons
				=<< caseOfCase =<< substituteVar =<< letInCase
				=<< inlineLet =<< letInLet =<< ( splitLet . appOfApp
					. removeLet . stripApp ) p
			if p == p' then return p else simplify p'
	
	simplifyInlines :: Exp -> Fresh Exp
	simplifyInlines e = do
		e' <- (return . removeLet <=< inlineLet <=< return . stripApp) e
		if e == e' then return e else simplifyInlines e'
		
	simplifyProg :: Prog -> Fresh Prog
	simplifyProg p = sequence [ return . Func f a =<< (\x -> t_S' (show x) (simplify x)) =<< simplifyInlines (t_S' ("Simplifying " ++ f) r) | Func f a r <- p]