{-# language TypeFamilies #-}

module Ersatz.Relation.Data ( 
-- * The 'Relation' type
  Relation
-- * Construction
, relation, symmetric_relation
, build
, buildFrom
, identity
-- * Components
, bounds, (!), indices, assocs, elems
-- *
, table
)  where

import Prelude hiding ( and )

import Ersatz.Bit
import Ersatz.Codec
import Ersatz.Variable (exists)
import Ersatz.Problem (MonadSAT)

import Control.Monad (guard)
import qualified Data.Array as A
import Data.Array ( Array, Ix )


-- | @Relation a b@ represents a binary relation \(R \subseteq A \times B \),
-- where the domain \(A\) is a finite subset of the type @a@,
-- and the codomain \(B\) is a finite subset of the type @b@.
--
-- A relation is stored internally as @Array (a,b) Bit@,
-- so @a@ and @b@ have to be instances of 'Ix',
-- and both \(A\) and \(B\) are intervals.

newtype Relation a b = Relation (A.Array (a, b) Bit)

instance (Ix a, Ix b) => Codec (Relation a b) where
  type Decoded (Relation a b) = A.Array (a, b) Bool
  decode :: forall (f :: * -> *).
MonadPlus f =>
Solution -> Relation a b -> f (Decoded (Relation a b))
decode Solution
s (Relation Array (a, b) Bit
a) = forall a (f :: * -> *).
(Codec a, MonadPlus f) =>
Solution -> a -> f (Decoded a)
decode Solution
s Array (a, b) Bit
a
  encode :: Decoded (Relation a b) -> Relation a b
encode Decoded (Relation a b)
a = forall a b. Array (a, b) Bit -> Relation a b
Relation forall a b. (a -> b) -> a -> b
$ forall a. Codec a => Decoded a -> a
encode Decoded (Relation a b)
a


-- | @relation ((amin,bmin),(amax,mbax))@ constructs an indeterminate relation \( R \subseteq A \times B \)
-- where \(A\) is @{amin .. amax}@ and \(B\) is @{bmin .. bmax}$.
relation :: ( Ix a, Ix b, MonadSAT s m ) =>
  ((a,b),(a,b)) 
  -> m ( Relation a b )
relation :: forall a b s (m :: * -> *).
(Ix a, Ix b, MonadSAT s m) =>
((a, b), (a, b)) -> m (Relation a b)
relation ((a, b), (a, b))
bnd = do
    [((a, b), Bit)]
pairs <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence forall a b. (a -> b) -> a -> b
$ do
        (a, b)
p <- forall a. Ix a => (a, a) -> [a]
A.range ((a, b), (a, b))
bnd
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
            Bit
x <- forall a s (m :: * -> *). (Variable a, MonadSAT s m) => m a
exists
            forall (m :: * -> *) a. Monad m => a -> m a
return ( (a, b)
p, Bit
x )
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Ix b) =>
((a, b), (a, b)) -> [((a, b), Bit)] -> Relation a b
build ((a, b), (a, b))
bnd [((a, b), Bit)]
pairs

-- | Constructs an indeterminate relation \( R \subseteq B \times B \)
-- that it is symmetric, i.e., \( \forall x, y \in B: ((x,y) \in R) \rightarrow ((y,x) \in R) \).
--
-- A symmetric relation is an undirected graph, possibly with loops.
symmetric_relation ::
  (MonadSAT s m, Ix b) =>
  ((b, b), (b, b)) -- ^ Since a symmetric relation must be homogeneous, the domain must equal the codomain. 
                   -- Therefore, given bounds @((p,q),(r,s))@, it must hold that @p=q@ and @r=s@.
  -> m (Relation b b)
symmetric_relation :: forall s (m :: * -> *) b.
(MonadSAT s m, Ix b) =>
((b, b), (b, b)) -> m (Relation b b)
symmetric_relation ((b, b), (b, b))
bnd = do
    [[((b, b), Bit)]]
pairs <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence forall a b. (a -> b) -> a -> b
$ do
        (b
p,b
q) <- forall a. Ix a => (a, a) -> [a]
A.range ((b, b), (b, b))
bnd
        forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ b
