content: {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE PatternSynonyms            #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  -- | Functions to make environments easier to read
  module Language.Fixpoint.Solver.Prettify (savePrettifiedQuery) where
  
  import           Data.Bifunctor (first)
  import           Data.HashMap.Lazy (HashMap)
  import qualified Data.HashMap.Lazy as HashMap
  import           Data.HashSet (HashSet)
  import qualified Data.HashSet as HashSet
  import           Data.List (intersperse, sortOn)
  import           Data.Maybe (fromMaybe)
  import           Data.Text (Text)
  import qualified Data.Text as Text
  import           Language.Fixpoint.Misc (ensurePath)
  import           Language.Fixpoint.Solver.EnvironmentReduction
    ( dropLikelyIrrelevantBindings
    , inlineInSortedReft
    , mergeDuplicatedBindings
    , simplifyBooleanRefts
    , undoANFAndVV
    )
  import           Language.Fixpoint.Types.Config (Config, queryFile)
  import           Language.Fixpoint.Types.Constraints
  import           Language.Fixpoint.Types.Environments
    (BindEnv, elemsIBindEnv, lookupBindEnv)
  import           Language.Fixpoint.Types.Names
    ( Symbol
    , dropPrefixOfSym
    , prefixOfSym
    , suffixOfSym
    , suffixSymbol
    , symbol
    , symbolText
    )
  import           Language.Fixpoint.Types.PrettyPrint
  import           Language.Fixpoint.Types.Refinements
    ( ExprV(..)
    , pattern PFalse
    , Reft
    , SortedReft(..)
    , conjuncts
    , expr
    , reft
    , reftBind
    , reftPred
    , sortedReftSymbols
    , substf
    , substSortInExpr
    )
  import           Language.Fixpoint.Types.Sorts (Sort(..), substSort)
  import           Language.Fixpoint.Types.Substitutions (exprSymbolsSet)
  import qualified Language.Fixpoint.Utils.Files as Files
  import           System.FilePath (addExtension)
  import           Text.PrettyPrint.HughesPJ.Compat
    (Doc, braces, hang, nest, render, sep, text, vcat, (<+>), ($+$))
  
  
  savePrettifiedQuery :: Fixpoint a => Config -> FInfo a -> IO ()
  savePrettifiedQuery cfg info = do
    let fq   = queryFile Files.Fq cfg `addExtension` "prettified"
    putStrLn $ "Saving prettified Query: "   ++ fq ++ "\n"
    ensurePath fq
    writeFile fq $ render (prettyConstraints info)
  
  prettyConstraints :: Fixpoint a => FInfo a -> Doc
  prettyConstraints info =
    vcat $
    map
      (prettyConstraint (bs info) . snd)
      (sortOn fst $ HashMap.toList (cm info))
  
  prettyConstraint
    :: Fixpoint a
    => BindEnv a
    -> SubC a
    -> Doc
  prettyConstraint bindEnv c =
    let env = [ (s, ([(bId, a)], sr))
              | bId <- elemsIBindEnv $ senv c
              , let (s, sr, a) = lookupBindEnv bId bindEnv
              ]
        mergedEnv = mergeDuplicatedBindings env
        undoANFEnv = undoANFAndVV mergedEnv
        boolSimplEnvDiff = simplifyBooleanRefts undoANFEnv
        boolSimplEnv = HashMap.map snd $ HashMap.union boolSimplEnvDiff undoANFEnv
  
        simplifiedLhs = simplify $ inlineInSortedReft (`HashMap.lookup` boolSimplEnv) (slhs c)
        simplifiedRhs = simplify $ inlineInSortedReft (`HashMap.lookup` boolSimplEnv) (srhs c)
  
        prunedEnv =
          if expr simplifiedRhs /= PFalse then
            dropLikelyIrrelevantBindings
              (constraintSymbols simplifiedLhs simplifiedRhs)
              boolSimplEnv
          else
            boolSimplEnv
        (renamedEnv, c') =
          shortenVarNames prunedEnv c { slhs = simplifiedLhs, srhs = simplifiedRhs }
        prettyEnv =
          sortOn bindingSelector $
          eraseUnusedBindings (constraintSymbols (slhs c') (srhs c')) renamedEnv
     in hang (text "\n\nconstraint:") 2 $
            text "lhs" <+> toFix (slhs c')
        $+$ text "rhs" <+> toFix (srhs c')
        $+$ pprId (sid c) <+> text "tag" <+> toFix (stag c)
        $+$ toFixMeta (text "constraint" <+> pprId (sid c)) (toFix (sinfo c))
        $+$ hang (text "environment:") 2
              (vcat $ intersperse "" $ elidedMessage : map prettyBind prettyEnv)
    where
      prettyBind (s, sr) = sep [toFix s <+> ":", nest 2 $ prettySortedRef sr]
      prettySortedRef sr = braces $ sep
        [ toFix (reftBind (sr_reft sr)) <+> ":" <+> toFix (sr_sort sr) <+> "|"
        , nest 2 $ toFix $ conjuncts $ reftPred $ sr_reft sr
        ]
  
      elidedMessage = "// elided some likely irrelevant bindings"
  
      constraintSymbols sr0 sr1 =
        HashSet.union (sortedReftSymbols sr0) (sortedReftSymbols sr1)
  
      -- Sort by symbol then reft
      bindingSelector (s, sr) =
        ( -- put unused symbols at the end
          if "_" `Text.isPrefixOf` s then "{" <> s else s
        , reftBind (sr_reft sr)
        , sr_sort sr
        , reftPred $ sr_reft sr
        )
  
  pprId :: Show a => Maybe a -> Doc
  pprId (Just i)  = "id" <+> text (show i)
  pprId _         = ""
  
  toFixMeta :: Doc -> Doc -> Doc
  toFixMeta k v = text "// META" <+> k <+> text ":" <+> v
  
  -- | @eraseUnusedBindings ss env@ prefixes with @_@ the symbols that aren't
  -- refered from @ss@ or refinement types in the environment.
  eraseUnusedBindings
    :: HashSet Symbol -> [(Symbol, SortedReft)] -> [(Text, SortedReft)]
  eraseUnusedBindings ss env = map (first reachable) env
    where
      allSymbols = HashSet.union ss envSymbols
      envSymbols =
        HashSet.unions $
        map (exprSymbolsSet . reftPred . sr_reft . snd) env
  
      reachable s =
        if s `HashSet.member` allSymbols then
          symbolText s
        else
          "_" <> symbolText s
  
  -- | Shortens the names of refinements types in a given environment
  -- and constraint
  shortenVarNames
    :: HashMap Symbol SortedReft
    -> SubC a
    -> ([(Symbol, SortedReft)], SubC a)
  shortenVarNames env subc =
    let bindsRenameMap = proposeRenamings $ HashMap.keys env
        env' = map (renameBind bindsRenameMap) (HashMap.toList env)
     in
        (env', renameSubC bindsRenameMap subc)
    where
      renameSubC :: HashMap Symbol Symbol -> SubC a -> SubC a
      renameSubC symMap c =
        mkSubC
          (senv c)
          (renameSortedReft symMap $ slhs c)
          (renameSortedReft symMap $ srhs c)
          (sid c)
          (stag c)
          (sinfo c)
  
      renameBind
        :: HashMap Symbol Symbol -> (Symbol, SortedReft) -> (Symbol, SortedReft)
      renameBind symMap (s, sr) =
        (at symMap s, renameSortedReft symMap sr)
  
      renameSortedReft
        :: HashMap Symbol Symbol -> SortedReft -> SortedReft
      renameSortedReft symMap (RR t r) =
        let sortSubst = FObj . at symMap
         in RR (substSort sortSubst t) (renameReft symMap r)
  
      renameReft :: HashMap Symbol Symbol -> Reft -> Reft
      renameReft symMap r =
        let m = HashMap.insert (reftBind r) (prefixOfSym $ reftBind r) symMap
            sortSubst = FObj . at symMap
         in reft (at m (reftBind r)) $
              substSortInExpr sortSubst $
              substf (EVar . at m) (reftPred r)
  
      at :: HashMap Symbol Symbol -> Symbol -> Symbol
      at m k = fromMaybe k $ HashMap.lookup k m
  
  -- | Given a list of symbols with no duplicates, it proposes shorter
  -- names to use for them.
  --
  -- All proposals preserve the prefix of the original name. A prefix is
  -- the longest substring containing the initial character and no separator
  -- ## ('symSepName').
  --
  -- Suffixes are preserved, except when symbols with a given prefix always
  -- use the same suffix.
  --
  proposeRenamings :: [Symbol] -> HashMap Symbol Symbol
  proposeRenamings = toSymMap . toPrefixSuffixMap
  
  -- | Indexes symbols by prefix and then by suffix. If the prefix
  -- equals the symbol, then the suffix is the empty symbol.
  --
  -- For instance,
  --
  -- > toPrefixSuffixMap ["a##b##c"] ! "a" ! "c" == ["a##b##c"]
  -- > toPrefixSuffixMap ["a"] ! "a" ! "" == ["a"]
  --
  -- In general,
  --
  -- > forall ss.
  -- > Set.fromList ss == Set.fromList $ concat [ xs | m <- elems (toPrefixSuffixMap ss), xs <- elems m ]
  --
  -- > forall ss.
  -- > and [ all (pfx `isPrefixOfSym`) xs && all (sfx `isSuffixOfSym`) xs
  -- >     | (pfx, m) <- toList (toPrefixSuffixMap ss)
  -- >     , (sfx, xs) <- toList m
  -- >     ]
  --
  -- TODO: put the above in unit tests
  --
  toPrefixSuffixMap :: [Symbol] -> HashMap Symbol (HashMap Symbol [Symbol])
  toPrefixSuffixMap xs = HashMap.fromListWith (HashMap.unionWith (++))
    [ (pfx, HashMap.singleton sfx [s])
    | s <- xs
    , let pfx = prefixOfSym s
          sfx = suffixOfSym $ dropPrefixOfSym s
    ]
  
  -- | Suggests renamings for every symbol in a given prefix/suffix map.
  toSymMap :: HashMap Symbol (HashMap Symbol [Symbol]) -> HashMap Symbol Symbol
  toSymMap prefixMap = HashMap.fromList
    [ r
    | (pfx, h) <- HashMap.toList prefixMap
    , r <- renameWithSuffixes pfx (HashMap.toList h)
    ]
    where
      renameWithSuffixes pfx = \case
        [(_sfx, ss)] -> renameWithAppendages pfx ("", ss)
        sfxs -> concatMap (renameWithAppendages pfx) sfxs
  
      renameWithAppendages pfx (sfx, xs) = zip xs $ case xs of
        [_s] -> [pfx `suffixIfNotNull` sfx]
        ss -> zipWith (rename pfx sfx) [1 :: Integer ..] ss
  
      rename pfx sfx i _s =
        pfx `suffixIfNotNull` sfx `suffixSymbol` symbol (show i)
  
      suffixIfNotNull a b =
        if Text.null (symbolText b) then a else a `suffixSymbol` b
  