{- | Private types for Bayes and Factors.

Those type are not exported

-}

module Bayes.PrivateTypes( 
 -- * Classes
   BayesianDiscreteVariable(..)
 , Set(..)
 -- * Variables
 , DV(..)
 , DVSet(..)
 -- * Instantiations
 , DVI(..)
 , setDVValue
 , (=:)
 , instantiationValue
 , instantiationVariable
 -- * Vertices 
 , Vertex(..)
 -- * Misc
 , getMinBound
 -- * Indices 
 , Index(..)
 , forAllInstantiations 
 , indicesForDomain
 , fromIndex
 , instantiationDetails
 , allInstantiationsForOneVariable
 ) where


import qualified Data.List as L 

{-
	Set
-}

-- | A Set of variables used in a factor. s is the set and a the variable
class Set s where
    -- | Empty set
    emptySet :: s a
    -- | Union of two sets
    union :: Eq a => s a -> s a -> s a
    -- | Intersection of two sets
    intersection :: Eq a => s a -> s a -> s a
    -- | Difference of two sets
    difference :: Eq a => s a -> s a -> s a
    -- | Check if the set is empty
    isEmpty :: s a -> Bool
    -- | Check if an element is member of the set
    isElem :: Eq a => a -> s a -> Bool
    -- | Add an element to the set
    addElem :: Eq a => a -> s a -> s a
    -- | Number of elements in the set
    nbElements :: s a -> Int

    -- | Check if a set is subset of another one
    subset :: Eq a => s a -> s a -> Bool

    -- | Check set equality
    equal :: Eq a => s a -> s a -> Bool
    equal sa sb = (sa `subset` sb) && (sb `subset` sa)

instance Set [] where
    emptySet = []
    union = L.union
    intersection = L.intersect
    difference a b = a L.\\ b
    isEmpty [] = True 
    isEmpty _ = False
    isElem = L.elem 
    addElem a l = if a `elem` l then l else a:l
    nbElements = length
    subset sa sb = all (`elem` sb) sa

{-

Misc

-}
-- | Vertex type used to identify a vertex in a graph
newtype Vertex = Vertex {vertexId :: Int} deriving(Eq,Ord)

instance Show Vertex where 
    show (Vertex v) = "v" ++ show v

-- | A discrete variable has a number of levels which is required to size the factors
class BayesianDiscreteVariable v where
    dimension :: v -> Int 

-- | Get the minimum bound for a type
getMinBound :: Bounded a => a -> a 
getMinBound _ = minBound



{-

Variables
	
-}

-- | A discrete variable
data DV = DV !Vertex !Int deriving(Eq,Ord)

-- | A set of discrete variables
-- The tag is used to check that an index is used with the right set of DV
newtype DVSet s = DVSet [DV] deriving(Eq)

-- | Remove the type tag when not needed
fromDVSet :: DVSet s -> [DV]
fromDVSet (DVSet l) = l

instance Show DV where
    show (DV v d) = show v ++ "(" ++ show d ++ ")"

instance BayesianDiscreteVariable DV where
    dimension (DV _ d) = d

{-

Index

-}

-- | An index with meaning only for a given DVSet
newtype Index s = Index Int deriving(Eq)

-- | Used to forget the type tag
fromIndex :: Index s -> Int 
fromIndex (Index i) = i 

-- | Generate all the indices for a set of variables
indicesForDomain :: DVSet s -> [[Index s]]
indicesForDomain (DVSet l) = mapM indicesForOneDomain l
 where 
 	indicesForOneDomain (DV _ d) = map Index [0..d-1]

allInstantiationsForOneVariable :: DV -> [DVI Int]
allInstantiationsForOneVariable v@(DV _ d) = map (setDVValue v) [0..d-1]

-- | Generate all instantiations of variables
-- The DVInt can be in any order so the tag s is not used
forAllInstantiations :: DVSet s -> [[DVI Int]]
forAllInstantiations (DVSet l) = mapM allInstantiationsForOneVariable l
 

{- 

Instantiations

-}
-- | Discrete Variable instantiation. A variable and its value
data DVI a = DVI DV !a deriving(Eq)

instance Show a => Show (DVI a) where 
   show (DVI (DV v _) i) = show v ++ "=" ++ show i

   -- | A set of variable instantiations
type DVISet a = [DVI a]

-- | Create a discrete variable instantiation for a given discrete variable
setDVValue :: DV -> a -> DVI a
setDVValue v a = DVI v a

-- | Create a variable instantiation using values from
-- an enumeration
(=:) :: (Bounded b, Enum b) => DV -> b -> DVI Int 
(=:) a b = setDVValue a (fromEnum b - fromEnum (getMinBound b))

instance BayesianDiscreteVariable (DVI a) where
    dimension (DVI v _) = dimension v

-- | Get the variables and their values with a type constraint
instantiationDetails :: [DVI Int] -> (DVSet s, [Index s])
instantiationDetails l = (DVSet $ map instantiationVariable l, map (Index . instantiationValue) l)

-- | Extract value of the instantiation
instantiationValue (DVI _ v) = v

-- | Discrete variable from the instantiation
instantiationVariable (DVI dv _) = dv