module Target where

import Data.Char (isDigit)

type Variable = String

data Literal = Str String | Int Integer | Frac Rational | Char Char 
               deriving (Show,Eq)

data Pattern = PVar Variable | WildCard | PCon Variable [Pattern]
               deriving (Show,Eq)

data Target = TLambda [Pattern] Target
            | TApp Target [Target]
            | TInfix Target Variable Target
            | TVar Variable  
            | TBound Int   -- de Bruijn index                   
            | TLit Literal
            | TLet [(Pattern,Target)] Target
            | TIf Target Target Target
            | TCase Target [Alt]

            -- temporarily used constructors
            | TApp1 Target Target
            | TLambda1 Variable Target
            | TLambda0 Target
            | TReturn Target
            | Target `TBind` Target
            deriving (Show,Eq)
                        
type Alt = (Pattern, Target) 

hsName :: String -> String
hsName "%"  = "`mod`" 
hsName "//" = "`div`"  -- from python
hsName n   = n

mkOp :: String -> Target
mkOp name = TLambda1 "x"
                (TReturn (TLambda1 "y"
                     (TReturn (TInfix (TBound 1) (hsName name) (TBound 0)))))

divOp :: String -> Target
divOp name = TLambda1 "x"
                 (TReturn (TLambda1 "y"
                      (TIf (TInfix (TBound 0) "==" (TLit (Int 0)))
                           (TVar "mzero")
                           (TReturn (TInfix (TBound 1) (hsName name) (TBound 0))))))

mkPrimN :: Int -> Variable -> Target
mkPrimN 0 f = TLambda0 (TVar f)
mkPrimN 1 f = TVar f
mkPrimN n f = let m        = TApp (TVar f) (map TBound [n-1,n-2..0])
                  TReturn x = foldr (\ i t -> TReturn (TLambda1 ('x':show i) t)) m [1..n]
                in x

mkPureN n f = let m        = TReturn (TApp (TVar f) (map TBound [n-1,n-2..0]))
                  TReturn x = foldr (\ i t -> TReturn (TLambda1 ('x':show i) t)) m [1..n]
                in x

failFun :: Target
failFun = TLambda0 (TVar "mzero")

commonPrimitives 
              = map (\ n -> (n, mkOp n)) ["**", "*", "+", "-", "++", ":", "==", "/=", "<", "<=", ">=", ">"]
             ++ map (\ n -> (n, divOp n)) ["/", "%", "//"]
             ++ map (\ n -> (n, mkPureN 1 n)) ["not", "null", "head", "tail", "show", "atoi", "atof"]
             ++ map (\ n -> (n, mkPrimN 1 n)) ["get", "writeStr", "write"]
             ++ map (\ n -> (n, mkPrimN 0 n)) ["readChar", "eof"]
             ++ map (\ n -> (n, mkPrimN 2 n)) [{- "try", -} "append", "set"]

commonPrimitivesPlus = commonPrimitives              
                    ++ map (\ n -> (n, mkPrimN 2 n)) ["mplus"]
                    ++ [("fail", failFun)]

removeSugar :: Bool -> Target -> Target
removeSugar b (TApp (TApp f a1) a2)
                             = removeSugar b (TApp f (a1++a2))
removeSugar b (TApp (TVar o) [l, r]) | head o `elem` ":!#$%&*+./<=>?@\\^|-~"
                             = removeSugar b (TInfix l o r)
removeSugar b (TApp1 f a)    = removeSugar b (TApp f [a])
removeSugar b (TLambda1 v t) = TLambda [PVar v] (removeSugar b t)
removeSugar b (TLambda0 t)   = TLambda [WildCard] (removeSugar b t)
removeSugar b (TReturn t)     = TApp (TVar "return") [removeSugar b t]
removeSugar b (m `TBind` k) = if b then TInfix (removeSugar b m) ">>=" (removeSugar b k)
                                    else removeSugar b m `TBind` removeSugar b k 
removeSugar b (TLambda p t)  = TLambda p (removeSugar b t)
removeSugar b (TApp f xs)    = TApp (removeSugar b f) (map (removeSugar b) xs)
removeSugar b (TInfix l o r) = TInfix (removeSugar b l) o (removeSugar b r)
removeSugar b (TLet ds t)    = TLet (removeSugarDecls b ds) (removeSugar b t)
removeSugar b (TIf c t e)    = TIf (removeSugar b c) (removeSugar b t) (removeSugar b e)
removeSugar b (TCase t alts) = TCase (removeSugar b t) (map (\ (p,t) -> (p,removeSugar b t)) alts)
removeSugar _ t              = t

