{-# LANGUAGE DeriveDataTypeable  #-}
{-# LANGUAGE DeriveFoldable      #-}
{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE DeriveTraversable   #-}
{-# LANGUAGE Safe                #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Algebra.Lattice.Free (
    Free (..),
    liftFree,
    lowerFree,
    substFree,
    retractFree,
    toExpr,
    ) where

import Algebra.Lattice
import Algebra.PartialOrd

import Control.Applicative          (liftA2)
import Control.Monad                (ap)
import Data.Data                    (Data, Typeable)
import GHC.Generics                 (Generic, Generic1)
import Math.NumberTheory.Logarithms (intLog2)

import qualified Algebra.Heyting.Free.Expr as E
import qualified Test.QuickCheck           as QC

-- $setup
-- >>> import Algebra.Lattice

-------------------------------------------------------------------------------
-- Free
-------------------------------------------------------------------------------

-- | Free distributive lattice.
--
-- `Eq` and `PartialOrd` instances aren't structural.
--
-- >>> (Var 'x' /\ Var 'y') == (Var 'y' /\ Var 'x' /\ Var 'x')
-- True
--
-- >>> Var 'x' == Var 'y'
-- False
--
-- This is /distributive/ lattice.
--
-- >>> import Algebra.Lattice.M3 -- non distributive lattice
-- >>> let x = M3a; y = M3b; z = M3c
-- >>> let lhs = Var x \/ (Var y /\ Var z)
-- >>> let rhs = (Var x \/ Var y) /\ (Var x \/ Var z)
--
-- 'Free' is distributive so
--
-- >>> lhs == rhs
-- True
--
-- but when retracted, values are inequal
--
-- >>> retractFree lhs == retractFree rhs
-- False
--
-- >>> (retractFree lhs, retractFree rhs)
-- (M3a,M3i)
--
data Free a
    = Var a
    | Free a :/\: Free a
    | Free a :\/: Free a
  deriving (Show, Functor, Foldable, Traversable, Generic, Generic1, Data, Typeable)

infixr 6 :/\:
infixr 5 :\/:

liftFree :: a -> Free a
liftFree = Var

retractFree :: Lattice a => Free a -> a
retractFree = lowerFree id

substFree :: Free a -> (a -> Free b) -> Free b
substFree z k = go z where
    go (Var x)    = k x
    go (x :/\: y) = go x /\ go y
    go (x :\/: y) = go x \/ go y

lowerFree :: Lattice b => (a -> b) -> Free a -> b
lowerFree f = go where
    go (Var x)    = f x
    go (x :/\: y) = go x /\ go y
    go (x :\/: y) = go x \/ go y

toExpr :: Free a -> E.Expr a
toExpr (Var a)    = E.Var a
toExpr (x :/\: y) = toExpr x E.:/\: toExpr y
toExpr (x :\/: y) = toExpr x E.:\/: toExpr y

-------------------------------------------------------------------------------
-- Monad
-------------------------------------------------------------------------------

instance Applicative Free where
    pure = liftFree
    (<*>) = ap

instance Monad Free where
    return = pure
    (>>=)  = substFree

-------------------------------------------------------------------------------
-- Instances
-------------------------------------------------------------------------------

instance Lattice (Free a) where
    x /\ y = x :/\: y
    x \/ y = x :\/: y

instance Ord a => Eq (Free a) where
    (==) = partialOrdEq

instance Ord a => PartialOrd (Free a) where
    leq x y = E.proofSearch (toExpr x E.:=>: toExpr y)

-------------------------------------------------------------------------------
-- Other instances
-------------------------------------------------------------------------------

instance QC.Arbitrary a => QC.Arbitrary (Free a) where
    arbitrary = QC.sized arb where
        arb n | n <= 0    = prim
              | otherwise = QC.oneof (prim : compound)
          where
            arb' = arb (intLog2 (max 1 n))

            compound =
                [ liftA2 (:/\:) arb' arb'
                , liftA2 (:\/:) arb' arb'
                ]

        prim = Var <$> QC.arbitrary

    shrink (Var c)    = map Var (QC.shrink c)
    shrink (x :/\: y) = x : y : map (uncurry (:/\:)) (QC.shrink (x, y))
    shrink (x :\/: y) = x : y : map (uncurry (:\/:)) (QC.shrink (x, y))
