{-
Contains the abstract syntax tree for initial parsing.
The parser transforms input file into a list of "Recursion"s defined in this module.
-}

module Tree where

import Data.List
import Data.String
import Data.Monoid
import Data.Char

data Program = Program { dpLabel :: String, lhs :: [Expr], rhs :: Expr }

--data Recursion = Rec { lhs :: Expr, rhs :: Expr }

data Op = Plus | Minus | Times | Div | And | IsEqual deriving Eq

-- comprehension : f(i,j...), (i <- ..., j <- ...), (predicate1, pred2...)
--data Set = Finite [Expr] | Union Set Set | Comprehension Expr [(Expr, Set)] [Expr] | ArraySet Expr deriving (Eq, Show)

data Expr = BinOp Op Expr Expr
			| Max Expr
			| FiniteSet [Expr]
			| Union Expr Expr
			| Comprehension Expr [(Expr, Expr)] Expr
			| DPAccess [Expr]
			| IfWithoutElse [Expr] [Expr]
			| ArrayAccess Expr [Expr]
			| Int Int
			| Var Type String
			| Negate Expr 
			| MyTrue deriving (Eq, Show)

-- Array n refers to type of n-dimensional arrays of integers
data Type = AnInt | Array Int deriving Eq

-- "folds" an Expr by applying a function to leaves and concatenating results from left to right
-- Set comprehension is considered a dead end just like a variable or int
foldExpr :: (Monoid m) => (Expr -> Bool) -> (Expr -> m) -> Expr -> m
foldExpr isLeaf f expr 
	| isLeaf expr = f expr 
	| otherwise = case expr of
		BinOp op l r -> (foldExpr isLeaf f l) `mappend` (foldExpr isLeaf f r)
		Negate e -> foldExpr isLeaf f e
		DPAccess inds -> foldl mappend mempty $ map (foldExpr isLeaf f) inds
		IfWithoutElse conds bodies -> foldr mappend mempty $ map (foldExpr isLeaf f) $ conds ++ bodies
		ArrayAccess name inds -> foldl mappend mempty $ map (foldExpr isLeaf f) (name : inds)
		Max expr -> foldExpr isLeaf f expr
		FiniteSet exprs -> foldr mappend mempty $ map (foldExpr isLeaf f) exprs
		Union l r -> (foldExpr isLeaf f l) `mappend` (foldExpr isLeaf f r)
		-- at this point, expr can only be an Int or a Var or a set comprehension.
		_ -> mempty

{-
foldSet :: (Monoid m) => (Set -> Bool) -> (Set -> m) -> Set -> m	
foldSet isLeaf f (Finite exprs) = foldr mappend mempty $ map (foldExpr isLeaf f) exprs
foldSet isLeaf f (Union l r) = mappend (foldSet isLeaf f l) (foldSet isLeaf f r)
foldSet isLeaf f (Comprehension e pairs conds) = 
-}
	
-- transform "leaves" by applying function f to them
-- in the same vein as fmap of the Functor class, but Expr isn't exactly a Functor.
myfmap :: (Expr -> Bool) -> (Expr -> Expr) -> Expr -> Expr
myfmap isLeaf f expr
	| isLeaf expr = f expr 
	| otherwise = case expr of
		BinOp op l r -> BinOp op (myfmap isLeaf f l) (myfmap isLeaf f r)
		Negate e -> Negate (myfmap isLeaf f e)
		DPAccess inds -> DPAccess $ map (myfmap isLeaf f) inds
		IfWithoutElse conds bodies -> IfWithoutElse (map (myfmap isLeaf f) conds) (map (myfmap isLeaf f) bodies)
		ArrayAccess name inds -> ArrayAccess name $ map (myfmap isLeaf f) inds
		Max set -> Max $ myfmap isLeaf f set
		FiniteSet es -> FiniteSet $ map (myfmap isLeaf f) es
		Union l r -> Union (myfmap isLeaf f l) (myfmap isLeaf f r)
		-- at this point, expr can only be an Int or a Var or a set comprehension; Hence further mapping is not possible.
		somethingElse -> somethingElse