removeSugarDecls :: Bool -> [(t, Target)] -> [(t, Target)]
removeSugarDecls b ds = map (\ (p,t) -> (p,removeSugar b t)) ds

simplify :: Target -> Target
simplify (TReturn a `TBind` k)                             = simplify (TApp1 k a)
simplify (m `TBind` TLambda1 v (TReturn (TBound 0)))       = simplify m
simplify (m `TBind` TLambda [PVar v] (TReturn (TBound 0))) = simplify m
simplify ((m `TBind` TLambda1 v m1) `TBind` k2) =
            simplify (m `TBind` TLambda1 v (m1 `TBind` shift 1 0 k2))
simplify ((m `TBind` TLambda0 m1) `TBind` k2) =
            simplify (m `TBind` TLambda0 (m1 `TBind` k2))
simplify ((m `TBind` TLambda [p] m1) `TBind` k2) =
            let j = length $ fvPat p in  
            simplify (m `TBind` TLambda [p] (m1 `TBind` shift j 0 k2))
simplify (TApp1 (TLambda1 v t) a)     = simplify (subst 0 a t)
simplify (TApp1 (TLambda0 t) a)       = simplify t
simplify (TApp (TLambda1 v t) [a])    = simplify (subst 0 a t)
simplify (TApp (TLambda0 t) [a])      = simplify t
simplify (TApp (TLambda1 v t) (x:xs)) = simplify (TApp (subst 0 x t) xs)
simplify (TApp (TLambda0 t) (x:xs))   = simplify (TApp t xs)

simplify (TLambda1 v t) = TLambda1 v (simplify t)
simplify (TLambda0 t)   = TLambda0 (simplify t)
simplify (TReturn t)     = TReturn (simplify t)
simplify (m `TBind` k) = let m1 = simplify m
                             k1 = simplify k
                           in if m1==m && k1==k
                              then m1 `TBind` k1 
                              else simplify (m1 `TBind` k1)
simplify (TLambda p t)  = TLambda p (simplify t)
simplify (TApp f xs)    = let f1  = simplify f
                              xs1 = map simplify xs
                            in if f1==f && xs1==xs
                               then TApp f1 xs1
                               else simplify (TApp f1 xs1)
simplify (TApp1 f x)    = let f1 = simplify f
                              x1 = simplify x
                            in if f1==f && x1==x
                               then TApp1 f1 x1
                               else simplify (TApp1 f1 x1)
simplify (TInfix l o r) = TInfix (simplify l) o (simplify r)
simplify (TLet ds t)    = TLet (simplifyDecls ds) (simplify t)
simplify (TIf c t e)    = TIf (simplify c) (simplify t) (simplify e)
simplify (TCase t alts) = TCase (simplify t) (map (\ (p,t) -> (p,simplify t)) alts)
simplify t              = t

simplifyDecls :: [(Pattern, Target)] -> [(Pattern, Target)]
simplifyDecls ds = map (\ (p,t) -> (p,simplify t)) ds

fvPat :: Pattern -> [Variable]
fvPat p = fvPatAux p []

fvPatAux :: Pattern -> [Variable] -> [Variable]
fvPatAux (PVar v) acc    = v:acc
fvPatAux WildCard acc    = acc
fvPatAux (PCon _ ps) acc = foldr fvPatAux acc ps

-- Convert occurrences of b to bound index i in a term
abstract :: Int -> String -> Target -> Target
abstract i b (TLambda pats t) = TLambda pats (abstract (i+j) b t) 
    where j = sum $ map (length . fvPat) pats
abstract i b (TApp t u)       = TApp (abstract i b t) (map (abstract i b) u); 
abstract i b (TApp1 t u)      = TApp1 (abstract i b t) (abstract i b u); 
abstract i b (TInfix l o r)   = TInfix (abstract i b l) o (abstract i b r); 
abstract i b (TVar a)         = if a == b then TBound i else TVar a;
abstract i b (TBound j)       = TBound j;
abstract i b (TLet bs t)      = TLet (map (\(p,u)->(p,abstract (i+j) b u)) bs) (abstract (i+j) b t)
   where ps = map fst bs
         j  = sum $ map (length . fvPat) ps
