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) = liftM (foldr1 timesType) (mapM (h arg) lBangType)
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)       = liftM (foldr1 timesType) (mapM (ii arg) lType)
i arg (TyApp hsType1 hsType2) = liftM (appType hsType1) (i arg hsType2)
-- 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  = liftM genPat (mapM (g' arg) l)

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] Nothing
                        (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 = TyApp . TyApp (TyVar $ Symbol op)
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 (const ((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")] Nothing (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")) Nothing (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')