{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}

-- |
-- Module      : Data.Diagram.Simple
-- Copyright   : (c) Eddie Jones 2020
-- License     : BSD-3
-- Maintainer  : eddiejones2108@gmail.com
-- Stability   : experimental
--
-- Reduced Ordered Binary Decision Diagrams
module Data.Diagram.Simple
  ( -- * Diagram
    Diagram,
    runDiagram,

    -- * Propositions
    Prop,
    mkProp,
    true,
    false,
    atom,
    notAtom,
    mapAtom,
    bindAtom,

    -- * Combinations
    Data.Diagram.Simple.and,
    Data.Diagram.Simple.or,
    Data.Diagram.Simple.not,
    restrict,

    -- * Summary
    fold,
    anySat,
  )
where

import Control.Monad.State
import qualified Data.Diagram as D
import Data.Functor.Classes
import Data.Functor.Identity
import Data.Hashable
import Data.Hashable.Lifted
import qualified Data.Map as M

data Memo l s
  = And (Prop l s) (Prop l s)
  | Or (Prop l s) (Prop l s)
  | Not (Prop l s)
  | Restr l Bool (Prop l s)
  deriving (Eq, Ord)

-- | A binary decision diagram
newtype Diagram l s a = Diagram
  { unDiag :: D.Diagram (Node l) Bool s (State (M.Map (Memo l s) (Prop l s))) a
  }
  deriving (Functor, Applicative, Monad)

-- | Extract non-diagrammatic information
runDiagram :: (forall s. Diagram l s a) -> a
runDiagram d = runIdentity $ D.runDiagram ((`evalState` M.empty) <$> D.compress (unDiag d))

-- | A diagramatic proposition build on atomic propositions of type @l@
newtype Prop l s = Prop
  { unProp :: D.Free (Node l) Bool s
  }
  deriving (Eq, Ord)

-- | An internal node of the diagram
data Node l k = Node
  { label :: l,
    lo :: k,
    hi :: k
  }
  deriving (Functor, Foldable, Traversable)

instance Eq l => Eq1 (Node l) where
  liftEq eq n m = label n == label m && eq (lo n) (lo m) && eq (hi n) (hi m)

instance Hashable l => Hashable1 (Node l) where
  liftHashWithSalt l s n = hashWithSalt s (label n, l s (lo n), l s (hi n))

-- | Make a proposition (if not already present) from it's lo and hi cases
mkProp :: (Eq l, Hashable l) => l -> Prop l s -> Prop l s -> Diagram l s (Prop l s)
mkProp l lo hi
  | lo == hi = return lo
  | otherwise = Diagram (Prop <$> D.free Node {label = l, lo = unProp lo, hi = unProp hi})

-- Map a function over a proposition but push the result through mkProp
mapProp :: (Eq l, Hashable l) => (Bool -> Bool) -> Prop l s -> Diagram l s (Prop l s)
mapProp f (Prop p) = Diagram $ D.fold (\n -> unDiag $ mkProp (label n) (lo n) (hi n)) (return . Prop . D.Pure . f) p

-- | Map atoms in proposition
mapAtom :: (Eq l, Hashable l) => (l -> l) -> Prop l s -> Diagram l s (Prop l s)
mapAtom f (Prop p) =
  Diagram $
    D.fold
      ( \n -> unDiag $ do
          lo' <- mapAtom f (lo n)
          hi' <- mapAtom f (hi n)
          mkProp (f $ label n) lo' hi'
      )
      (return . Prop . D.Pure)
      p

-- | Replace an atom with a non-terminal in a 'Prop' sub-diagram
bindAtom :: (Ord l, Hashable l) => Prop l s -> (l -> Diagram l s (Prop l s)) -> Diagram l s (Prop l s)
bindAtom (Prop p) f =
  Diagram $
    D.fold
      ( \n -> unDiag $ do
          a <- f (label n)
          b <- apply (&&) a (hi n)
          apply (||) b (lo n)
      )
      (return . Prop . D.Pure)
      p

