module FunctorOf (
   getCtx
 , functorOfInst
 , getSeed
 , hsPat2Exp
 , typeable
 , observable
 ) where

import Data.Map hiding (map) 
import Data.List
import Data.Maybe
import Language.Haskell.Exts.Syntax as Exts

import Control.Monad.State
import Data.Char (isUpper)
import Data.Generics hiding (Unit,(:*:),Inl,Inr)

import Language.Haskell.Exts.Parser
import Language.Haskell.Exts.Pretty

import Language.Pointwise.Parser (hsPat2Exp)
import Language.Pointwise.Syntax as Pointwise hiding (swap)

-- Associates a constructor to definition

type Ctx = [(String,Term)]

getCtx :: Module -> Ctx
getCtx (Module _ _ _ _ _ _ decls)
  = concat $ catMaybes $ map getCtxDecl decls

getCtxDecl :: Decl -> Maybe Ctx
getCtxDecl (DataDecl _ _ _ _ _ cons _) = getCtxCons cons In
getCtxDecl _ = fail "Only applies to data declarations"

getCtxCons :: [QualConDecl] -> (Term -> Term) -> Maybe Ctx
getCtxCons [QualConDecl _ _ _ (ConDecl (Ident name) l)] f = 
    return [(name, buildArgs f (length l))]