abstract i b (TIf c t e)      = TIf (abstract i b c) (abstract i b t) (abstract i b e)
abstract i b (TCase t as)     = TCase (abstract i b t) 
                                      (map (\(p,u)->(p,abstract (i+length (fvPat p)) b u)) as)
abstract i b (TLambda1 v t)   = TLambda1 v (abstract (i+1) b t)
abstract i b (TLambda0 t)     = TLambda0 (abstract i b t)
abstract i b (TReturn t)       = TReturn (abstract i b t)
abstract i b (l `TBind` r)   = abstract i b l `TBind` abstract i b r                                                              
abstract i b t                = t                                       

-- Abstraction over several free variables
absList :: [String] -> Target -> Target
absList bs t = let n = length bs 
               in foldr (\ (i,b) u -> abstract i b u) t (zip [n-1,n-2 .. 0] bs)

mkLambda :: [Pattern] -> Target -> Target
mkLambda ps t =
    let vs = foldr fvPatAux [] ps
    in  TLambda ps (absList vs t)

mkLambda1 :: Variable -> Target -> Target
mkLambda1 v t = TLambda1 v (absList [v] t)

mkLambda0 :: Target -> Target
mkLambda0 t   = TLambda0 t

mkLet :: [(Pattern, Target)] -> Target -> Target
mkLet bs t    = let ps = map fst bs
                    vs = foldr fvPatAux [] ps
                in TLet (map (\ (p,u) -> (p, absList vs u)) bs)
                        (absList vs t) 

mkAlt :: Pattern -> Target -> (Pattern, Target)
mkAlt p t     = (p, absList (fvPat p) t)

deBruijn :: Target -> Target
deBruijn (TLambda ps t) = mkLambda ps (deBruijn t)
deBruijn (TLambda1 v t) = mkLambda1 v (deBruijn t)
deBruijn (TLambda0 t)   = mkLambda0 (deBruijn t)
deBruijn (TApp t ts)    = TApp (deBruijn t) (map deBruijn ts)
deBruijn (TApp1 t s)    = TApp1 (deBruijn t) (deBruijn s)
deBruijn (TInfix l o r) = TInfix (deBruijn l) o (deBruijn r)
deBruijn (TLet ps t)    = mkLet (map (\ (p,u) -> (p, deBruijn u)) ps) (deBruijn t)
deBruijn (TIf c t e)    = TIf (deBruijn c) (deBruijn t) (deBruijn e)
deBruijn (TCase t as)   = TCase (deBruijn t) (map (\ (p, u) -> mkAlt p (deBruijn u)) as)
deBruijn (TReturn t)     = TReturn (deBruijn t)
deBruijn (m `TBind` k) = deBruijn m `TBind` deBruijn k
deBruijn t              = t

deBruijnDecls :: [(Pattern, Target)] -> [(Pattern, Target)]
deBruijnDecls ds = let (ps, us) = unzip ds
                       vs = foldr fvPatAux [] ps
                   in zip ps (map (\ u -> absList vs (deBruijn u)) us)
               
-- Shift a term's non-local indices by i; d is the depth of abstractions
shift :: Int -> Int -> Target -> Target
shift 0 d u = u
shift i d (TLambda ps t) = TLambda ps (shift i (d+j) t) where j = sum $ map (length . fvPat) ps
shift i d (TLambda1 v t) = TLambda1 v (shift i (d+1) t)
shift i d (TLambda0 t)   = TLambda0 (shift i d t)
shift i d (TApp t u)     = TApp (shift i d t) (map (shift i d) u)
shift i d (TApp1 t u)    = TApp1 (shift i d t) (shift i d u)
shift i d (TInfix l o r) = TInfix (shift i d l) o (shift i d r)
shift i d (TVar a)       = TVar a
shift i d (TBound j)     = if j >= d then TBound (j + i) else TBound j
shift i d (TLet bs t)    = TLet (map (\(p,u) -> (p, shift i (d + j) u)) bs) (shift i (d + j) t)
   where ps = map fst bs
         j  = sum $ map (length . fvPat) ps