-- | Simple propositions
true, false :: Prop l s
true = Prop $ D.Pure True
false = Prop $ D.Pure False

-- | Construct a diagramatic proposition from an atom
atom, notAtom :: (Hashable l, Eq l) => l -> Diagram l s (Prop l s)
atom l = mkProp l false true
notAtom l = mkProp l true false

-- | Lift a binary operation on booleans to propositions
apply :: (Hashable l, Ord l) => (Bool -> Bool -> Bool) -> Prop l s -> Prop l s -> Diagram l s (Prop l s)
apply op p@(Prop p') q@(Prop q') =
  Diagram $
    D.fromFree
      p'
      ( \n ->
          D.fromFree
            q'
            ( \m ->
                unDiag $
                  case compare (label n) (label m) of
                    LT -> do
                      lo' <- apply op (Prop $ lo n) q
                      hi' <- apply op (Prop $ hi n) q
                      mkProp (label n) lo' hi'
                    EQ -> do
                      lo' <- apply op (Prop $ lo n) (Prop $ lo m)
                      hi' <- apply op (Prop $ hi n) (Prop $ hi m)
                      mkProp (label n) lo' hi'
                    GT -> do
                      lo' <- apply op p (Prop $ lo m)
                      hi' <- apply op p (Prop $ hi m)
                      mkProp (label m) lo' hi'
            )
            (\b -> unDiag $ mapProp (`op` b) p)
      )
      (\a -> unDiag $ mapProp (op a) q)

-- | Assign an atomic proposition a boolean value
restrict :: (Hashable l, Ord l) => l -> Bool -> Prop l s -> Diagram l s (Prop l s)
restrict l b (Prop p) = do
  memo <- Diagram $ lift get
  case M.lookup (Restr l b (Prop p)) memo of
    Just r -> return r
    Nothing -> do
      r <-
        Diagram $
          D.fold
            ( \n ->
                if label n == l
                  then if b then return (hi n) else return (lo n)
                  else unDiag $ mkProp (label n) (lo n) (hi n)
            )
            (return . Prop . D.Pure)
            p
      Diagram $ lift $ put $ M.insert (Restr l b (Prop p)) r memo
      return r

-- | Lift boolean and to diagramatic proposition
and :: (Ord l, Hashable l) => Prop l s -> Prop l s -> Diagram l s (Prop l s)
and p q = do
  memo <- Diagram $ lift get
  case M.lookup (And p q) memo of
    Just r -> return r
    Nothing -> do
      r <- apply (&&) p q
      Diagram $ lift $ put $ M.insert (And p q) r memo
      return r

-- | Lift boolean or to diagramatic proposition
or :: (Ord l, Hashable l) => Prop l s -> Prop l s -> Diagram l s (Prop l s)
or p q = do
  memo <- Diagram $ lift get
  case M.lookup (Or p q) memo of
    Just r -> return r
    Nothing -> do
      r <- apply (||) p q
      Diagram $ lift $ put $ M.insert (Or p q) r memo
      return r

-- | Lift boolean not to diagramatic proposition
not :: (Ord l, Hashable l) => Prop l s -> Diagram l s (Prop l s)
not p = do
  memo <- Diagram $ lift get
  case M.lookup (Not p) memo of
    Just r -> return r
    Nothing -> do
      r <- apply (\a b -> Prelude.not a && Prelude.not b) p p
      Diagram $ lift $ put $ M.insert (Not p) r memo
      return r

-- | Determine if a proposition is satisfiable
--
--     >  anySat = fold (\_ p q -> p || q) id
anySat :: (Hashable l, Eq l) => Prop l s -> Diagram l s Bool
anySat = fold (\_ p q -> return (p || q)) return

-- | Create a summary value of a proposition
fold :: (Hashable l, Eq l) => (l -> b -> b -> Diagram l s b) -> (Bool -> Diagram l s b) -> Prop l s -> Diagram l s b
fold f g (Prop p) = Diagram $ D.fold (\n -> unDiag $ f (label n) (lo n) (hi n)) (unDiag . g) p