{-# LANGUAGE ViewPatterns, PatternGuards #-}
{-# OPTIONS_GHC -fno-warn-overlapping-patterns #-}

module Desugar(
    records, untyped, irrefutable, core, core2, singleCase
    ) where

import HSE
import Data.Generics.PlateData
import Data.Maybe
import Data.List hiding (find)
import Control.Monad.State


find x xs = fromMaybe (error $ "Couldn't find: " ++ show x) $ lookup x xs


records :: [Decl] -> [Decl]
records xs = transformBi fPat $ transformBi fExp xs
    where
        recs = [(x, concatMap fst ys) | RecDecl x ys <- universeBi xs]
        typs = [(x, replicate (length ys) (Ident "")) | ConDecl x ys <- universeBi xs]

        fExp o@(RecConstr (UnQual name) xs) = Paren $ apps (Con $ UnQual name) [Paren $ find l lbls | l <- find name recs]
            where lbls = [(n,x) | FieldUpdate (UnQual n) x <- reverse xs]
        fExp x = x

        fPat (PRec (UnQual name) [PFieldWildcard]) = PParen $ PApp (UnQual name) $ map PVar $ find name recs
        fPat (PRec (UnQual name) []) = PParen $ PApp (UnQual name) $ replicate (length $ find name $ recs++typs) PWildCard
        fPat x = x


untyped :: [Decl] -> [Decl]
untyped = map (descendBi untyped) . filter (not . isTypeSig)
    where isTypeSig TypeSig{} = True
          isTypeSig _ = False


irrefutable :: [Decl] -> [Decl]
irrefutable = concatMap f . descend irrefutable
    where
        f (PatBind sl (PTuple [PVar (Ident x), PVar (Ident y)]) a b c) =
            [PatBind sl (pvar xy) a b c
            ,PatBind sl (pvar x) Nothing (UnGuardedRhs $ App (var "fst") (var xy)) (BDecls [])
            ,PatBind sl (pvar y) Nothing (UnGuardedRhs $ App (var "snd") (var xy)) (BDecls [])]
            where xy = x++"_"++y
        f x = [x]


singleCase :: [String] -> [Decl] -> [Decl]
singleCase single = transformBi f
    where
        f (Case (Var (UnQual (Ident "otherwise"))) (Alt _ (PApp (UnQual (Ident "True")) []) (UnGuardedAlt x) (BDecls []):_)) = x
        f (Case on as@(Alt _ (PApp (UnQual (Ident x)) _) _ _:_)) | x `elem` single = Case on [head as]
        f x = x


core :: [Decl] -> [Decl]
core = id {- transformBi simpleCase . transformBi removeGuards . lambdaLift . transformBi flatMatches . transformBi removeWhere
    where
        -- afterwards, no where
        removeWhere (Match a b c d bod (BDecls whr)) | whr /= [] = Match a b c d (f bod) (BDecls [])
            where f (UnGuardedRhs x) = UnGuardedRhs $ Let (BDecls whr) x
                  f (GuardedRhss xs) = GuardedRhss [GuardedRhs a b $ Let (BDecls whr) c | GuardedRhs a b c <- xs]
        removeWhere x = x

        -- afterwards, no multiple equations for one function
        flatMatches (FunBind xs@(Match a b c d e f:_)) | length xs > 1 || any isPAsPat (universeBi c) =
                FunBind [Match a b (map PVar ps) Nothing (UnGuardedRhs bod) (BDecls [])]
            where
                ps = map (Ident . (++) "p_" . show) [1..length c]
                bod = Case (tuple $ map (Var . UnQual) ps) [Alt sl (ptuple p) (f b) d | Match _ _ p _ b d <- xs]
                f (UnGuardedRhs x) = UnGuardedAlt x
                f (GuardedRhss xs) = GuardedAlts $ map g xs
                g (GuardedRhs x y z) = GuardedAlt x y z
        flatMatches x = x

        removeGuards (Alt s p (GuardedAlts xs) d:rest) | length xs > 1 = removeGuards $ [Alt s p (GuardedAlts [x]) d | x <- xs] ++ rest
        removeGuards (Alt s p (GuardedAlts [GuardedAlt _ tsts e]) d:rest) =
            [Alt s (PAsPat (Ident "u_") p) (UnGuardedAlt
                (Case (tuple as) $ removeGuards
                    [Alt s (ptuple bs) (UnGuardedAlt e) d] ++ removeGuards
                    [Alt s PWildCard (UnGuardedAlt $ Case (var "u_") rest) d | rest /= []])) d]
            where (as,bs) = unzip $ map f tsts
                  f (Qualifier x) = (x, PApp (UnQual $ Ident "True") [])
                  f (Generator _ p x) = (x, p)
        removeGuards x = x


-- afterwards no PatBind or Lambda, only FunBind
lambdaLift :: [Decl] -> [Decl]
lambdaLift = concatMap f . transformBi add
    where
        f (PatBind a (PVar b) c (UnGuardedRhs (Lambda x y z)) d) | not $ any isLambda $ universe z =
            [FunBind [Match a b y c (UnGuardedRhs z) d]]
        f (PatBind a (PVar b) c (UnGuardedRhs z) d) | not $ any isLambda $ universe z =
            [FunBind [Match a b [] c (UnGuardedRhs z) d]]
        f x@(PatBind _ (PVar (Ident name)) _ _ _) = pat $ uniques (name++"_") $ descendBiM (push []) $ root x
        f x = [x]

        -- convert a PatBind in to several FunBind's
        pat x = map f $ transformBi dropLam $ x : filter isPatLambda (universeBi x)
            where dropLam (Let (BDecls x) y) = let ds = filter (not . isPatLambda) x in if null ds then y else Let (BDecls ds) y
                  dropLam x = x
                  f (PatBind a (PVar b) c (UnGuardedRhs (Lambda x y z)) d) = FunBind [Match a b y c (UnGuardedRhs z) d]
                  f x = x

        -- include all variables at lambdas
        push :: [String] -> Exp -> Unique Exp
        push seen o@(Let (BDecls xs) y) = descendM (push (nub $ seen++concat [pats v | PatBind _ v _ _ _ <- xs])) o
        push seen o@(Lambda a xs y) = do
            v <- fresh
            let used = [x | Var (UnQual (Ident x)) <- universe y]
                next = concatMap pats xs `intersect` used
                now = (seen \\ next) `intersect` used
            y <- push (nub $ now++next) y
            xs <- return $ map (PVar . Ident) now ++ xs
            o <- return $ Lambda a xs y
            return $ Let
                (BDecls [PatBind sl (PVar $ Ident v) Nothing (UnGuardedRhs o) (BDecls [])])
                (apps (var v) (map var now))
        push seen o = descendM (push seen) o

        -- introduce a root
        root (PatBind a b@(PVar (Ident name)) c (UnGuardedRhs d) e) = PatBind a b c (UnGuardedRhs bod) e
            where bod = Let (BDecls [PatBind sl (PVar $ Ident $ name ++ "_root_") Nothing (UnGuardedRhs d) (BDecls [])]) (var $ name ++ "_root_")

        -- introduce lambdas
        add (FunBind [Match a b c d (UnGuardedRhs e) f]) = PatBind a (PVar b) d (UnGuardedRhs $ Lambda a c e) f
        add x = x


simpleCase :: Exp -> Exp
simpleCase (Case on alts) | any (\(Alt _ p _ _) -> isPAsPat $ fromPParen p) alts =
    case fromParen on of
        Var v -> f on
        _ -> let1 "a_" on $ f $ var "a_"
    where f v = Case v $ map (g v) alts
          g v (Alt a (fromPParen -> PAsPat (Ident n) x) (UnGuardedAlt c) d) = Alt a x (UnGuardedAlt $ let1 n on c) d
          g v x = x
simpleCase x = x




isLambda Lambda{} = True; isLambda _ = False
isPatLambda (PatBind a b c (UnGuardedRhs d) e) = isLambda d; isPatLambda _ = False
isPAsPat PAsPat{} = True; isPAsPat _ = False


lets xy z = Let (BDecls [PatBind sl (PVar $ Ident x) Nothing (UnGuardedRhs y) (BDecls []) | (x,y) <- xy]) z


type Unique a = State [String] a

uniques :: String -> (Unique a) -> a
uniques s o = evalState o $ map ((++) s . show) [1..]

fresh :: Unique String
fresh = do
    s:ss <- get
    put ss
    return s


-}

{-

core3 :: [Decl] -> [Decl]
core3 = concatMap (uniques "a_" . fDecl [])


fDecl :: [String] -> Decl -> Unique [Decl]
fDecl seen (PatBind a (PVar b) c d e) = fDecl seen $ FunBind [Match a b [] c d e]
fDecl seen (FunBind xs@(Match _ nam ps _ _ _:_)) = do
    vs <- replicateM (length ps) fresh
    let f (Match _ _ ps _ bod whr) = Alt sl (ptuple ps) (g bod) whr
        g (UnGuardedRhs x) = UnGuardedAlt x
        g (GuardedRhss xs) = GuardedAlts [GuardedAlt x y z | GuardedRhs x y z <- xs]
    bod <- fExp vs $ Case (tuple $ map var vs) $ map f xs
    return [FunBind [Match sl nam (map pvar vs) Nothing (UnGuardedRhs bod) (BDecls [])]]
fDecl seen x = return [x]


fExp :: [String] -> Exp -> Unique Exp

-- case
-- if the constructor matches the first pattern, take it
fExp seen (Case on (Alt _ pat (UnGuardedAlt x) decs:rest)) | Just bind <- patMatch pat on = fExp seen $ lets bind x
fExp seen x = return x



patMatch :: Pat -> Exp -> Maybe [(String, Exp)]
patMatch (PTuple xs) (Tuple ys) = do guard (length xs == length ys) ; concatZipWithM patMatch xs ys
patMatch (PVar (Ident x)) y@(Var _) = Just [(x,y)]
patMatch _ _ = Nothing


patMatch1 :: Pat -> Exp -> Maybe (String,Exp)
patMatch1 = undefined


concatZipWithM f xs ys = fmap concat $ sequence $ zipWith f xs ys

-}





lambdaLift :: Decl -> [Decl]
lambdaLift (PatBind _ (PVar (Ident name)) _ (UnGuardedRhs bod) _) = transformBi remLams $
        PatBind sl (PVar (Ident name)) Nothing (UnGuardedRhs bod2) (BDecls []) :
        filter isPatLambda (universeBi bod2)
    where
        bod2 = descendBi (addVars []) bod

        remLams (Let (BDecls xs) y) = if null xs2 then y else Let (BDecls xs2) y
            where xs2 = filter (not . isPatLambda) xs
        remLams x = x

        addVars seen (Let (BDecls xs) y) = g (concat ns) $ lets ys $ addVars seen2 y
            where
                xs2 = [(a,b) | PatBind _ (PVar (Ident a)) _ (UnGuardedRhs b) _ <- xs]
                seen2 = nub $ seen ++ map fst (filter (not . isLambda . snd) xs2)
                (ns,ys) = unzip $ map f xs2

                f (n, Lambda sl ps bod) = ([(n,Paren $ apps (var n2) (map var now))], (n2, addVars [] $ Lambda sl (map pvar now ++ ps) bod))
                    where now = delete "fail_" seen
                          n2 = name++"_"++n
                f (n, x) = ([], (n, addVars seen2 x))

                g ((x,y):xs) q = subst x y $ g xs q
                g [] q = q

        addVars seen (Lambda sl ps x) = Lambda sl ps $ addVars (nub $ seen ++ concatMap pats ps) x
        addVars seen (Case on alts) = Case (addVars seen on) [alt ps $ addVars (nub $ seen ++ pats ps) x | Alt _ ps (UnGuardedAlt x) _ <- alts]
        addVars seen x = descend (addVars seen) x


{-
lambdaLift x@(PatBind _ (PVar (Ident name)) _ _ _) = pat $ uniques (name++"_") $ descendBiM (push []) $ root x
    where
        -- convert a PatBind in to several FunBind's
        pat x = map f $ transformBi dropLam $ x : filter isPatLambda (universeBi x)
            where dropLam (Let (BDecls x) y) = let ds = filter (not . isPatLambda) x in if null ds then y else Let (BDecls ds) y
                  dropLam x = x
                  f (PatBind a (PVar b) c (UnGuardedRhs (Lambda x y z)) d) = FunBind [Match a b y c (UnGuardedRhs z) d]
                  f x = x

        -- include all variables at lambdas
        push :: [String] -> Exp -> Unique Exp
        push seen o@(Let (BDecls xs) y) = descendM (push (nub $ seen++concat [pats v | PatBind _ v _ _ _ <- xs])) o
        push seen o@(Lambda a xs y) = do
            v <- fresh
            let used = [x | Var (UnQual (Ident x)) <- universe y]
                next = concatMap pats xs `intersect` used
                now = (seen \\ next) `intersect` used
            y <- push (nub $ now++next) y
            xs <- return $ map (PVar . Ident) now ++ xs
            o <- return $ Lambda a xs y
            return $ Let
                (BDecls [PatBind sl (PVar $ Ident v) Nothing (UnGuardedRhs o) (BDecls [])])
                (apps (var v) (map var now))
        push seen o = descendM (push seen) o

        -- introduce a root
        root (PatBind a b@(PVar (Ident name)) c (UnGuardedRhs d) e) = PatBind a b c (UnGuardedRhs bod) e
            where bod = Let (BDecls [PatBind sl (PVar $ Ident $ name ++ "_root_") Nothing (UnGuardedRhs d) (BDecls [])]) (var $ name ++ "_root_")
-}


core2 :: [Decl] -> [Decl]
core2 = concatMap f . transformBi qname . transformBi name
    where
        -- transform
        f (PatBind a (PVar name) b c d) =
            lambdaLift $ PatBind a (PVar name) Nothing (UnGuardedRhs $ run $ fExp bod) (BDecls [])
            where bod = Let (BDecls [PatBind a (pvar "root") b c d]) (var "root")
        f (FunBind ms@(Match a name b c d e : _)) =
            lambdaLift $ PatBind a (PVar name) Nothing (UnGuardedRhs $ run $ fExp bod) (BDecls [])
            where bod = Let (BDecls [FunBind [Match a (Ident "root") b c d e | Match a _ b c d e <- ms]]) (var "root")
        f x = [x]

        -- pre simplify
        qname (Qual _ x) = UnQual x
        qname (Special x) = UnQual $ Ident $ "("++prettyPrint x++")"
        qname x = x
        name (Symbol x) = Ident $ "("++x++")"
        name x = x



fExp (Let (BDecls whr) x) = fDecls whr >> fExp x
fExp x@(Var _) = ret x
fExp x@(Con _) = ret x
fExp x@(Lit _) = ret x
fExp (Tuple xs) = fExp $ apps (con $ "(" ++ replicate (length xs - 1) ',' ++ ")") xs
fExp (Paren x) = fExp x
fExp (InfixApp x y z) = fExp $ App (App (opExp y) x) z
fExp (LeftSection x op) = fExp $ App (opExp op) x
fExp (RightSection op x) = fExp $ App (App (var "flip") (opExp op)) x
fExp (App x y) = app (fExp x) (fExp y)
fExp (Case on alts) = do
    v <- share $ fExp on
    patFail
    branch $ map (fAlt v) alts
fExp (If x y z) = branch [match (fExp x) ("True",0) >> fExp y, fExp z]
fExp (List []) = ret $ con "[]"
fExp (List (x:xs)) = fExp $ App (App (var "(:)") x) (List xs)
fExp (Lambda _ ps x) = do
    v <- share $ do
        vs <- lam (length ps)
        fPats vs ps
        fExp x
    ret $ var v
fExp x = err ("fExp",x)



fDecls = binds . map fDecl

fDecl (PatBind _ (PVar (Ident name)) _ (UnGuardedRhs bod) (BDecls whr)) = (name, fDecls whr >> fExp bod)
fDecl (FunBind xs@(Match _ (Ident name) ps _ _ _ : _)) = (name, do vs <- lam $ length ps ; patFail ; branch $ map (fMatch vs) xs)
fDecl x = err ("fDecl",x)

fMatch vs (Match _ _ ps _ rhs (BDecls whr)) = fPats vs ps >> fDecls whr >> fRhs rhs

fPats vs xs | length vs == length xs = sequence_ $ zipWith fPat vs xs

fPat v (PVar (Ident x)) = bind x (ret $ var v)
fPat v (PParen x) = fPat v x
fPat v (PAsPat (Ident x) y) = bind x (ret $ var v) >> fPat v y
fPat v (PLit x) = matchLit (ret $ var v) x
fPat v (PWildCard) = return ()
fPat v (PTuple xs) = fPat v $ PApp (UnQual $ Ident $ "(" ++ replicate (length xs - 1) ',' ++ ")") xs
fPat v (PList []) = fPat v $ PApp (UnQual $ Ident "[]") []
fPat v (PList (x:xs)) = fPat v $ PApp (UnQual $ Ident "(:)") [x, PList xs]
fPat v (PInfixApp x y z) = fPat v $ PApp y [x,z]
fPat v (PApp (UnQual (Ident x)) xs) = do
    vs <- match (ret $ var v) (x,length xs) 
    fPats vs xs
fPat v x = err ("fPat",v,x)


fRhs (UnGuardedRhs x) = fExp x
fRhs (GuardedRhss xs) = branch $ map fGRhs xs

fRhs x = err ("fRhs",x)

fAlt v (Alt _ p bod (BDecls whr)) = fPat v p >> fDecls whr >> fGAlts bod

fGAlts (UnGuardedAlt x) = fExp x
fGAlts (GuardedAlts xs) = branch $ map fGAlt xs

fGAlt (GuardedAlt _ xs x) = mapM_ fStmt xs >> fExp x
fGRhs (GuardedRhs _ xs x) = mapM_ fStmt xs >> fExp x


fStmt (Qualifier x) = match (fExp x) ("True",0) >> return ()
fStmt (Generator _ p x) = do
    v <- share $ fExp x
    fPat v p
fStmt x = err ("fStmt",x)




share :: M Exp -> M String
share x = do
    x <- go x
    case fromParen x of
        Var (UnQual (Ident v)) -> return v
        _ -> do
            v <- fresh
            bind v (ret x)
            return v

bind :: String -> M Exp -> M ()
bind name x = do
    x <- go x
    addExp $ let1Simplify name $ fromParen x

binds :: [(String, M Exp)] -> M ()
binds xs = do
    let (vs,bs) = unzip xs
    bs <- mapM go bs
    addExp $ lets $ zip vs $ map fromParen bs

patFail :: M ()
patFail = do
    fl <- freshFail
    bind fl (ret $ var "patternMatchFail")

branch :: [M Exp] -> M Exp
branch [x] = x
branch (x:xs) = do
    xs <- go $ branch xs
    fl <- freshFail
    bind fl (return xs)
    x

match :: M Exp -> (String,Int) -> M [String]
match x (s,n) = do
    x <- go x
    vs <- replicateM n fresh
    fl <- getFail
    addExp $ \bod -> Case x [alt (PApp (UnQual $ Ident s) (map pvar vs)) bod, alt PWildCard fl]
    return vs

matchLit :: M Exp -> Literal -> M ()
matchLit x lit = do
    x <- go x
    fl <- getFail
    addExp $ \bod -> Case x [alt (PLit lit) bod, alt PWildCard fl]

lam :: Int -> M [String]
lam n = do
    v <- fresh
    vs <- replicateM n fresh
    addExp $ lams vs
    return vs


app :: M Exp -> M Exp -> M Exp
app x y = do
    x <- go x
    y <- go y
    ret $ App x y




type M a = State (Exp -> Exp, Int, Int) a

run :: M Exp -> Exp
run x = evalState x (id, 0, 1)


go :: M Exp -> M Exp
go x = do
    (k,s,i) <- get
    put (id,s,i)
    res <- x
    (k2,s2,i2) <- get
    put (k,s,i2)
    return res

addExp :: (Exp -> Exp) -> M ()
addExp x = modify $ \(k,s,i) -> (k . paren . x, s, i)

freshFail :: M String
freshFail = do
    (k,s,i) <- get
    put (k,i,i+1)
    return $ "f_" ++ show i

getFail :: M Exp
getFail = do
    (k,s,i) <- get
    return $ var $ "f_" ++ show s

fresh :: M String
fresh = do
    (k,s,i) <- get
    put (k,s,i+1)
    return $ "v_" ++ show i

ret :: Exp -> M Exp
ret x = do
    (k,s,i) <- get
    return $ k $ paren x

err x = error $ show x




let1 x y z = Let (BDecls [PatBind sl (pvar x) Nothing (UnGuardedRhs y) (BDecls [])]) z

lets [] z = z
lets xy z = Let (BDecls [PatBind sl (pvar x) Nothing (UnGuardedRhs y) (BDecls []) | (x,y) <- xy]) z

let1Simplify name bind@(fromParen -> Var _) = subst name bind
let1Simplify name bind = let1 name bind

subst name bind = f
    where
        f (App x y) = App (f x) (f y)
        f (Paren x) = Paren $ f x
        f (Var (UnQual (Ident x))) = if x == name then bind else var x
        f (Lit x) = Lit x
        f (Con x) = Con x
        f (Case on alts) = Case (f on) $ map g alts
        f (Let (BDecls bind) x) = if name `elem` map h bind then Let (BDecls bind) x else
            Let (BDecls [PatBind a b c (UnGuardedRhs $ f d) e | PatBind a b c (UnGuardedRhs d) e <- bind]) (f x)
        f (Lambda a ps x) = Lambda a ps $ if name `elem` concatMap pats ps then x else f x
        f x = err ("let1Simplify", x)

        g a@(Alt sl pat (UnGuardedAlt bod) (BDecls [])) = if name `elem` pats pat then a else Alt sl pat (UnGuardedAlt $ f bod) (BDecls [])
        h (PatBind _ (PVar (Ident x)) _ (UnGuardedRhs _) (BDecls [])) = x


lams xs y = Lambda sl (map pvar xs) y

alt p x = Alt sl p (UnGuardedAlt x) (BDecls [])

opExp (QConOp x) = Con x
opExp (QVarOp x) = Var x


paren (Paren x) = Paren x
paren (Var x) = Var x
paren x = Paren x

pats = concatMap f . universe
    where f (PVar (Ident v)) = [v]
          f (PAsPat (Ident v) _) = [v]
          f _ = []

isLambda Lambda{} = True; isLambda _ = False
isPatLambda (PatBind a b c (UnGuardedRhs d) e) = isLambda d; isPatLambda _ = False
namePatBind (PatBind _ (PVar (Ident x)) _ _ _) = x