-- Defined to be used as the first argument to foldExpr and myfmap 
isInt (Int _) = True
isInt _ = False

isVar (Var _ _ ) = True
isVar _ = False

isComplexSet (FiniteSet _) = True
isComplexSet (Comprehension _ _ _) = True

isComplex (DPAccess _) = True
isComplex (Max _) = True
isComplex (FiniteSet _) = True
isComplex (Comprehension _ _ _) = True
isComplex (Union _ _) = True
isComplex (IfWithoutElse _ _) = True
isComplex _ = False

--makeExpr :: Rec -> Expr
--makeExpr (Rec _ pairs) = let (conds, bodies) = (map fst pairs, map snd pairs) in
	

-- If an expression is "simple", then it can be translated into a C++ counterpart
-- in a straightforward fashion
isSimple :: Expr -> Bool
isSimple = getAll . isSimple' where	
	-- this works because foldExpr _ _ applied to a variable or integer returns All True
	isSimple' = foldExpr (\x -> True) (\x -> if isComplex x then All False else foldExpr (\x -> True) isSimple' x)
	
isArg (Var ctype name) = if isPrefixOf "arg" name && all isDigit (drop 3 name) then True else False 
isArg _ = False

indices (DPAccess inds) = inds
	
--These exist for debugging purposes only:

{-					
instance Show Recursion where
	show (Rec lhs rhs) = "Rec\n" ++ 
							indent 1 (show lhs) ++ "\n" ++ 
							indent 1 (show rhs)
-}

instance Show Op where
	show Plus = "+"
	show Minus = "-"
	show Times = "*"
	show Div = "/"
	show And = "&&"
	show IsEqual = "=="
	
show2 :: Expr -> String
show2 e = case e of
	BinOp op l r -> show2 l ++ show op ++ show2 r
	Max e -> "max " ++ show2 e
	FiniteSet es -> "{" ++ intercalate ", " (map show2 es) ++ "}"
	Union l r -> show2 l ++ " U " ++ show2 r
	Comprehension f forAlls cond -> "{ " ++ show2 f ++ " | " ++ intercalate ", " (map showPair forAlls) ++ ", " ++ show2 cond ++ " }" where
		showPair (l, r) = show2 l ++ " <- " ++ show2 r
	Int x -> show x
	Var typ name -> name
{-
instance Show Expr where
	show (BinOp op lhs rhs) = show op ++ "\n" ++ 
								indent 1 (show lhs) ++ "\n" ++ 
								indent 1 (show rhs)
	show (Max exprs) = "Max\n" ++ (indent 1 $ 
									intercalate "\n" $ 
										--map show exprs)
										show exprs)
	show (DPAccess exprs) = "DP\n" ++ (indent 1 $ 
											intercalate "\n" $ 
												map show exprs)
	show (ArrayAccess name exprs) = "Array " ++ show name ++ "\n" ++ 
										(indent 1 $ 
											intercalate "\n" $ 
												map show exprs)
	show (IfWithoutElse [c] [b]) = "if " ++ show c ++ " \n" ++ indent 1 (show b) ++ "\n"
	show (IfWithoutElse conds bodies) = show (IfWithoutElse (tail conds) (tail bodies)) ++ "if " ++ 
		show (head conds) ++ " \n" ++ indent 1 (show (head bodies)) ++ "\n" 
	-- show (IfWithoutElse _ _) = "wtf"
	show MyTrue = "true"
	show (Int x) = show x
	show (Var vtype name) = name
	show (Negate e) = "-\n" ++ indent 1 (show e)
-}
	
instance Show Type where
	show AnInt = "int"
	show (Array x) = iterate (\x -> "std::vector<" ++ x ++ ">") "int" !! x 
	
-- The first argument n is the depth of indentation;
-- Every line of the second argument is prepended n * 4 spaces
indent :: Int -> String -> String
indent 0 ss = ss
indent n ss = indent (n-1) $ 
				intercalate "\n" $ 
					map ("    "++) (lines ss)