p forall a. Ord a => a -> a -> Bool
<= b
q
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
            Bit
x <- forall a s (m :: * -> *). (Variable a, MonadSAT s m) => m a
exists
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$   ((b
p,b
q), Bit
x)
                   forall a. a -> [a] -> [a]
: [ ((b
q,b
p), Bit
x) | b
p forall a. Eq a => a -> a -> Bool
/= b
q ]
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b.
(Ix a, Ix b) =>
((a, b), (a, b)) -> [((a, b), Bit)] -> Relation a b
build ((b, b), (b, b))
bnd forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[((b, b), Bit)]]
pairs

-- | Constructs a relation \(R \subseteq A \times B \) from a list.
-- 
-- ==== __Example__
--
-- @
-- r = build ((0,'a'),(1,'b')) [((0,'a'), true), ((0,'b'), false), 
--                          ((1,'a'), false), ((1,'b'), true))]
-- @
build :: ( Ix a, Ix b )
      => ((a,b),(a,b))
      -> [ ((a,b), Bit ) ] -- ^ A list of tuples, where the first element represents an element
                           -- \( (x,y) \in A \times B \) and the second element is a positive 'Bit'
                           -- if \( (x,y) \in R \), or a negative 'Bit' if \( (x,y) \notin R \).
      -> Relation a b
build :: forall a b.
(Ix a, Ix b) =>
((a, b), (a, b)) -> [((a, b), Bit)] -> Relation a b
build ((a, b), (a, b))
bnd [((a, b), Bit)]
pairs = forall a b. Array (a, b) Bit -> Relation a b
Relation forall a b. (a -> b) -> a -> b
$ forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
A.array ((a, b), (a, b))
bnd [((a, b), Bit)]
pairs

-- | Constructs a relation \(R \subseteq A \times B \) from a function.
buildFrom :: (Ix a, Ix b)
          => (a -> b -> Bit) -- ^ A function with the specified signature, that assigns a 'Bit'-value 
                             -- to each element \( (x,y) \in A \times B \).
          -> ((a,b),(a,b))
          -> Relation a b
buildFrom :: forall a b.
(Ix a, Ix b) =>
(a -> b -> Bit) -> ((a, b), (a, b)) -> Relation a b
buildFrom a -> b -> Bit
p ((a, b), (a, b))
bnd = forall a b.
(Ix a, Ix b) =>
((a, b), (a, b)) -> [((a, b), Bit)] -> Relation a b
build ((a, b), (a, b))
bnd forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> [a] -> [b]
map (forall a. Ix a => (a, a) -> [a]
A.range ((a, b), (a, b))
bnd) forall a b. (a -> b) -> a -> b
$ \ (a
i,b
j) -> ((a
i, b
j), a -> b -> Bit
p a
i b
j)

-- | Constructs the identity relation \(I \subseteq A \times A, I = \{ (x,x) ~|~ x \in A \} \).
identity :: (Ix a)
         => ((a,a),(a,a)) -- ^ Since the identity relation is homogeneous, the domain must equal the codomain. 
                          -- Therefore, given bounds @((p,q),(r,s))@, it must hold that @p=q@ and @r=s@.
         -> Relation a a
identity :: forall a. Ix a => ((a, a), (a, a)) -> Relation a a
identity = forall a b.
(Ix a, Ix b) =>
(a -> b -> Bit) -> ((a, b), (a, b)) -> Relation a b
buildFrom (\ a
i a
j -> forall b. Boolean b => Bool -> b
bool forall a b. (a -> b) -> a -> b
$ a
i forall a. Eq a => a -> a -> Bool
== a
j)


-- | The bounds of the array that correspond to the matrix representation of the given relation.
--
-- ==== __Example__
--
-- >>> r = build ((0,0),(1,1)) [((0,0), false), ((0,1), true), ((1,0), true), ((1,1), false))]
-- >>> bounds r
-- ((0,0),(1,1))
bounds :: (Ix a, Ix b) => Relation a b -> ((a,b),(a,b))
bounds :: forall a b. (Ix a, Ix b) => Relation a b -> ((a, b), (a, b))
bounds ( Relation Array (a, b) Bit
r ) = forall i e. Array i e -> (i, i)
A.bounds Array (a, b) Bit
r