shift i d (TIf c t e)    = TIf (shift i d c) (shift i d t) (shift i d e)
shift i d (TCase t as)   = TCase (shift i d t)
                                 (map (\(p,u) -> (p,shift i (d+length (fvPat p)) u)) as)
shift i d (TReturn t)     = TReturn (shift i d t)
shift i d (l `TBind` r) = shift i d l `TBind` shift i d r
shift i d t              = t

-- Substitute u for bound variable i in a term t
subst :: Int -> Target -> Target -> Target
subst i u (TVar a)       = TVar a
subst i u (TBound j) =
        if j<i then TBound j    {- locally bound -}
        else if j == i then shift i 0 u
        else {- j>i -}    TBound (j - 1)  {- non-local to t -}
subst i u (TLambda1 a t) = TLambda1 a (subst (i+1) u t)
subst i u (TLambda ps t) = TLambda ps (subst (i+j) u t) where j = sum $ map (length . fvPat) ps
subst i u (TLambda0 t)   = TLambda0 (subst i u t)
subst i u (TApp t1 t2)   = TApp (subst i u t1) (map (subst i u) t2)
subst i u (TApp1 t1 t2)  = TApp1 (subst i u t1) (subst i u t2)
subst i u (TInfix l o r) = TInfix (subst i u l) o (subst i u r)         
subst i u (TLet bs t)    = TLet (map (\(p,v) -> (p,subst (i+j) u v)) bs) (subst (i+j) u t)
   where ps = map fst bs
         j  = sum $ map (length . fvPat) ps
subst i u (TIf c t e)    = TIf (subst i u c) (subst i u t) (subst i u e)
subst i u (TCase t as)   = TCase (subst i u t)
                                 (map (\(p,v) -> (p, subst (i+length (fvPat p)) u v)) as)
subst i u (TReturn t)     = TReturn (subst i u t)
subst i u (l `TBind` r) = subst i u l `TBind` subst i u r
subst i u t              = t

-- Substitution for free variables (recursive version)
inst :: [(Variable, Target)] -> Target -> Target
inst env (TVar a)       = case lookup a env of 
                            Just b  -> inst env b
                            Nothing -> TVar a
inst env (TBound i)     = TBound i
inst env (TLambda ps t) = TLambda ps (inst env t)
inst env (TLambda1 v t) = TLambda1 v (inst env t)
inst env (TLambda0 t)   = TLambda0 (inst env t)
inst env (TApp t1 t2)   = TApp (inst env t1) (map (inst env) t2)
inst env (TApp1 t1 t2)  = TApp1 (inst env t1) (inst env t2)
inst env (TInfix l o r) = TInfix (inst env l) o (inst env r)
inst env (TLet bs t)    = TLet (map (\(p,u) -> (p,inst env u)) bs) (inst env t)
inst env (TIf c t e)    = TIf (inst env c) (inst env t) (inst env e)
inst env (TCase t as)   = TCase (inst env t) (map (\(p,u) -> (p, inst env u)) as)
inst env (TReturn t)     = TReturn (inst env t)
inst env (l `TBind` r) = inst env l `TBind` inst env r
inst env t              = t


-- Substitution for free variables (non-recursive version)
instOnce :: [(Variable, Target)] -> Target -> Target
instOnce env (TVar a)       = case lookup a env of 
                                Just b  -> b
                                Nothing -> TVar a
instOnce env (TBound i)     = TBound i
instOnce env (TLambda ps t) = TLambda ps (instOnce env t)
instOnce env (TLambda1 v t) = TLambda1 v (instOnce env t)
instOnce env (TLambda0 t)   = TLambda0 (instOnce env t)
instOnce env (TApp t1 t2)   = TApp (instOnce env t1) (map (instOnce env) t2)
instOnce env (TApp1 t1 t2)  = TApp1 (instOnce env t1) (instOnce env t2)
instOnce env (TInfix l o r) = TInfix (instOnce env l) o (instOnce env r)
instOnce env (TLet bs t)    = TLet (instOnceDecls env bs) (instOnce env t)
instOnce env (TIf c t e)    = TIf (instOnce env c) (instOnce env t) (instOnce env e)
instOnce env (TCase t as)   = TCase (instOnce env t) (map (\(p,u) -> (p, instOnce env u)) as)
instOnce env (TReturn t)     = TReturn (instOnce env t)
instOnce env (l `TBind` r) = instOnce env l `TBind` instOnce env r
instOnce env t              = t

