    -- |
    -- Module      :  $Header$
    -- Copyright   :  (c) 2014-2016 Galois, Inc.
    -- License     :  BSD3
    -- Maintainer  :  cryptol@galois.com
    -- Stability   :  provisional
    -- Portability :  portable
    --
    -- Desugar into SMTLIB Terminology

    {-# LANGUAGE Safe #-}

    module Cryptol.TypeCheck.Solver.Numeric.SMT
      ( desugarProp
      , smtName
      , smtFinName
      , ifPropToSmtLib
      , cryImproveModel
      , getVal
      , getVals
      ) where

    import           Cryptol.TypeCheck.AST (TVar(TVFree,TVBound))
    import           Cryptol.TypeCheck.Solver.InfNat
    import           Cryptol.TypeCheck.Solver.Numeric.AST
    import           Cryptol.TypeCheck.Solver.Numeric.Simplify(crySimplify)
    import           Cryptol.Utils.Misc ( anyJust )
    import           Cryptol.Utils.Panic ( panic )

    import           Data.List ( partition, unfoldr )
    import           Data.Map ( Map )
    import qualified Data.Map as Map
    import qualified Data.Set as Set
    import           SimpleSMT ( SExpr )
    import qualified SimpleSMT as SMT
    import           MonadLib

    --------------------------------------------------------------------------------
    -- Desugar to SMT
    --------------------------------------------------------------------------------

    -- XXX: Expanding the if-then-elses could make things large.
    -- Perhaps keep them as first class things, in hope that the solver
    -- can do something more clever with that?

    -- | Assumes simplified, linear, finite, defined expressions.
    desugarExpr :: Expr -> IfExpr Expr
    desugarExpr expr =
      do es <- mapM desugarExpr (cryExprExprs expr)
         case (expr,es) of
           (Min {}, [x,y]) -> If (x :>: y) (return y) (return x)
           (Max {}, [x,y]) -> If (x :>: y) (return x) (return y)
           (LenFromThenTo {}, [ x@(K (Nat a)), K (Nat b), z ])

              -- going down
              | a > b -> If (z :>: x) (return zero)
                                      (return (Div (x :- z) step :+ one))

              -- goind up
              | b > a -> If (x :>: z) (return zero)
                                      (return (Div (z :- x) step :+ one))

              where step = K (Nat (abs (a - b)))

           _ -> return (cryRebuildExpr expr es)

    -- | Assumes simplified, linear, defined.
    desugarProp :: Prop -> IfExpr Prop
    desugarProp prop =
      case prop of
        PFalse      -> return PFalse
        PTrue       -> return PTrue
        Not p       -> Not   `fmap` desugarProp p
        p :&& q     -> (:&&) `fmap` desugarProp p `ap` desugarProp q
        p :|| q     -> (:||) `fmap` desugarProp p `ap` desugarProp q
        Fin (Var _) -> return prop
        x :==: y    -> (:==:) `fmap` desugarExpr x `ap` desugarExpr y
        x :>: y     -> (:>:)  `fmap` desugarExpr x `ap` desugarExpr y

        Fin _     -> unexpected
        _ :== _   -> unexpected
        _ :>= _   -> unexpected
        _ :> _    -> unexpected

      where
      unexpected = panic "desugarProp" [ show (ppProp prop) ]

    ifPropToSmtLib :: IfExpr Prop -> SExpr
    ifPropToSmtLib ifProp =
      case ifProp of
        Impossible -> SMT.bool False -- Sholdn't really matter
        Return p   -> propToSmtLib p
        If p q r   -> SMT.ite (propToSmtLib p) (ifPropToSmtLib q) (ifPropToSmtLib r)

    propToSmtLib :: Prop -> SExpr
    propToSmtLib prop =
      case prop of
        PFalse       -> SMT.bool False
        PTrue        -> SMT.bool True
        Not p        -> case p of
                          Fin _   -> SMT.not (propToSmtLib p)

                          -- It is IMPORTANT that the fin constraints are outside
                          -- the not.
                          x :>: y -> addFin $ SMT.geq (exprToSmtLib y)
                                                      (exprToSmtLib x)
                          _ -> unexpected

        p :&& q     -> SMT.and (propToSmtLib p) (propToSmtLib q)
        p :|| q     -> SMT.or  (propToSmtLib p) (propToSmtLib q)
        Fin (Var x) -> fin x

        {- For the linear constraints, if the term is finite, then all of
           its variables must have been finite.

           XXX: Adding the `fin` decls at the leaves leads to some duplication:
           We could add them just once for each conjunctoin of simple formulas,
           but I am not sure how much this matters.
        -}
        x :==: y    -> addFin $ SMT.eq (exprToSmtLib x) (exprToSmtLib y)
        x :>: y     -> addFin $ SMT.gt (exprToSmtLib x) (exprToSmtLib y)

        Fin _       -> unexpected
        _ :== _     -> unexpected
        _ :>= _     -> unexpected
        _ :> _      -> unexpected

      where
      unexpected = panic "propToSmtLib" [ show (ppProp prop) ]
      fin x      = SMT.const (smtFinName x)

      addFin e   = foldr (\x' e' -> SMT.and (fin x') e') e
                         (Set.toList (cryPropFVS prop))

    exprToSmtLib :: Expr -> SExpr
    exprToSmtLib expr =

      case expr of
        K (Nat n)           -> SMT.int n
        Var a               -> SMT.const (smtName a)
        x :+ y              -> SMT.add (exprToSmtLib x) (exprToSmtLib y)
        x :- y              -> SMT.sub (exprToSmtLib x) (exprToSmtLib y)
        x :* y              -> SMT.mul (exprToSmtLib x) (exprToSmtLib y)
        Div x y             -> SMT.div (exprToSmtLib x) (exprToSmtLib y)
        Mod x y             -> SMT.mod (exprToSmtLib x) (exprToSmtLib y)

        K Inf               -> unexpected
        _ :^^ _             -> unexpected
        Min {}              -> unexpected
        Max {}              -> unexpected
        Width {}            -> unexpected
        LenFromThen {}      -> unexpected
        LenFromThenTo {}    -> unexpected

      where
      unexpected = panic "exprToSmtLib" [ show (ppExpr expr) ]

    -- | The name of a variable in the SMT translation.
    smtName :: Name -> String
    smtName a = case a of
                  SysName n -> name "s" n
                  UserName tv -> case tv of
                                   TVFree n _ _ _ -> name "u" n
                                   TVBound n _    -> name "k" n

      where
      name p n = case divMod n 26 of
                   (q,r) -> p ++ toEnum (fromEnum 'a' + r) :
                                  (if q == 0 then "" else show q)

    -- | The name of a boolean variable, representing `fin x`.
    smtFinName :: Name -> String
    smtFinName x = "fin_" ++ smtName x

    --------------------------------------------------------------------------------
    -- Models
    --------------------------------------------------------------------------------

    {- | Get the value for the given name.
          * Assumes that we are in a SAT state (i.e., there is a model)
          * Assumes that the name is in the model -}
    getVal :: SMT.Solver -> Name -> IO Nat'
    getVal s a =
      do yes <- isInf a
         if yes
           then return Inf
           else do v <- SMT.getConst s (smtName a)
                   case v of
                     SMT.Int x | x >= 0 -> return (Nat x)
                     _ -> panic "cryCheck.getVal" [ "Not a natural number", show v ]

      where
      isInf v = do yes <- SMT.getConst s (smtFinName v)
                   case yes of
                     SMT.Bool ans -> return (not ans)
                     _            -> panic "cryCheck.isInf"
                                           [ "Not a boolean value", show yes ]

    -- | Get the values for the given names.
    getVals :: SMT.Solver -> [Name] -> IO (Map Name Nat')
    getVals s xs =
      do es <- mapM (getVal s) xs
         return (Map.fromList (zip xs es))

    -- | Convert a bunch of improving equations into an idempotent substitution.
    -- Assumes that the equations don't have loops.
    toSubst :: Map Name Expr -> Subst
    toSubst m0 = last (m0 : unfoldr step m0)
      where step m = do m1 <- anyJust (apSubst m) m
                        return (m1,m1)

    {- | Given a model, compute an improving substitution, implied by the model.
    The entries in the substitution look like this:

      * @x = A@         variable @x@ must equal constant @A@

      * @x = y@         variable @x@ must equal variable @y@

      * @x = A * y + B@ (coming soon)
                        variable @x@ is a linear function of @y@,
                        @A@ and @B@ are natural numbers.
    -}

    {- | We are mostly interested in improving unification variables.
    However, it is also useful to improve skolem variables, as this could
    turn non-linear constraints into linear ones.  For example, if we
    have a constraint @x * y = z@, and we can figure out that @x@ must be 5,
    then we end up with a linear constraint @5 * y = z@.
    -}
    cryImproveModel :: SMT.Solver -> SMT.Logger -> Map Name Nat'
                    -> IO (Map Name Expr, [Prop])
    cryImproveModel solver logger model =
      do (imps,subGoals) <- go Map.empty [] initialTodo
         return (toSubst imps, subGoals)

      where
      -- Process unification variables first.  That way, if we get `x = y`, we'd
      -- prefer `x` to be a unification variable.
      initialTodo    = uncurry (++) $ partition (isUniVar . fst) $ Map.toList model
      isUniVar x     = case x of
                         UserName (TVFree {}) -> True
                         _                    -> False

      -- done:  the set of known improvements
      -- extra: the collection of inferred sub-goals
      go done extra []             = return (done,extra)
      go done extra ((x,e) : rest) =

        -- x = K?
        do mbCounter <- cryMustEqualK solver (Map.keys model) x e
           case mbCounter of
             Nothing -> go (Map.insert x (K e) done) extra rest
             Just ce -> goV ce done extra [] x e rest

      -- ce:    a counter example to `x = e`
      -- done:  known improvements
      -- extra: known sub-goals
      -- todo:  more work to process once we are done with `x`.
      goV _  done extra todo _ _ [] = go done extra (reverse todo)
      goV ce done extra todo x e ((y,e') : more)
        -- x = y?
        | e == e' = do yesK <- cryMustEqualV solver x y
                       if yesK then go (Map.insert x (Var y) done)
                                       extra
                                       (reverse todo ++ more)
                               else tryLR

        | otherwise = tryLR

        where
        next = goV ce done extra ((y,e'):todo) x e more

        tryLR =
          do mb <- tryLR_with x e y e'
             case mb of
               Just (r,subGoals) -> go (Map.insert x r done)
                                       (subGoals ++ extra)
                                       (reverse todo ++ more)
               Nothing ->
                 do mb1 <- tryLR_with y e' x e
                    case mb1 of
                      Nothing -> next
                      Just (r, subGoals) -> go (Map.insert y r done)
                                               (subGoals ++ extra)
                                               (reverse todo ++ more)

        tryLR_with v1 v1Expr v2 v2Expr =
          case ( isUniVar v1
               , v1Expr
               , v2Expr
               , Map.lookup v1 ce
               , Map.lookup v2 ce
               ) of
            (True, x1, y1, Just x2, Just y2) ->
              cryCheckLinRel solver logger v2 v1 (y1,x1) (y2,x2)
            _ -> return Nothing

    -- | Try to prove the given expression.
    checkUnsat :: SMT.Solver -> SExpr -> IO Bool
    checkUnsat s e =
      do SMT.push s
         SMT.assert s e
         res <- SMT.check s
         SMT.pop s
         return (res == SMT.Unsat)

    -- | Try to prove the given expression.
    -- If we fail, we try to give a counter example.
    -- If the answer is unknown, then we return an empty counter example.
    getCounterExample :: SMT.Solver -> [Name] -> SExpr -> IO (Maybe (Map Name Nat'))
    getCounterExample s xs e =
      do SMT.push s
         SMT.assert s e
         res <- SMT.check s
         ans <- case res of
                  SMT.Unsat   -> return Nothing
                  SMT.Unknown -> return (Just Map.empty)
                  SMT.Sat     -> Just `fmap` getVals s xs
         SMT.pop s
         return ans

    -- | Is this the only possible value for the constant, under the current
    -- assumptions.
    -- Assumes that we are in a 'Sat' state.
    -- Returns 'Nothing' if the variables must always match the given value.
    -- Otherwise, we return a counter-example (which may be empty, if uniknown)
    cryMustEqualK :: SMT.Solver -> [Name] ->
                     Name -> Nat' -> IO (Maybe (Map Name Nat'))
    cryMustEqualK solver xs x val =
      case val of
        Inf   -> getCounterExample solver xs (SMT.const (smtFinName x))
        Nat n -> getCounterExample solver xs $
                 SMT.not $
                 SMT.and (SMT.const $ smtFinName x)
                         (SMT.eq (SMT.const (smtName x)) (SMT.int n))

    -- | Do these two variables need to always be the same, under the current
    -- assumptions.
    -- Assumes that we are in a 'Sat' state.
    cryMustEqualV :: SMT.Solver -> Name -> Name -> IO Bool
    cryMustEqualV solver x y =
         checkUnsat solver $
            SMT.not $
            SMT.or (SMT.not (fin x) `SMT.and` SMT.not (fin y))
                   (fin x `SMT.and` fin y `SMT.and` SMT.eq (var x) (var y))

      where fin a = SMT.const (smtFinName a)
            var a = SMT.const (smtName a)

    -- | Try to find a linear relation between the two variables, based
    -- on two observed data points.
    -- NOTE:  The variable being defined is the SECOND name.
    cryCheckLinRel :: SMT.Solver -> SMT.Logger ->
             {- x -} Name {- ^ Definition in terms of this variable. -} ->
             {- y -} Name {- ^ Define this variable. -} ->
                     (Nat',Nat') {- ^ Values in one model (x,y) -} ->
                     (Nat',Nat') {- ^ Values in another model (x,y) -} ->
                     IO (Maybe (Expr,[Prop]))
                     {- ^ Either nothing, or an improving expression, and any
                          additional obligations -}
    cryCheckLinRel s logger x y p1 p2 =
      -- First, try to find a linear relation that holds in all finite cases.
      do SMT.push s
         SMT.assert s (isFin x)
         SMT.assert s (isFin y)
         ans <- case (p1,p2) of
                  ((Nat x1, Nat y1), (Nat x2, Nat y2)) ->
                      checkLR x1 y1 x2 y2

                  ((Inf,    Inf),    (Nat x2, Nat y2)) ->
                     mbGoOn (getFinPt x2) $ \(x1,y1) -> checkLR x1 y1 x2 y2

                  ((Nat x1, Nat y1), (Inf,    Inf)) ->
                     mbGoOn (getFinPt x1) $ \(x2,y2) -> checkLR x1 y1 x2 y2

                  _ -> return Nothing

         SMT.pop s

         -- Next, check the infinite cases: if @y = A * x + B@, then
         -- either both @x@ and @y@ must be infinite or they both must be finite.
         -- Note that we don't consider relations where A = 0: because they
         -- would be handled when we checked that @y@ is a constant.
         case ans of
           Nothing -> return Nothing
           Just e ->
             do SMT.push s
                SMT.assert s (SMT.not (SMT.eq (isFin x) (isFin y)))
                c <- SMT.check s
                SMT.pop s
                case c of
                  SMT.Unsat -> return (Just e)
                  _         -> return Nothing

      where
      isFin a = SMT.const (smtFinName a)

      -- XXX: Duplicates `cryDefined`
      -- The constraints are always of the form: x >= K, or K >= x
      wellDefined e =
        case e of
          (K (Nat a) :* t) :- K (Nat b) ->
            let c = div (b + a - 1) a
            in [ t :>= K (Nat c) ]

          K (Nat b) :- (K (Nat a) :* t)
            -> let c = div b a
               in [ K (Nat c) :>= t ]

          a  :- b -> [ a :>= b ]

          _ -> []

      checkLR x1 y1 x2 y2 =
        do SMT.logMessage logger ("checkLR: " ++ show (x1,y1) ++ "   "
                                              ++ show (x2,y2))
           mbGoOn (return (linRel (x1,y1) (x2,y2))) (\(a,b) ->
             do let sumTerm v
                       | b == 0    = v
                       | b < 0     = v :- K (Nat (negate b))
                       | otherwise = v :+ K (Nat b)

                    expr
                      | a == 1     = sumTerm (Var x)
                      | a <  0     = K (Nat b) :- K (Nat (negate a)) :* Var x
                      | otherwise  = sumTerm (K (Nat a) :* Var x)

                SMT.logMessage logger ("candidate: " ++ show (ppProp (Var y :==: expr)))

                proved <- checkUnsat s
                        $ propToSmtLib $ crySimplify
                        $ Not $ Var y :==: expr

                if not proved
                   then SMT.logMessage logger "failed" >> return Nothing
                   else return (Just (expr,wellDefined expr)))

      -- Try to get an example of another point, which is finite, and at
      -- different @x@ coordinate.
      getFinPt otherX =
        do SMT.push s
           SMT.assert s (SMT.not (SMT.eq (SMT.const (smtName x)) (SMT.int otherX)))
           smtAns <- SMT.check s
           ans <- case smtAns of
                    SMT.Sat ->
                      do vX <- SMT.getConst s (smtName x)
                         vY <- SMT.getConst s (smtName y)
                         case (vX, vY) of
                           (SMT.Int vx, SMT.Int vy)
                              | vx >= 0 && vy >= 0 -> return (Just (vx,vy))
                           _ -> return Nothing
                    _ -> return Nothing
           SMT.pop s
           return ans

      mbGoOn m k = do ans <- m
                      case ans of
                        Nothing -> return Nothing
                        Just a  -> k a

    {- | Compute a linear relation through two concrete points.
    Try to find a relation of the form @y = a * x + b@.
    Depending on the signs of @a@ and @b@, we need additional checks,
    to ensure tha @a * x + b@ is valid.

    y1 = A * x1 + B
    y2 = A * x2 + B
    (y2 - y1) = A * (x2 - x1)

    A = (y2 - y1) / (x2 - x1)
    B = y1 - A * x1
    -}
    linRel :: (Integer,Integer)       {- ^ First point -} ->
              (Integer,Integer)       {- ^ Second point -} ->
              Maybe (Integer,Integer) {- ^ (A,B) -}
    linRel (x1,y1) (x2,y2) =
      do guard (x1 /= x2)
         let (a,r) = divMod (y2 - y1) (x2 - x1)
         guard (r == 0 && a /= 0)    -- Not interested in A = 0
         let b = y1 - a * x1
         guard $ not $ a < 0 && b < 0   -- No way this will give a natural number.
         return (a,b)