getCtxCons ((QualConDecl _ _ _ (ConDecl (Ident name) l)):t) f =
    do t' <- getCtxCons t (f . Inr)
       return ((name, buildArgs (f . Inl) (length l)):t')
getCtxCons _ _ = fail "Case not covered"

buildArgs :: (Term -> Term) -> Int -> Term
buildArgs f 0 = f Unit
buildArgs f 1 = Lam "x1" (f (Pointwise.Var "x1"))
buildArgs f n = let v = "x"++(show n)
		    in Lam v (buildArgs (f . (\t -> (Pointwise.Var v :&: t))) (n-1))

{- Calculation of the instances of FunctorOf, when possible. -}

type St = StateT ((String,Int),[Type]) Maybe

conDecl :: QualConDecl -> ConDecl
conDecl (QualConDecl _ _ _ con) = con

functorOfInst :: Bool -> Module -> Module
functorOfInst ob (Module a b c d e i decls)
  = let seed = "v" --getSeed decls
        newDecls = concat $ catMaybes $
                   map (\x -> evalStateT (getInstances ob x) ((seed,0),[])) decls
    in Module a b c d e i (newDecls)

addConst t = modify (\(s,l) -> (s,t:l))

g :: Type -> QualConDecl -> St Type
g arg = gCon arg . conDecl

gCon :: Type -> ConDecl -> St Type
gCon arg (ConDecl _ []) = return constNil
gCon arg (ConDecl _ lBangType) = mapM (h arg) lBangType >>= return . foldr1 timesType
gCon arg (RecDecl hsName l ) = gCon arg (ConDecl hsName (concatMap unRec l))
-- fail "data types records not yet supported"

h :: Type -> BangType -> St Type
h arg (BangedTy   hsType ) = ii arg hsType
h arg (UnBangedTy hsType ) = ii arg hsType
-- what's the diference between Banged and UnBanged?

i :: Type -> Type -> St Type
i arg (TyFun hsType1 hsType2) = fail "TyFun not yet supported"
i arg (TyTuple _ lType)       = mapM (ii arg) lType >>= return . foldr1 timesType
i arg (TyApp hsType1 hsType2) = i arg hsType2 >>= return . appType hsType1
-- fail "TyApp not yet supported" --g :@: h
i arg t@(TyVar hsName)        = addConst t >> (return $ TyApp (TyCon $ UnQual $ Ident "Const") t)
i arg t@(TyCon hsQName)       = addConst t >> (return $ TyApp (TyCon $ UnQual $ Ident "Const") t)

ii :: Type -> Type -> St Type
ii arg typ | typ == arg = return $ TyCon $ UnQual $ Ident "Id"
           | otherwise   = i arg typ

genInOut :: Type -> [QualConDecl] -> St [(Pat, Pat)]
genInOut arg [] = fail "empty list of constructors"
genInOut arg l  = mapM (g' arg) l >>= return . genPat

genPat :: [(Pat, Pat)] -> [(Pat, Pat)]
genPat [] = error "empty list found when generating instance of 'Mu' "
genPat l  = genPatAux id l

genPatAux ac [(a,b)] = [(ac a, b)]
genPatAux ac ((a,b):t) = (ac . inl $ a, b) : genPatAux (ac . inr) t

g' :: Type -> QualConDecl -> St (Pat,Pat)
g' arg = g'Con arg . conDecl

g'Con :: Type -> ConDecl -> St (Pat, Pat)
g'Con arg (ConDecl hsName []) = return (constPNil,PApp (UnQual hsName) [])
g'Con arg (ConDecl hsName lBangType) = do
   (l2,l3) <- mapAndUnzipM (h' arg) lBangType
   return (foldr1 times l2,PApp (UnQual hsName) l3)
g'Con arg (RecDecl hsName l) = g'Con arg (ConDecl hsName (concatMap unRec l))
--fail "data types records not yet supported"

h' :: Type -> BangType -> St (Pat, Pat)
h' arg (BangedTy   hsType) = ii' arg hsType
h' arg (UnBangedTy hsType) = ii' arg hsType
-- what's the diference between Banged and UnBanged?

i' :: Type -> Type -> St (Pat, Pat)
i' arg (TyFun hsType1 hsType2) = fail "TyFun not yet supported"
i' arg (TyTuple _ lType) = do
   (l1,l2) <- mapAndUnzipM (ii' arg) lType
   return (foldr1 times l1,PTuple l2)
i' arg (TyApp hsType1 hsType2) = do
   a <- getFreshVar
   let av = PVar a
   return (av,av)
i' arg t@(TyVar hsName) = do
   a <- getFreshVar
   let av = PVar a
   return (av,av)
i' arg t@(TyCon hsQName) = do
   a <- getFreshVar
   let av = case hsQName of
               Special _ -> PApp hsQName []
               _         -> PVar a
   return (av,av)

ii' :: Type -> Type -> St (Pat, Pat)
ii' arg typ | typ == arg = do
   a <- getFreshVar
   let av = PVar a
   return (av,av)
        | otherwise = i' arg typ

mkDataName :: Decl -> Type
mkDataName (DataDecl _ _ _ hsName lName _ _) = foldl TyApp (TyCon $ UnQual hsName) . map TyVar $ lName

getDataDeclFunctor :: Type -> [QualConDecl] -> St (Type,[Type])
getDataDeclFunctor arg lConDecl = withStateT (\((s,n),_) -> ((s,n),[])) $ do
   l1 <- mapM (g arg) lConDecl
   let functor = foldr1 plusType l1
   (_,consts) <- get
   return (functor,consts)

deriveTypeable :: Decl -> Decl
deriveTypeable (DataDecl loc dn ctx hsName lName lConDecl derive) =
   DataDecl loc dn ctx hsName lName lConDecl (nub $ UnQual typeable : derive)

getInstances :: Bool      -- ^ Observable or not
             -> Decl      -- ^ Data types
             -> St [Decl] -- ^ Instances for functor representation
getInstances ob d@(DataDecl loc _ [] hsName lName lConDecl _) = do
     let arg = mkDataName d
     (functor,consts) <- getDataDeclFunctor arg lConDecl
     l <- genInOut arg lConDecl
     let innOut = [InsDecl (FunBind (map inMatch l)), InsDecl (FunBind (map outMatch l))]
     let pfTInst = TypeInsDecl loc (TyApp pfType arg) functor
     let muInst = InstDecl loc [] mu [arg] innOut
     let observableInst = getObservableInst loc (nub consts) arg
     if ob then return [deriveTypeable d,pfTInst,muInst,observableInst]
           else return [d,pfTInst,muInst]
  where
    match str (a,b) = Exts.Match mkLoc (Ident str) [a]
                        (UnGuardedRhs $ hsPat2Exp b) (BDecls [])
    mkLoc = SrcLoc "" 0 0
    inMatch  = match "inn"
    outMatch = match "out" . swap

getInstances _ (DataDecl _ _ _ _ _ _ _ ) =
                 fail "type context not treated"
getInstances _ d = return [d]

-- types
opType op a b = TyApp (TyApp (TyVar $ Symbol op) a ) b
plusType  = opType ":+:"
timesType = opType ":*:"
appType   = opType ":@:"
constNil  = TyApp (TyCon $ UnQual $ Ident "Const") nil
nil = TyCon $ UnQual $ Ident "One"
mu = UnQual $ Ident "Mu"
pfType = TyCon $ UnQual $ Ident "PF"

-- patterns
constPNil = PWildCard
inl = PApp (UnQual $ Ident "Left") . singl . PParen
inr = PApp (UnQual $ Ident "Right") . singl . PParen
singl = (:[])
times a b = PTuple [a,b]

unRec (a,b) = replicate (length a) b

getFreshVar = do
   ((seed,n),l) <- gets id
   modify (\_->((seed,n+1),l))
   return $ Ident $ seed ++ show (n+1)

getObservableInst :: SrcLoc -> [Type] -> Type -> Decl
getObservableInst loc cts a = InstDecl loc ctx (UnQual (Ident "Observable")) [a] [InsDecl (FunBind [Exts.Match loc (Ident "observer") [PVar (Ident "x")] (UnGuardedRhs (App (App (Exts.Var (UnQual (Ident "send"))) (Lit (String ""))) (Paren (InfixApp (InfixApp (App (App (App (Exts.Var (UnQual (Ident "omap"))) (Paren (ExpTypeSig loc (Exts.Var (UnQual (Ident "_L"))) a))) (Exts.Var (UnQual (Ident "thk")))) (Paren (App (Exts.Var (UnQual (Ident "out"))) (Exts.Var (UnQual (Ident "x")))))) (QVarOp (UnQual (Symbol ">>="))) (Exts.Var (UnQual (Ident "return")))) (QVarOp (UnQual (Symbol "."))) (Exts.Var (UnQual (Ident "inn"))))))) (BDecls [PatBind loc (PVar (Ident "thk")) (UnGuardedRhs (ExpTypeSig loc (Exts.Var (UnQual (Ident "thunk"))) thunkSig )) (BDecls [])])])]
   where ctx = foldr (\c b -> mkIns typeable c : mkIns observable c : b) [] cts
         thunkSig = TyForall Nothing ctx $ TyFun a (TyApp (TyCon (UnQual (Ident "ObserverM"))) a)

typeable :: Name
typeable = Ident "Typeable"
observable :: Name
observable = Ident "Observable"

mkIns :: Name -> Type -> Asst
mkIns cl t = ClassA (UnQual cl) [t]

---- auxiliary functions
swap (a,b) = (b,a)

showName (Ident s) = s
showQName (UnQual (Ident s)) = s
showQName (Qual (ModuleName s1) (Ident s2)) = s1++('.':s2)

getSeed :: Data a => a -> String
getSeed = flip replicate 'x' .
          maximum . (1:) .
          everything (++) (mkQ [] aux)
 where aux = (:[]) . (+1) . length . takeWhile (=='x')