instOnceDecls :: [(Variable, Target)] -> [(t, Target)] -> [(t, Target)]
instOnceDecls env bs = map (\(p,u) -> (p,instOnce env u)) bs

-- Free variable in a term
vars :: Target -> [String] -> [String]
vars (TVar a) acc        = a : acc
vars (TBound _) acc      = acc
vars (TLambda ps t) acc  = vars t acc
vars (TLambda1 v t) acc  = vars t acc
vars (TLambda0 t) acc    = vars t acc
vars (TApp t1 t2) acc    = foldr vars acc (t1:t2)
vars (TApp1 t1 t2) acc   = vars t1 (vars t2 acc)
vars (TInfix l o r) acc  = vars l (vars r acc)
vars (TLet bs t) acc     = foldr vars (vars t acc) (map snd bs)
vars (TIf c t e) acc     = vars c (vars t (vars e acc))
vars (TCase t as) acc    = foldr vars acc (t:(map snd as))
vars (TReturn t) acc     = vars t acc
vars (l `TBind` r) acc   = vars l (vars r acc)
vars (TLit _)  acc       = acc

name1 :: String -> String;
name1 str = let  (body, num) = break isDigit str
                 n           = if num =="" then 0 else read num 
            in body ++ show (n+1)

-- Rename variable "a" to avoid clashes with the strings bs.
rename :: [String] -> String -> String
rename bs a = if a `elem ` bs then rename bs (name1 a) else a

renamePat :: [String] -> Pattern -> Pattern
renamePat bs (PVar a)    = PVar (rename bs a)
renamePat bs WildCard    = WildCard   
renamePat bs (PCon c ps) = PCon c (map (renamePat bs) ps)

-- 
stripPatAux :: [Pattern] -> [String] -> [Pattern]
stripPatAux pat vs = map (renamePat vs) pat 
                  
stripPat :: [Pattern] -> Target -> ([Pattern],Target)
stripPat pat t = let pat1 = stripPatAux pat (vars t [])
                     vs   = foldr fvPatAux [] pat1
                 in (pat1, stripVars vs t)

stripVars :: Foldable t => t Variable -> Target -> Target
stripVars vs t = foldr (\ v t -> subst 0 (TVar v) t) t vs

data Assoc = LeftAssoc | RightAssoc | NonAssoc deriving Eq

lookupAssoc :: String -> (Integer, Assoc)
lookupAssoc o = case lookup o assocList of
                  Just (p,f) -> (p,f)
                  Nothing    -> (9,LeftAssoc)  -- Any operator lacking a fixity declaration is assumed to be infixl 9.
              where 
                 assocList =  map (\o -> (o, (0, RightAssoc))) ["$", "$!", "`seq`"] 
                           ++ [("=<<", (1, RightAssoc))]
                           ++ map (\o -> (o, (1, LeftAssoc))) [">>", ">>="] 
                           ++ [("||", (2, RightAssoc)), ("&&", (3, RightAssoc))] 
                           ++ map (\o -> (o, (4, NonAssoc))) ["==", "/=", "<", "<=", ">=", ">"]
                           ++ map (\o -> (o, (5, RightAssoc))) [":", "++"]  
                           ++ map (\o -> (o, (6, LeftAssoc))) ["+", "-"]   
                           ++ map (\o -> (o, (7, LeftAssoc))) ["*", "/", "`quot`", "`rem`", "`div`", "`mod`"] 
                           ++ map (\o -> (o, (8, RightAssoc))) ["^", "^^", "**"]
                           ++ [(".", (9, RightAssoc))]

dangerousLeft :: (Integer, Assoc) -> Target -> Bool
dangerousLeft _ (TLambda _ _)  = True
dangerousLeft _ (TApp _ _)     = False
dangerousLeft _ (TVar _)       = False
dangerousLeft _ (TLit _)       = False
dangerousLeft _ (TLet _ _)     = True
dangerousLeft _ (TIf _ _ _)    = True
dangerousLeft _ (TCase _ _)    = True
dangerousLeft (p0,f0) (TInfix l o r) = let (p1,f1) = lookupAssoc o
                                       in if p1 > p0 || p1 == p0 && f0 == RightAssoc && f1 == RightAssoc
                                          then dangerousLeft (p1,f1) r
                                          else False