-- | The list of indices, where each index represents an element \( (x,y) \in A \times B \) 
-- that may be contained in the given relation \(R \subseteq A \times B \).
--
-- ==== __Example__
--
-- >>> r = build ((0,0),(1,1)) [((0,0), false), ((0,1), true), ((1,0), true), ((1,1), false))]
-- >>> indices r
-- [(0,0),(0,1),(1,0),(1,1)]
indices :: (Ix a, Ix b) => Relation a b -> [(a, b)]
indices :: forall a b. (Ix a, Ix b) => Relation a b -> [(a, b)]
indices ( Relation Array (a, b) Bit
r ) = forall i e. Ix i => Array i e -> [i]
A.indices Array (a, b) Bit
r

-- | The list of tuples for the given relation \(R \subseteq A \times B \), 
-- where the first element represents an element \( (x,y) \in A \times B \) 
-- and the second element indicates via a 'Bit' , if \( (x,y) \in R \) or not.
-- 
-- ==== __Example__
--
-- >>> r = build ((0,0),(1,1)) [((0,0), false), ((0,1), true), ((1,0), true), ((1,1), false))]
-- >>> assocs r
-- [((0,0),Var (-1)),((0,1),Var 1),((1,0),Var 1),((1,1),Var (-1))]
assocs :: (Ix a, Ix b) => Relation a b -> [((a, b), Bit)]
assocs :: forall a b. (Ix a, Ix b) => Relation a b -> [((a, b), Bit)]
assocs ( Relation Array (a, b) Bit
r ) = forall i e. Ix i => Array i e -> [(i, e)]
A.assocs Array (a, b) Bit
r

-- | The list of elements of the array
-- that correspond to the matrix representation of the given relation.
--
-- ==== __Example__
--
-- >>> r = build ((0,0),(1,1)) [((0,0), false), ((0,1), true), ((1,0), true), ((1,1), false))]
-- >>> elems r
-- [Var (-1),Var 1,Var 1,Var (-1)]
elems :: (Ix a, Ix b) => Relation a b -> [Bit]
elems :: forall a b. (Ix a, Ix b) => Relation a b -> [Bit]
elems ( Relation Array (a, b) Bit
r ) = forall i e. Array i e -> [e]
A.elems Array (a, b) Bit
r

-- | The 'Bit'-value for a given element \( (x,y) \in A \times B \) 
-- and a given relation \(R \subseteq A \times B \) that indicates
-- if \( (x,y) \in R \) or not.
-- 
-- ==== __Example__
--
-- >>> r = build ((0,0),(1,1)) [((0,0), false), ((0,1), true), ((1,0), true), ((1,1), false))]
-- >>> r ! (0,0)
-- Var (-1)
-- >>> r ! (0,1)
-- Var 1
(!) :: (Ix a, Ix b) => Relation a b -> (a, b) -> Bit
Relation Array (a, b) Bit
r ! :: forall a b. (Ix a, Ix b) => Relation a b -> (a, b) -> Bit
! (a, b)
p = Array (a, b) Bit
r forall i e. Ix i => Array i e -> i -> e
A.! (a, b)
p

-- | Print a satisfying assignment from a SAT solver, where the assignment is interpreted as a relation.
-- @putStrLn $ table \</assignment/\>@ corresponds to the matrix representation of this relation.
table :: (Enum a, Ix a, Enum b, Ix b)
      => Array (a,b) Bool -> String
table :: forall a b.
(Enum a, Ix a, Enum b, Ix b) =>
Array (a, b) Bool -> String
table Array (a, b) Bool
r = [String] -> String
unlines forall a b. (a -> b) -> a -> b
$ do
    let ((a
a,b
b),(a
c,b
d)) = forall i e. Array i e -> (i, i)
A.bounds Array (a, b) Bool
r
    a
x <- [ a
a .. a
c ]
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [String] -> String
unwords forall a b. (a -> b) -> a -> b
$ do
        b
y <- [ b
b .. b
d ]
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ if Array (a, b) Bool
r forall i e. Ix i => Array i e -> i -> e
A.! (a
x,b
y) then String
"*" else String
"."