content: {-# LANGUAGE DerivingVia #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE NoMonomorphismRestriction #-}
  
  -- TODO: the calls to 'error' should be replaced with logging/error capabilities.
  
  -- |
  -- Module      :  Disco.Pretty
  -- Copyright   :  disco team and contributors
  -- Maintainer  :  byorgey@gmail.com
  --
  -- SPDX-License-Identifier: BSD-3-Clause
  --
  -- Various pretty-printing facilities for disco.
  module Disco.Pretty (
    module Disco.Pretty.DSL,
    module Disco.Pretty,
    module Disco.Pretty.Prec,
    Doc,
  )
  where
  
  import Prelude hiding ((<>))
  
  import Data.Bifunctor
  import Data.Char (isAlpha)
  import Data.List.NonEmpty (NonEmpty)
  import qualified Data.List.NonEmpty as NE
  import Data.Map (Map)
  import qualified Data.Map as M
  import Data.Ratio
  import Data.Set (Set)
  import qualified Data.Set as S
  import Disco.Effects.LFresh
  import Disco.Pretty.DSL
  import Disco.Pretty.Prec
  import Disco.Syntax.Operators
  import Polysemy
  import Polysemy.Reader
  import Prettyprinter (Doc)
  import Unbound.Generics.LocallyNameless (Name)
  
  ------------------------------------------------------------
  -- Utilities for handling precedence and associativity
  
  -- | Convenience function combining 'setPA' and 'mparens', since we
  --   often want to simultaneously indicate what the precedence and
  --   associativity of a term is, and optionally surround it with
  --   parentheses depending on the precedence and associativity of its
  --   parent.
  withPA :: Member (Reader PA) r => PA -> Sem r (Doc ann) -> Sem r (Doc ann)
  withPA pa = mparens pa . setPA pa
  
  -- | Locally set the precedence and associativity within a
  --   subcomputation.
  setPA :: Member (Reader PA) r => PA -> Sem r a -> Sem r a
  setPA = local . const
  
  -- | Mark a subcomputation as pretty-printing a term on the left of an
  --   operator (so parentheses can be inserted appropriately, depending
  --   on the associativity).
  lt :: Member (Reader PA) r => Sem r (Doc ann) -> Sem r (Doc ann)
  lt = local (\(PA p _) -> PA p InL)
  
  -- | Mark a subcomputation as pretty-printing a term on the right of
  --   an operator (so parentheses can be inserted appropriately,
  --   depending on the associativity).
  rt :: Member (Reader PA) r => Sem r (Doc ann) -> Sem r (Doc ann)
  rt = local (\(PA p _) -> PA p InR)
  
  -- | Optionally surround a pretty-printed term with parentheses,
  --   depending on its precedence and associativity (given as the 'PA'
  --   argument) and that of its context (given by the ambient 'Reader
  --   PA' effect).
  mparens :: Member (Reader PA) r => PA -> Sem r (Doc ann) -> Sem r (Doc ann)
  mparens pa doc = do
    parentPA <- ask
    (if pa `lowerPrec` parentPA then parens else id) doc
  
  ------------------------------------------------------------
  -- Pretty type class
  
  class Pretty t where
    pretty :: Members '[Reader PA, LFresh] r => t -> Sem r (Doc ann)
  
  prettyStr :: Pretty t => t -> Sem r String
  prettyStr = renderDoc . runLFresh . pretty
  
  pretty' :: Pretty t => t -> Sem r (Doc ann)
  pretty' = runReader initPA . runLFresh . pretty
  
  ------------------------------------------------------------
  -- Some standard instances
  
  instance Pretty a => Pretty [a] where
    pretty = brackets . intercalate "," . map pretty
  
  instance Pretty a => Pretty (NonEmpty a) where
    pretty = pretty . NE.toList
  
  instance (Pretty k, Pretty v) => Pretty (Map k v) where
    pretty m = do
      let es = map (\(k, v) -> pretty k <+> "->" <+> pretty v) (M.assocs m)
      ds <- setPA initPA $ punctuate "," es
      braces (hsep ds)
  
  instance Pretty a => Pretty (Set a) where
    pretty = braces . intercalate "," . map pretty . S.toList
  
  ------------------------------------------------------------
  -- Some Disco instances
  
  instance Pretty (Name a) where
    pretty = text . show
  
  instance Pretty TyOp where
    pretty = \case
      Enumerate -> text "enumerate"
      Count -> text "count"
  
  -- | Pretty-print a unary operator, by looking up its concrete syntax
  --   in the 'uopMap'.
  instance Pretty UOp where
    pretty op = case M.lookup op uopMap of
      Just (OpInfo _ (syn : _) _) ->
        text $ syn ++ (if all isAlpha syn then " " else "")
      _ -> error $ "UOp " ++ show op ++ " not in uopMap!"
  
  -- | Pretty-print a binary operator, by looking up its concrete syntax
  --   in the 'bopMap'.
  instance Pretty BOp where
    pretty (Should op) = pretty op <> "!"
    pretty op = case M.lookup op bopMap of
      Just (OpInfo _ (syn : _) _) -> text syn
      _ -> error $ "BOp " ++ show op ++ " not in bopMap!"
  
  --------------------------------------------------
  -- Pretty-printing rationals + decimals
  
  -- | Pretty-print a rational number as a fraction.
  prettyRational :: Rational -> String
  prettyRational r
    | denominator r == 1 = show (numerator r)
    | otherwise = show (numerator r) ++ "/" ++ show (denominator r)
  
  -- | Pretty-print a rational number using its decimal expansion, in
  --   the format @nnn.prefix[rep]...@, with any repeating digits enclosed
  --   in square brackets.
  prettyDecimal :: Rational -> String
  prettyDecimal r = wholePart ++ "." ++ decimalPart
   where
    (n', d') = properFraction r :: (Integer, Rational)
    d = abs d'
    n = abs n'
    (expan, len) = digitalExpansion 10 (numerator d) (denominator d)
    wholePart = (if d' < 0 then "-" else "") ++ show n
    decimalPart
      | length first102 > 101 || length first102 == 101 && last first102 /= 0 =
          concatMap show (take 100 expan) ++ "..."
      | rep == [0] =
          if null pre then "0" else concatMap show pre
      | otherwise =
          concatMap show pre ++ "[" ++ concatMap show rep ++ "]"
     where
      (pre, rep) = splitAt len expan
      first102 = take 102 expan
  
  -- Given a list, find the indices of the list giving the first and
  -- second occurrence of the first element to repeat, or Nothing if
  -- there are no repeats.
  findRep :: Ord a => [a] -> ([a], Int)
  findRep = findRep' M.empty 0
  
  findRep' :: Ord a => M.Map a Int -> Int -> [a] -> ([a], Int)
  findRep' _ _ [] = error "Impossible. Empty list in findRep'"
  findRep' prevs ix (x : xs)
    | x `M.member` prevs = ([], prevs M.! x)
    | otherwise = first (x :) $ findRep' (M.insert x ix prevs) (ix + 1) xs
  
  -- | @digitalExpansion b n d@ takes the numerator and denominator of a
  --   fraction n/d between 0 and 1, and returns a pair of (1) a list of
  --   digits @ds@, and (2) a nonnegative integer k such that @splitAt k
  --   ds = (prefix, rep)@, where the infinite base-b expansion of
  --   n/d is 0.@(prefix ++ cycle rep)@.  For example,
  --
  --   > digitalExpansion 10 1 4  = ([2,5,0], 2)
  --   > digitalExpansion 10 1 7  = ([1,4,2,8,5,7], 0)
  --   > digitalExpansion 10 3 28 = ([1,0,7,1,4,2,8,5], 2)
  --   > digitalExpansion 2  1 5  = ([0,0,1,1], 0)
  --
  --   It works by performing the standard long division algorithm, and
  --   looking for the first time that the remainder repeats.
  digitalExpansion :: Integer -> Integer -> Integer -> ([Integer], Int)
  digitalExpansion b n d = digits
   where
    longDivStep (_, r) = (b * r) `divMod` d
    res = drop 1 $ iterate longDivStep (0, n)
    digits = first (map fst) (findRep res)
  