{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}

-- | This module implements the @M^X@ functor, where @M@ is a monoid that
-- doesn't have to be a group like in "Copar.Functors.GroupValued".
--
-- The 'RefinementInterface' implementation for such functors doesn't fulfil the
-- same runtime complexity criteria as the other functors and it also uses tons
-- of space, but it works and satisfies the axioms.
module Copar.Functors.MonoidValued
  ( SlowMonoidValued(..)
  , maxIntValued
  , minIntValued
  , maxRealValued
  , minRealValued
  , andWordValued 
  , orWordValued
  )
where

import           Data.List                      ( foldl', intersperse )
import           Data.Semigroup                 ( Max(..), Min(..) )
import           Data.Functor.Classes
import           Control.Monad
import           Data.Foldable
import           Data.Word

import qualified Data.Vector                   as V
import           Text.Megaparsec
import qualified Data.Text.Prettyprint         as Doc
import           Data.Text.Prettyprint          ( (<+>) )
import           Data.Text                      ( Text )

import qualified Data.Vector.Utils             as V
import           Copar.RefinementInterface
import           Copar.FunctorDescription
import qualified Copar.Parser.Lexer               as L
import           Copar.FunctorExpression.Parser
import           Copar.Coalgebra.Parser
import           Data.Float.Utils               ( MaxDouble(..)
                                                , MinDouble(..)
                                                )
import           Copar.Parser.Types
import           Data.SumBag (SumBag)
import qualified Data.SumBag as SumBag
import           Data.Bits.Monoid


newtype SlowMonoidValued m a = SlowMonoidValued a

instance Eq1 (SlowMonoidValued m) where
  liftEq f (SlowMonoidValued a1) (SlowMonoidValued a2) = f a1 a2

instance Show1 (SlowMonoidValued m) where
  liftShowsPrec f _ i (SlowMonoidValued a) = f i a

deriving instance Show (SlowMonoidValued m ())
deriving instance Functor (SlowMonoidValued m)
deriving instance Foldable (SlowMonoidValued m)
deriving instance Traversable (SlowMonoidValued m)

data MonoidValuedDescription m = MonoidValued
  { mvName :: Text
  , mvDescription :: Text
  , mvSet :: [Text] -- often [ascii, unicode]
  , mvOperation :: Text
  , mvTerm :: Text
  }

makeMonoidValued
  :: MonoidValuedDescription m -> FunctorDescription (SlowMonoidValued m)
makeMonoidValued desc = FunctorDescription
  { name = mvName desc <> "-valued"
  , syntaxExample = fold (intersperse " | " (map syntax (mvSet desc)))
  , description = Just (makeMVHelp desc)
  , functorExprParser =
    prefix
      -- We need this try here, so that parenthesis can still be parsed as
      -- normal if they don't contain exactly (Z, max)
      (  try
          (L.parens
            ((choice (map L.symbol (mvSet desc)))
            >> L.comma
            >> L.symbol (mvOperation desc)
            )
          )
      >> L.symbol "^"
      >> pure SlowMonoidValued
      )
  }
  where syntax x = functorSyntax x (mvOperation desc)

functorSyntax :: Text -> Text -> Text
functorSyntax s o = "(" <> s <> ", " <> o <> ")^X"

makeMVHelp :: MonoidValuedDescription m -> Doc.Doc Doc.AnsiStyle
makeMVHelp desc =
  Doc.reflow ("Weighted systems with " <> mvDescription desc)
  <> Doc.line <> Doc.line
  <> Doc.reflow "This is similar to all the group valued functors (int, real, \
                \etc) but isn't actually a group as it lacks an inverse. Thus \
                \the refinement interface implementation for this functor is\
                \asymptotically slower than for the others."
  <> Doc.line <> Doc.line
  <> Doc.annotate Doc.bold "Functor syntax:"
  <+> Doc.reflow (fold (intersperse " | " (map syntax (mvSet desc))))
  <> Doc.line <> Doc.line
  <> Doc.annotate Doc.bold "Coalgebra syntax:"
  <+> Doc.reflow ("'{' X ':' " <> mvTerm desc <> ", ... '}'")
  where syntax x = functorSyntax x (mvOperation desc)

-- | The @(ℤ, max)^X@ functor
maxIntValued :: FunctorDescription (SlowMonoidValued (Max Int))
maxIntValued = makeMonoidValued $ MonoidValued
  { mvName = "Max"
  , mvSet = ["Z", "ℤ"]
  , mvOperation = "max"
  , mvDescription = "the monoid (Z, max)"
  , mvTerm = "int"
  }

-- | The @(ℤ, min)^X@ functor
minIntValued :: FunctorDescription (SlowMonoidValued (Min Int))
minIntValued = makeMonoidValued $ MonoidValued
  { mvName = "Min"
  , mvSet = ["Z", "ℤ"]
  , mvOperation = "min"
  , mvDescription = "the monoid (Z, min)"
  , mvTerm = "int"
  }


-- | The @(ℝ, max)^X@ functor
maxRealValued :: FunctorDescription (SlowMonoidValued MaxDouble)
maxRealValued = makeMonoidValued $ MonoidValued
  { mvName = "Max"
  , mvSet = ["R", "ℝ"]
  , mvOperation = "max"
  , mvDescription = "the monoid (R, max)"
  , mvTerm = "real"
  }


-- | The @(ℝ, max)^X@ functor
minRealValued :: FunctorDescription (SlowMonoidValued MinDouble)
minRealValued = makeMonoidValued $ MonoidValued
  { mvName = "Min"
  , mvSet = ["R", "ℝ"]
  , mvOperation = "min"
  , mvDescription = "the monoid (R, min)"
  , mvTerm = "real"
  }


-- | The @(Word, and)^X@ functor
andWordValued :: FunctorDescription (SlowMonoidValued (BitAnd Word64))
andWordValued = makeMonoidValued $ MonoidValued
  { mvName = "BitAnd"
  , mvSet = ["Word"]
  , mvOperation = "and"
  , mvDescription = "bitvectors and bitwise 'and' as monoid weight"
  , mvTerm = "0xCAFE"
  }


-- | The @(Word, or)^X@ functor
orWordValued :: FunctorDescription (SlowMonoidValued (BitOr Word64))
orWordValued = makeMonoidValued $ MonoidValued
  { mvName = "BitOr"
  , mvSet = ["Word"]
  , mvOperation = "or"
  , mvDescription = "bitvectors and bitwise 'or' as monoid weight"
  , mvTerm = "0xCAFE"
  }


type instance Label (SlowMonoidValued m) = m
type instance Weight (SlowMonoidValued m) = (m, SumBag m)
type instance F1 (SlowMonoidValued m) = m
type instance F3 (SlowMonoidValued m) = (m, m, m)

instance (Monoid m, Ord m) => RefinementInterface (SlowMonoidValued m) where
  init
    :: F1 (SlowMonoidValued m)
    -> [Label (SlowMonoidValued m)]
    -> Weight (SlowMonoidValued m)
  init _ labels =
    (mempty, foldl' (flip SumBag.insert) SumBag.empty labels)

  update
    :: [Label (SlowMonoidValued m)]
    -> Weight (SlowMonoidValued m)
    -> ( Weight (SlowMonoidValued m)
       , F3 (SlowMonoidValued m)
       , Weight (SlowMonoidValued m)
       )
  update labels (sumRest, counts) =
    let toS  = foldl' (flip SumBag.insert) SumBag.empty labels
        toCWithoutS = foldl' (flip SumBag.delete) counts labels
        sumS = fold toS
        sumCWithoutS = fold toCWithoutS
        f3   = (sumRest, sumCWithoutS, sumS)
        w1   = (sumRest <> sumCWithoutS, toS)
        w2   = (sumRest <> sumS, toCWithoutS)
    in  (w1, f3, w2)


instance ParseMorphism (SlowMonoidValued (Max Int)) where
  parseMorphismPoint = parseMorphismPointHelper (Max <$> (L.signed L.decimal))

instance ParseMorphism (SlowMonoidValued (Min Int)) where
  parseMorphismPoint = parseMorphismPointHelper (Min <$> (L.signed L.decimal))

instance ParseMorphism (SlowMonoidValued MaxDouble) where
  parseMorphismPoint = parseMorphismPointHelper (MaxDouble <$> L.signed L.float)

instance ParseMorphism (SlowMonoidValued MinDouble) where
  parseMorphismPoint = parseMorphismPointHelper (MinDouble <$> L.signed L.float)

instance ParseMorphism (SlowMonoidValued (BitAnd Word64)) where
  parseMorphismPoint = parseMorphismPointHelper L.hex

instance ParseMorphism (SlowMonoidValued (BitOr Word64)) where
  parseMorphismPoint = parseMorphismPointHelper L.hex


parseMorphismPointHelper :: (HasSanityChecks m, MonadParser m, Ord x, Monoid w) => m w -> SlowMonoidValued w (m x) -> m (w, V.Vector (x, w))
parseMorphismPointHelper weightParser (SlowMonoidValued inner) = do
  sanity <- not <$> noSanityChecks

  !successors <- case sanity of
    True -> do
      succs <- V.sortOn fst . V.fromList <$> L.braces (edge `sepBy` L.comma)

      when (V.hasDuplicates (fmap fst succs)) $
        fail "monoid valued: Duplicate edges"

      return succs

    False -> V.fromList <$> L.braces (edge `sepBy` L.comma)

  let !f1 = fold (V.map snd successors)
  return (f1, successors)
  where edge = (,) <$> inner <*> (L.colon *> weightParser)