    {-# LANGUAGE CPP #-}

    {-| Generate GHC code for optimized datatypes and their values
    -}

    module Agda.Compiler.Agate.OptimizedPrinter where

    #include "../../undefined.h"
    import Agda.Utils.Impossible

    import Control.Monad.Trans
    import qualified Data.Map as Map
    import Data.Map ((!), Map)

    import Agda.Compiler.Agate.Classify
    import Agda.Compiler.Agate.TranslateName
    import Agda.Compiler.Agate.Common

    import Agda.Syntax.Common
    import Agda.Syntax.Internal
    import Agda.TypeChecking.Monad
    import Agda.TypeChecking.Reduce
    import Agda.TypeChecking.Free
    import Agda.Utils.Pretty
    import Agda.Utils.Size

    ----------------------------------------------------------------

    showAsOptimizedConstructor :: QName -> TCM Doc
    showAsOptimizedConstructor s =
        return $ text $ translateNameAsOptimizedConstructor $ show s

    ----------------------------------------------------------------

    showOptimizedDefinitions :: Definitions -> TCM ()
    showOptimizedDefinitions definitions = do
        liftIO $ putStrLn ("\n-----\n")
        typefams <- enumCompilableTypeFamilies definitions
        dtypedecls <- showTypeDeclarations definitions typefams
        liftIO $ putStrLn $ render dtypedecls
        optimizableConstants <- enumOptimizableConstants definitions typefams
        liftIO $ putStrLn $ ""
        dvaluedefs <- showValueDefinitions definitions typefams
        liftIO $ putStrLn $ render dvaluedefs
        return ()

    ----------------------------------------------------------------
    -- Generating GHC Type Declarations

    showTypeDeclarations :: Definitions -> [QName] -> TCM Doc
    showTypeDeclarations definitions typefams = do
        dtypedecls <- mapM (showTypeDeclaration definitions) typefams
        return $ vcat dtypedecls

    showTypeDeclaration :: Definitions -> QName -> TCM Doc
    showTypeDeclaration definitions name = do
        dtypename <- showAsOptimizedType name
        let defn = definitions ! name
        ty <- normalise $ defType defn
        case theDef defn of
    	Datatype{ dataPars = np, dataIxs = ni, dataCons = cnames } | ni == 0 -> do
    	    dparams <- underDatatypeParameters np ty showTypeParameter stop
    	    dargs <- mapM (showConstructorDeclaration np) cnames
    	    showDatatypeDeclaration dtypename dparams dargs
    	Record{ recPars = np, recFields = flds } -> do
    	    let arity = np `div` 2
    	    dparams <- underDatatypeParameters arity ty showTypeParameter stop
    	    dcon <- showAsOptimizedConstructor name
    	    dflds <- mapM (showRecordFieldDeclaration (arity + 1)) (map snd flds)
    	    showDatatypeDeclaration dtypename dparams [sep $ dcon : dflds]
         	Function { funClauses = [clause@(Clause{ clausePats = pat })] } -> do
    	    (args,_) <- splitType ty
    	    let dextra = map (\i -> text $ "a" ++ show i)
    			     [length pat + 1 .. length args]
    	    showTypeSynonymDeclaration dtypename dextra clause
         	_ -> return empty
        where
        stop ty = return []
        showDatatypeDeclaration dtypename dparams [] =
    	return $ (text "type") <+> (sep (dtypename : dparams) <+> equals <+>
    				    text "()")
        showDatatypeDeclaration dtypename dparams dargs =
    	return $ (text "data") <+> (sep $
    				    sep (dtypename : dparams) <+> equals :
    				    punctuate (text " |") dargs)
    				$+$ text "    -- deriving (Show)"
        showTypeSynonymDeclaration dtypename dextra = showClause fTerm fCon fBody
    	where
    	fTerm = showAsOptimizedType
    	fCon name dargs = __IMPOSSIBLE__
    	fBody dpats term = do
    	    dterm <- showAsOptimizedType term
    	    return $ (text "type") <+>
    		     (sep (dtypename : (dpats ++ dextra)) <+> equals <+>
    		      sep (dterm : dextra))

        showConstructorDeclaration np cname = do
    	let defn = definitions ! cname
    	ty <- normalise $ defType defn
    	dcon <- showAsOptimizedConstructor cname
    	underDatatypeParameters np ty (\d t -> return id) $ \ty -> do
    	    dargs <- forEachArgM showAsOptimizedType ty
    	    return $ sep $ dcon : dargs
        showRecordFieldDeclaration np fldname = do
    	let defn = definitions ! fldname
    	ty <- normalise $ defType defn
    	underDatatypeParameters np ty (\d t -> return id) showAsOptimizedType

    showTypeParameter :: Doc -> Type -> TCM ([Doc] -> [Doc])
    showTypeParameter dname ty = do
        ty <- instantiate ty
        (args,_) <- splitType ty
        case args of
    	[]  -> return $ (:) dname
    	_:_ -> do
    		dk <- showAsOptimizedKind ty
    		return $ (:) $ parens $ sep [ dname <+> text "::", dk ]

    underDatatypeParameters :: Nat -> Type -> (Doc -> Type -> TCM (a -> a)) ->
    			   (Type -> TCM a) -> TCM a
    underDatatypeParameters np ty f k = go 0 ty where
        go i ty        | i >= np = k ty
        go i (El s tm)           = do
          let varname = show (i + 1)
          tm <- reduce tm
          case tm of
    	Pi arg abs -> underAbstraction arg abs{absName = varname} $ \ty -> do
    	    dname <- showAsOptimizedTerm $ Var 0 []
    	    cont <- f dname $ unArg arg
    	    dparams <- go (i + 1) ty
    	    return $ cont dparams
    	Fun arg ty -> do
    	    let dname = text $ translateNameAsOptimizedTerm varname
    	    cont <- f dname $ unArg arg
    	    dparams <- go (i + 1) ty
    	    return $ cont dparams
    	Var _ _    -> __IMPOSSIBLE__
    	Def _ _    -> __IMPOSSIBLE__
    	Con _ _    -> __IMPOSSIBLE__
    	Sort _	   -> __IMPOSSIBLE__
    	Lit _	   -> __IMPOSSIBLE__
    	MetaV _ _  -> __IMPOSSIBLE__
    	Lam _ _	   -> __IMPOSSIBLE__

    ----------------------------------------------------------------
    -- Generating GHC Value Definitions

    showValueDefinitions :: Definitions -> [QName] -> TCM Doc
    showValueDefinitions definitions typefams = do
        let defs = Map.toList definitions
        dvaluedefs <- mapM (showValueDefinition definitions typefams) defs
        return $ vcat dvaluedefs

    showValueDefinition :: Definitions -> [QName] -> (QName,Definition) -> TCM Doc
    showValueDefinition definitions typefams (name, defn) = do
        dname <- showAsOptimizedTerm name
        ty <- normalise $ defType defn
        dty <- showAsOptimizedType ty
        let dtypedecl = sep [ dname, text "::" ] <+> dty
        case theDef defn of
    	Axiom{} -> do
    	    let dvaluedef = sep [ dname, equals ] <+>
    			    sep [ text "undefined {- postulate -}" ]
    	    return $ dtypedecl $+$ dvaluedef
    	Function { funClauses = clauses } -> do
    	    dclauses <- mapM (showOptimizedClause dname) clauses
    	    return $ vcat $ dtypedecl : dclauses
    	Datatype{ dataPars = np, dataIxs = ni, dataCons = cnames } -> do
    	    let dvars = map (\i -> text ("v" ++ show i)) [1 .. np + ni]
    	    let dvaluedef = sep [ sep (dname : dvars), equals ] <+> text "()"
    	    return $ dtypedecl $+$ dvaluedef
    	Record{} ->
    	    return empty  -- no o_xxx since we always use D_xxx
    	Constructor{} ->
    	    return empty  -- no o_xxx since we always use D_xxx
    	Primitive { primName = pf } -> do
    	    let dvaluedef = sep [ dname, equals ] <+>
    			    sep [ text (show name), text "{- primitive -}" ]
    	    return $ dtypedecl $+$ dvaluedef

    showOptimizedClause :: Doc -> Clause -> TCM Doc
    showOptimizedClause dfuncname = showClause fTerm fCon fBody where
        fTerm = showAsOptimizedTerm
        fCon name dargs = do
    	dname <- showAsOptimizedConstructor name
    	return $ parens $ sep (dname : dargs)
        fBody dpats term = do
    	dterm <- showAsOptimizedTerm term
    	return $ (sep (dfuncname : dpats) <+> equals) <+> nest 2 dterm

    ----------------------------------------------------------------
    -- implementation of the "K" function

    class ShowAsOptimizedKind a where
        showAsOptimizedKind :: a -> TCM Doc

    instance ShowAsOptimizedKind Type where
        showAsOptimizedKind (El s t) = showAsOptimizedKind t

    instance ShowAsOptimizedKind Term where
        showAsOptimizedKind t = do
    	t <- reduce t
    	case t of
    	    Pi arg abs -> do
    		dk1 <- showAsOptimizedKind $ unArg arg
    		dk2 <- showAsOptimizedKind $ absBody abs
    		return $ parens $ sep [dk1 <+> text "->", dk2]
    	    Fun arg ty -> do
    		dk1 <- showAsOptimizedKind $ unArg arg
    		dk2 <- showAsOptimizedKind $ ty
    		return $ parens $ sep [dk1 <+> text "->", dk2]
    	    Sort s     -> return $ text "*"
    	    Var _ _    -> return $ text "*" -- ok?
    	    Def _ _    -> __IMPOSSIBLE__
    	    Con _ _    -> __IMPOSSIBLE__
    	    Lit _      -> __IMPOSSIBLE__
    	    MetaV _ _  -> __IMPOSSIBLE__
    	    Lam _ _    -> __IMPOSSIBLE__

    ----------------------------------------------------------------
    -- implementation of the "T" function

    class ShowAsOptimizedType a where
        showAsOptimizedType :: a -> TCM Doc

    instance ShowAsOptimizedType QName where
        showAsOptimizedType s = return $ text $ translateNameAsOptimizedType $ show s

    instance ShowAsOptimizedType Type where
        showAsOptimizedType (El s t) = showAsOptimizedType t

    instance ShowAsOptimizedType Term where
        showAsOptimizedType t = do
    	t <- reduce t
    	case t of
    	    Var n args -> do
    		varname <- nameOfBV n
    		dvar <- showAsOptimizedTerm varname
    		dargs <- mapM (showAsOptimizedType . unArg) args
    		return $ psep $ dvar : dargs
    	    Def name args -> do
    		dname <- showAsOptimizedType name
    		dargs <- mapM (showAsOptimizedType . unArg) args
    		return $ psep $ dname : dargs
    	    Con name args -> do
    		return $ text "<constructor term (impossible)>"
    	    Pi arg abs -> do
    		dt1 <- showAsOptimizedType $ unArg arg
    		underAbstraction_ abs $ \body -> do
    		    dvar <- showAsOptimizedTerm $ Var 0 []
    		    dt2 <- showAsOptimizedType body
    		    case freeIn 0 body of
    			True  -> return $ parens $ sep $
    			    [ sep [ text "forall",
    				    dvar <> text ".", dt1, text "->" ],
    			      dt2 ]
    			False -> return $ parens $ sep [dt1 <+> text "->", dt2]
    	    Fun arg ty -> do
    		dt1 <- showAsOptimizedType $ unArg arg
    		dt2 <- showAsOptimizedType $ ty
    		return $ parens $ sep [dt1, text "->", dt2]
    	    Sort s        -> return $ text "()"
    	    MetaV id args -> return $ text "<meta (impossible)>"
    	    Lam h t       -> __IMPOSSIBLE__
    	    Lit lit       -> __IMPOSSIBLE__

    ----------------------------------------------------------------
    -- implementation of the "O" function

    class ShowAsOptimizedTerm a where
        showAsOptimizedTerm :: a -> TCM Doc

    instance ShowAsOptimizedTerm Name where
        showAsOptimizedTerm s = return $ text $ translateNameAsOptimizedTerm $ show s

    instance ShowAsOptimizedTerm QName where
        showAsOptimizedTerm s = return $ text $ translateNameAsOptimizedTerm $ show s

    instance ShowAsOptimizedTerm Term where
        showAsOptimizedTerm (Var n args) = do
        	varname <- nameOfBV n
        	dvar <- showAsOptimizedTerm varname
        	dargs <- mapM (showAsOptimizedTerm . unArg) args
        	return $ psep $ dvar : dargs
        showAsOptimizedTerm (Lam h abs) = underAbstraction_ abs $ \body -> do
        	dvar <- showAsOptimizedTerm $ Var 0 []
        	dt <- showAsOptimizedTerm body
        	return $ parens $ sep [ text "\\" <> dvar, text "->", dt ]
        showAsOptimizedTerm (Def name args) = do
    	dname <- showAsOptimizedTerm name
    	dargs <- mapM (showAsOptimizedTerm . unArg) args
    	return $ psep $ dname : dargs
        showAsOptimizedTerm (Con name args) = do
        	dname <- showAsOptimizedConstructor name
        	dargs <- mapM (showAsOptimizedTerm . unArg) args
        	return $ psep $ dname : dargs
        showAsOptimizedTerm (Lit lit)        = return $ showOptimizedLiteral lit
        showAsOptimizedTerm (Pi _ _)	 = return $ text "()"
        showAsOptimizedTerm (Fun _ _)	 = return $ text "()"
        showAsOptimizedTerm (Sort _)	 = return $ text "()"
        showAsOptimizedTerm (MetaV id args)  = __IMPOSSIBLE__

    ----------------------------------------------------------------
    --