    {-# LANGUAGE CPP #-}

    -- | Check that a datatype is strictly positive.
    module Agda.TypeChecking.Positivity where

    import Control.Applicative hiding (empty)
    import Data.Set (Set)
    import qualified Data.Set as Set
    import Data.Map (Map)
    import qualified Data.Map as Map
    import Data.List as List

    import Agda.Syntax.Position
    import Agda.Syntax.Common
    import Agda.Syntax.Internal
    import Agda.TypeChecking.Monad
    import Agda.TypeChecking.Reduce
    import Agda.TypeChecking.Pretty
    import Agda.TypeChecking.Substitute

    import Agda.Utils.Impossible
    import Agda.Utils.Permutation
    import Agda.Utils.Size
    import Agda.Utils.Monad
    import Agda.Utils.SemiRing
    import qualified Agda.Utils.Graph as Graph
    import Agda.Utils.Graph (Graph)

    #include "../undefined.h"

    -- | Check that the datatypes in the given mutual block
    --   are strictly positive.
    checkStrictlyPositive :: MutualId -> TCM ()
    checkStrictlyPositive mi = do
      qs <- lookupMutualBlock mi
      reportSDoc "tc.pos.tick" 100 $ text "positivity of" <+> prettyTCM (Set.toList qs)
      g  <- buildOccurrenceGraph qs
      reportSDoc "tc.pos.tick" 100 $ text "constructed graph"
      reportSLn "tc.pos.graph" 5 $ "Positivity graph: N=" ++ show (size $ Graph.nodes g) ++
                                   " E=" ++ show (length $ Graph.edges g)
      reportSDoc "tc.pos.graph" 10 $ vcat
        [ text "positivity graph for" <+> prettyTCM (Set.toList qs)
        , nest 2 $ prettyGraph g
        ]
      mapM_ (setArgs g) $ Set.toList qs
      reportSDoc "tc.pos.tick" 100 $ text "set args"
      whenM positivityCheckEnabled $
        mapM_ (checkPos g) $ Set.toList qs
      reportSDoc "tc.pos.tick" 100 $ text "checked positivity"
      where
        checkPos g q = whenM (isDatatype q) $ do
          reportSDoc "tc.pos.check" 10 $ text "Checking positivity of" <+> prettyTCM q
          case Graph.findPath isNegative (DefNode q) (DefNode q) g of
            Nothing                  -> return ()
            Just (Edge Positive _)   -> __IMPOSSIBLE__
            Just (Edge Unused _)     -> __IMPOSSIBLE__
            Just (Edge Negative how) -> do
              err <- fsep $
                [prettyTCM q] ++ pwords "is not strictly positive, because it occurs" ++
                [prettyTCM how]
              setCurrentRange (getRange q) $ typeError $ GenericError (show err)

        isNegative (Edge o _) = o == Negative

        isDatatype q = do
          def <- theDef <$> getConstInfo q
          return $ case def of
            Datatype{dataClause = Nothing} -> True
            _ -> False

        getArity q = do
          def <- theDef <$> getConstInfo q
          return $ case def of
            Function{ funClauses = Clause{ clausePats = ps } : _ } -> size ps
            Datatype{ dataPars = n }                     -> n
            Record{ recPars = n }                        -> n
            _                                            -> 0

        -- Set the polarity of the arguments to a definition
        setArgs g q = do
          reportSDoc "tc.pos.args" 5 $ text "checking args of" <+> prettyTCM q
          n <- getArity q
          let nArgs = maximum $ n :
                        [ i + 1 | (ArgNode q1 i) <- Set.toList $ Graph.nodes g
                        , q1 == q ]
              findOcc i = case Graph.allPaths isNegative (ArgNode q i) (DefNode q) g of
                []                     -> Positive  -- Unused is not good for datatypes
                es | any isNegative es -> Negative
                   | otherwise         -> Positive
              args = map findOcc [0..nArgs - 1]
          reportSDoc "tc.pos.args" 10 $ sep
            [ text "args of" <+> prettyTCM q <+> text "="
            , nest 2 $ prettyList $ map (text . show) args
            ]
          setArgOccurrences q args

    -- Specification of occurrences -------------------------------------------

    instance SemiRing Occurrence where
      oplus Negative _        = Negative
      oplus _ Negative        = Negative
      oplus Unused o          = o
      oplus o Unused          = o
      oplus Positive Positive = Positive

      otimes Unused _          = Unused
      otimes _ Unused          = Unused
      otimes Negative _        = Negative
      otimes _ Negative        = Negative
      otimes Positive Positive = Positive

    -- | Description of an occurrence.
    data OccursWhere
      = LeftOfArrow OccursWhere
      | DefArg QName Nat OccursWhere -- ^ in the nth argument of a define constant
      | VarArg OccursWhere           -- ^ as an argument to a bound variable
      | MetaArg OccursWhere          -- ^ as an argument of a metavariable
      | ConArgType QName OccursWhere -- ^ in the type of a constructor
      | InClause Nat OccursWhere     -- ^ in the nth clause of a defined function
      | InDefOf QName OccursWhere    -- ^ in the definition of a constant
      | Here
      | Unknown                      -- ^ an unknown position (treated as negative)
      deriving (Show, Eq)

    (>*<) :: OccursWhere -> OccursWhere -> OccursWhere
    Here            >*< o  = o
    Unknown         >*< o  = Unknown
    LeftOfArrow o1  >*< o2 = LeftOfArrow (o1 >*< o2)
    DefArg d i o1   >*< o2 = DefArg d i (o1 >*< o2)
    VarArg o1       >*< o2 = VarArg (o1 >*< o2)
    MetaArg o1      >*< o2 = MetaArg (o1 >*< o2)
    ConArgType c o1 >*< o2 = ConArgType c (o1 >*< o2)
    InClause i o1   >*< o2 = InClause i (o1 >*< o2)
    InDefOf d o1    >*< o2 = InDefOf d (o1 >*< o2)

    instance PrettyTCM OccursWhere where
      prettyTCM o = prettyOs $ map maxOneLeftOfArrow $ uniq $ splitOnDef o
        where
          nth 0 = pwords "first"
          nth 1 = pwords "second"
          nth 2 = pwords "third"
          nth n = pwords $ show (n + 1) ++ "th"

          uniq (x:y:xs)
            | x == y  = uniq (x:xs)
          uniq (x:xs) = x : uniq xs
          uniq []     = []

          prettyOs [] = __IMPOSSIBLE__
          prettyOs [o] = prettyO o <> text "."
          prettyOs (o:os) = prettyO o <> text ", which occurs" <+> prettyOs os

          prettyO o = case o of
            Here           -> empty
            Unknown        -> empty
            LeftOfArrow o  -> explain o $ pwords "to the left of an arrow"
            DefArg q i o   -> explain o $ pwords "in the" ++ nth i ++ pwords "argument to" ++
                                          [prettyTCM q]
            VarArg o       -> explain o $ pwords "in an argument to a bound variable"
            MetaArg o      -> explain o $ pwords "in an argument to a metavariable"
            ConArgType c o -> explain o $ pwords "in the type of the constructor" ++ [prettyTCM c]
            InClause i o   -> explain o $ pwords "in the" ++ nth i ++ pwords "clause"
            InDefOf d o    -> explain o $ pwords "in the definition of" ++ [prettyTCM d]

          explain o ds = prettyO o $$ fsep ds

          maxOneLeftOfArrow o = case o of
            LeftOfArrow o  -> LeftOfArrow $ purgeArrows o
            Here           -> Here
            Unknown        -> Unknown
            DefArg q i o   -> DefArg q i   $ maxOneLeftOfArrow o
            InDefOf d o    -> InDefOf d    $ maxOneLeftOfArrow o
            VarArg o       -> VarArg       $ maxOneLeftOfArrow o
            MetaArg o      -> MetaArg      $ maxOneLeftOfArrow o
            ConArgType c o -> ConArgType c $ maxOneLeftOfArrow o
            InClause i o   -> InClause i   $ maxOneLeftOfArrow o

          purgeArrows o = case o of
            LeftOfArrow o -> purgeArrows o
            Here           -> Here
            Unknown        -> Unknown
            DefArg q i o   -> DefArg q i   $ purgeArrows o
            InDefOf d o    -> InDefOf d    $ purgeArrows o
            VarArg o       -> VarArg       $ purgeArrows o
            MetaArg o      -> MetaArg      $ purgeArrows o
            ConArgType c o -> ConArgType c $ purgeArrows o
            InClause i o   -> InClause i   $ purgeArrows o

          splitOnDef o = case o of
            Here           -> [Here]
            Unknown        -> [Unknown]
            InDefOf d o    -> sp (InDefOf d) o
            LeftOfArrow o  -> sp LeftOfArrow o
            DefArg q i o   -> sp (DefArg q i) o
            VarArg o       -> sp VarArg o
            MetaArg o      -> sp MetaArg o
            ConArgType c o -> sp (ConArgType c) o
            InClause i o   -> sp (InClause i) o
            where
              sp f o = case splitOnDef o of
                os@(InDefOf _ _:_) -> f Here : os
                o:os               -> f o : os
                []                 -> __IMPOSSIBLE__

    -- Computing occurrences --------------------------------------------------

    data Item = AnArg Nat
              | ADef QName
      deriving (Eq, Ord, Show)

    type Occurrences = Map Item [OccursWhere]

    (>+<) :: Occurrences -> Occurrences -> Occurrences
    (>+<) = Map.unionWith (++)

    concatOccurs :: [Occurrences] -> Occurrences
    concatOccurs = Map.unionsWith (++)

    occursAs :: (OccursWhere -> OccursWhere) -> Occurrences -> Occurrences
    occursAs f = Map.map (map f)

    here :: Item -> Occurrences
    here i = Map.singleton i [Here]

    class ComputeOccurrences a where
      -- | The first argument is the items corresponding to the free variables.
      occurrences :: [Maybe Item] -> a -> Occurrences

    instance ComputeOccurrences Clause where
      occurrences vars (Clause{ clausePats = ps, clauseBody = body }) =
        concatOccurs (zipWith match [0..] ps) >+<
        walk vars (patItems ps) body
        where
          walk _    _         NoBody     = Map.empty
          walk vars []        (Body v)   = occurrences vars v
          walk vars (i : pis) (Bind b)   = walk (i : vars) pis $ absBody b
          walk vars (_ : pis) (NoBind b) = walk vars pis b
          walk _    []        Bind{}     = __IMPOSSIBLE__
          walk _    []        NoBind{}   = __IMPOSSIBLE__
          walk _    (_ : _)   Body{}     = __IMPOSSIBLE__

          match i (Arg _ VarP{}) = Map.empty
          match i _              = Map.singleton (AnArg i) [Unknown]

          patItems ps = concat $ zipWith patItem [0..] $ map unArg ps
          patItem i (VarP _) = [Just (AnArg i)]
          patItem i p        = replicate (nVars p) Nothing

          nVars p = case p of
            VarP{}    -> 1
            DotP{}    -> 1
            ConP _ ps -> sum $ map (nVars . unArg) ps
            LitP{}    -> 0

    instance ComputeOccurrences Term where
      occurrences vars v = case v of
        Var i args ->
          maybe Map.empty here (vars ! fromIntegral i)
          >+< occursAs VarArg (occurrences vars args)
        Def d args   ->
          here (ADef d) >+<
          concatOccurs (zipWith (occursAs . DefArg d) [0..] $ map (occurrences vars) args)
        Con c args   -> occurrences vars args
        MetaV _ args -> occursAs MetaArg $ occurrences vars args
        Pi a b       -> occursAs LeftOfArrow (occurrences vars a) >+<
                        occurrences vars b
        Fun a b      -> occursAs LeftOfArrow (occurrences vars a) >+<
                        occurrences vars b
        Lam _ b      -> occurrences vars b
        Lit{}        -> Map.empty
        Sort{}       -> Map.empty
        where
          vs ! i
            | i < length vs = vs !! i
            | otherwise     = error $ show vs ++ " ! " ++ show i ++ "  (" ++ show v ++ ")"

    instance ComputeOccurrences Type where
      occurrences vars (El _ v) = occurrences vars v

    instance ComputeOccurrences Telescope where
      occurrences vars EmptyTel        = Map.empty
      occurrences vars (ExtendTel a b) = occurrences vars (a, b)

    instance ComputeOccurrences a => ComputeOccurrences (Abs a) where
      occurrences vars = occurrences (Nothing : vars) . absBody

    instance ComputeOccurrences a => ComputeOccurrences (Arg a) where
      occurrences vars = occurrences vars . unArg

    instance ComputeOccurrences a => ComputeOccurrences [a] where
      occurrences vars = concatOccurs . map (occurrences vars)

    instance (ComputeOccurrences a, ComputeOccurrences b) => ComputeOccurrences (a, b) where
      occurrences vars (x, y) = occurrences vars x >+< occurrences vars y

    -- | Compute the occurrences in a given definition.
    computeOccurrences :: QName -> TCM Occurrences
    computeOccurrences q = do
      def <- getConstInfo q
      occursAs (InDefOf q) <$> case theDef def of
        Function{funClauses = cs} -> do
          n  <- arity <$> instantiateFull (defType def)
          cs <- map (etaExpandClause n) <$> instantiateFull cs
          return
            $ concatOccurs
            $ zipWith (occursAs . InClause) [0..]
            $ map (occurrences []) cs
        Datatype{dataClause = Just c} -> occurrences [] <$> instantiateFull c
        Datatype{dataPars = np, dataCons = cs}       -> do
          let conOcc c = do
                a <- defType <$> getConstInfo c
                TelV tel _ <- telView <$> normalise a
                let tel' = telFromList $ genericDrop np $ telToList tel
                    vars = reverse [ Just (AnArg i) | i <- [0..np - 1] ]
                return $ occursAs (ConArgType c) $ occurrences vars tel'
          concatOccurs <$> mapM conOcc cs
        Record{recClause = Just c} -> occurrences [] <$> instantiateFull c
        Record{recPars = np, recTel = tel} -> do
          let tel' = telFromList $ genericDrop np $ telToList tel
              vars = reverse [ Just (AnArg i) | i <- [0..np - 1] ]
          occurrences vars <$> instantiateFull tel'

        -- Arguments to other kinds of definitions are hard-wired.
        Constructor{} -> return Map.empty
        Axiom{}       -> return Map.empty
        Primitive{}   -> return Map.empty

    -- | Eta expand a clause to have the given number of variables.
    --   Warning: doesn't update telescope or permutation!
    --   This is used instead of special treatment of lambdas
    --   (which was unsound: issue 121)
    etaExpandClause :: Nat -> Clause -> Clause
    etaExpandClause n c@Clause{ clausePats = ps, clauseBody = b }
      | m <= 0    = c
      | otherwise = c { clausePats = ps ++ genericReplicate m (Arg NotHidden $ VarP "_")
                      , clauseBody = liftBody m b
                      }
      where
        m = n - genericLength ps

        bind 0 = id
        bind n = Bind . Abs "_" . bind (n - 1)

        vars = reverse [ Arg NotHidden $ Var i [] | i <- [0..m - 1] ]

        liftBody m (Bind b)   = Bind $ fmap (liftBody m) b
        liftBody m (NoBind b) = NoBind $ liftBody m b
        liftBody m NoBody     = bind m NoBody
        liftBody m (Body v)   = bind m $ Body $ raise m v `apply` vars

    -- Building the occurrence graph ------------------------------------------

    data Node = DefNode QName
              | ArgNode QName Nat
      deriving (Eq, Ord)

    instance Show Node where
      show (DefNode q)   = show q
      show (ArgNode q i) = show q ++ "." ++ show i

    instance PrettyTCM Node where
      prettyTCM (DefNode q)   = prettyTCM q
      prettyTCM (ArgNode q i) = prettyTCM q <> text ("." ++ show i)

    prettyGraph g = vcat $ map pr $ Map.assocs $ Graph.unGraph g
      where
        pr (n, es) = sep
          [ prettyTCM n
          , nest 2 $ vcat $ map prE $ Map.assocs es
          ]
        prE (n, Edge o w) = prO o <+> prettyTCM n <+> fsep (pwords $ show w)
        prO Positive = text "-[+]->"
        prO Negative = text "-[-]->"
        prO Unused   = text "-[ ]->"

    data Edge = Edge Occurrence OccursWhere
      deriving (Show)

    instance SemiRing Edge where
      oplus _                   e@(Edge Negative _) = e
      oplus e@(Edge Negative _) _                   = e
      oplus (Edge Unused _)     e                   = e
      oplus e                   (Edge Unused _)     = e
      oplus (Edge Positive _)   e@(Edge Positive _) = e

      otimes (Edge o1 w1) (Edge o2 w2) = Edge (otimes o1 o2) (w1 >*< w2)

    buildOccurrenceGraph :: Set QName -> TCM (Graph Node Edge)
    buildOccurrenceGraph qs = Graph.unions <$> mapM defGraph (Set.toList qs)
      where
        defGraph :: QName -> TCM (Graph Node Edge)
        defGraph q = do
          occs <- computeOccurrences q
          let onItem (item, occs) = do
                es <- mapM (computeEdge qs) occs
                return $ Graph.unions $
                    map (\(b, w) -> Graph.singleton (itemToNode item) b w) es
          Graph.unions <$> mapM onItem (Map.assocs occs)
          where
            itemToNode (AnArg i) = ArgNode q i
            itemToNode (ADef q)  = DefNode q

    -- | Given an 'OccursWhere' computes the target node and an 'Edge'. The first
    --   argument is the set of names in the current mutual block.
    computeEdge :: Set QName -> OccursWhere -> TCM (Node, Edge)
    computeEdge muts o = do
      (to, occ) <- mkEdge __IMPOSSIBLE__ Positive o
      return (to, Edge occ o)
      where
        mkEdge to pol o = case o of
          Here           -> return (to, pol)
          Unknown        -> return (to, Negative)
          VarArg o       -> negative o
          MetaArg o      -> negative o
          LeftOfArrow o  -> negative o
          DefArg d i o
            | Set.member d muts -> inArg d i o
            | otherwise         -> addPol o =<< getArgOccurrence d i
          ConArgType _ o -> keepGoing o
          InClause _ o   -> keepGoing o
          InDefOf d o    -> mkEdge (DefNode d) Positive o
          where
            keepGoing     = mkEdge to pol
            negative      = mkEdge to Negative
            addPol o pol' = mkEdge to (otimes pol pol') o

            -- Reset polarity when changing the target node
            -- D: (A B -> C) generates a positive edge B --> A.1
            -- even though the context is negative.
            inArg d i = mkEdge (ArgNode d i) Positive