module HyLo.Model ( Model, worlds, succs, valP, valN,
                    model, equiv, expand, setSignature, (??),
                    removeWorld, removeWorlds, countInModel,
                    namesOf, propsOf,
                    ModelsRel(..), (|/=)
)

where

import Prelude hiding ( (!!) )

import Control.Monad       ( filterM )
import Control.Applicative ( (<$>) )

import Data.Maybe ( fromMaybe )

import Data.Foldable ( Foldable, toList )

import Data.Map ( Map )
import qualified Data.Map as Map

import Data.Set ( Set )
import qualified Data.Set as Set

import Text.Read ( Read(..) )
import Text.ParserCombinators.ReadPrec ( lift )
import Text.ParserCombinators.ReadP ( string )

import HyLo.Signature ( Signature,
                        HasSignature(..),
                        nomSymbols, propSymbols, relSymbols, merge )

import HyLo.Signature.Simple ( NomSymbol, PropSymbol, RelSymbol )

import HyLo.Formula ( Formula(..), Where(..), CountOp(..) )
import qualified HyLo.Formula as F

data Model w n p r = Model{worlds :: Set w,
                           succs  :: r -> w -> Set w,
                           valP   :: p -> Set w,
                           valN   :: n -> w,
                           sig     :: Signature n p r}

instance HasSignature (Model w n p r) where
    type NomsOf  (Model w n p r) = n
    type PropsOf (Model w n p r) = p
    type RelsOf  (Model w n p r) = r
    getSignature = sig


preds :: (Ord w) => Model w n p r -> r -> w -> Set w
preds m r w = Set.fromList [ w' | w' <- Set.toList $ worlds m,
                                  w2 <- Set.toList $ succs m r w',
                                  w2 == w ]

{- equiv defines an equivalence realtion among models. it is weaker than
   equality: two equivalent models can give rise to non-equivalent ones
   when the signature is expanded -}
equiv :: (Eq w, Eq n, Eq p, Eq r) => Model w n p r -> Model w n p r -> Bool
m1 `equiv` m2 = worlds m1 == worlds m2
            &&     sig m1 == sig m2
            && and [   valN m1 i == valN m2 i    | i <- toList (nomSymbols  s)]
            && and [   valP m1 p == valP m2 p    | p <- toList (propSymbols s)]
            && and [succs m1 r w == succs m2 r w | r <- toList (relSymbols  s),
                                                   w <- toList (worlds m1)]
    where s = sig m1

instance (Show w, Show n, Show p, Show r) => Show (Model w n p r) where
    showsPrec _ m = showString "Model{\n" .
                    showString " worlds = " . shows (worlds m) . showSep .
                    showString " succs  = " . shows mkRels     . showSep .
                    showString " valP   = " . shows mkProps    . showSep .
                    showString " valN   = " . shows mkNoms     . showSep .
                    showString " sig    = " . shows (sig m)    . showChar '}'
        where mkRels  = [   (w, r, w') | r  <- toList (relSymbols $ sig m),
                                         w  <- toList (worlds m),
                                         w' <- toList (succs m r w)]
              mkProps = [(p, valP m p) | p  <- toList (propSymbols $ sig m)]
              mkNoms  = [(i, valN m i) | i  <- toList (nomSymbols  $ sig m)]
              --
              showSep = showString ",\n"


instance (Read w, Read n, Read p, Read r, Ord w, Ord n, Ord p, Ord r)
      => Read (Model w n p r) where
    readPrec = do _ <- s "Model{\n"
                  _ <- s " worlds = "; ws <- readPrec; _ <- fsep
                  _ <- s " succs  = "; rs <- readPrec; _ <- fsep
                  _ <- s " valP   = "; ps <- readPrec; _ <- fsep
                  _ <- s " valN   = "; ns <- readPrec; _ <- fsep
                  _ <- s " sig    = "; si <- readPrec; _ <- s "}"
                  return $ model ws
                                 (rs :: [(w,r,w)])
                                 (Map.fromList ps :: Map p (Set w))
                                 (Map.fromList ns :: Map n w, Set.findMin ws)
                                 si
        where s = lift . string
              fsep = s ",\n"

model :: (RelRepr rs r w, ValRepr ps p w, ValNomRepr ns n w)
      => Set w -> rs -> ps -> ns -> Signature n p r -> Model w n p r
model ws rs ps ns s = Model ws (asSuccs rs) (asVal ps) (asValN ns) s

expand :: (Ord n, Ord p, Ord r)
       => Signature n p r
       -> Model w n p r
       -> Model w n p r
expand s m = m{sig = sig m `merge` s}

setSignature :: (Ord n, Ord p, Ord r)
             => Signature n p r
             -> Model w n p r
             -> Model w n p r
setSignature s m = m{sig = s}

removeWorld :: Ord w => w -> Model w n p r -> Model w n p r
removeWorld = removeWorlds . Set.singleton

removeWorlds :: Ord w => Set w -> Model w n p r -> Model w n p r
removeWorlds ws m = m'
    where m' = m{worlds = worlds m `Set.difference` ws,
                 succs  = \r w -> if w `Set.member` ws
                                    then Set.empty
                                    else succs m r w `Set.difference` ws,
                 valP   = \p -> valP m p `Set.difference` ws,
                 valN   = \n -> let w = valN m n
                                in if w `Set.member` ws
                                     then maybe (error "valN remove empty")
                                                fst
                                                (Set.minView $ worlds m')
                                     else w,
                 sig    = sig m}


-- TODO: Make it O(log applications) instead O(applications)
(??) :: Eq n => Model w n p r -> (n, w) -> Model w n p r
m ?? (x,a) = m{valN = \n -> if n == x then a else valN m n}

infix 9 |=, |/=

class ModelsRel m f n p r | m -> n, m -> p, m -> r, f -> n, f -> p, f -> r where
    (|=) :: m -> f -> Bool

(|/=) :: ModelsRel m f n p r => m -> f -> Bool
m |/= f = not (m |= f)

instance (Ord w, Ord n, Ord p, Ord r)
      => ModelsRel (Model w n p r, w) (F.Formula n p r) n p r where
    (m,w) |= _ | not (w `Set.member` worlds m) = False
    --
    (_,_) |= Top          = True
    (_,_) |= Bot          = False
    --
    (m,w) |= (Prop p)     = w `Set.member` valP m p
    (m,w) |= (Nom  n)     = w == valN m n
    --
    (m,w) |= (Neg f)      = (m,w) |/= f
    --
    (m,w) |= (f1  :&: f2) = (m,w) |= f1  &&  (m,w) |= f2
    (m,w) |= (f1  :|: f2) = (m,w) |= f1  ||  (m,w) |= f2
    --
    (m,w) |= (f1 :-->:f2) = (m,w) |= (Neg f1 :|: f2)
    (m,w) |= (f1:<-->:f2) = (m,w) |= f1 == (m,w) |= f2
    --
    (m,w) |= (Diam r f)   = or  [(m,w') |= f | w' <- toList (succs m r w)]
    (m,w) |= (Box  r f)   = and [(m,w') |= f | w' <- toList (succs m r w)]
    --
    (m,w) |= (IDiam r f)  = or  [(m,w') |= f | w' <- toList (preds m r w)]
    (m,w) |= (IBox  r f)  = and [(m,w') |= f | w' <- toList (preds m r w)]
    --
    (m,_) |= (At  n f)    = (m,valN m n) |= f
    --
    (m,_) |= (E f)        = or  [(m,w') |= f | w' <- toList (worlds m)]
    (m,_) |= (A f)        = and [(m,w') |= f | w' <- toList (worlds m)]
    --
    (m,w) |= (D f)        = or  [(m,w') |= f | w' <- toList (Set.delete w $ worlds m)]
    (m,w) |= (B f)        = and [(m,w') |= f | w' <- toList (Set.delete w $ worlds m)]
    --
    (m,w) |= (Down x f)   = (m ?? (x,w),w) |= f
    --
    (m,w) |= (Count c l i f) = [(m,w') |= f | w' <- toList (whatToCount l)] `cmpLenWithC` i
                                 where whatToCount Global    = worlds m
                                       whatToCount (Local r) = succs m r w
                                       cmpLenWithC = F.cmpListLen c


instance (Ord w, Ord n, Ord p, Ord r)
      => ModelsRel (Model w n p r) (F.Formula n p r) n p r where
    m |= f = and [(m,w) |= f | w <- Set.toList $ worlds m]

class RelRepr a r w | a -> r, a -> w where
    asSuccs :: a -> (r -> w -> Set w)

instance RelRepr (r -> w -> Set w) r w where
    asSuccs = id

instance Ord w => RelRepr (r -> w -> [w]) r w where
    asSuccs s = \r w -> Set.fromList (s r w)

instance (Ord r, Ord w) => RelRepr (Map r (Map w (Set w))) r w where
    asSuccs rs = \r w -> fromMaybe Set.empty $
                           do ws <- Map.lookup r rs
                              Map.lookup w ws

instance (Ord r, Ord w, Foldable t) => RelRepr (Map r (t (w,w))) r w where
    asSuccs = asSuccs . fmap mapRep

instance (Ord r, Ord w) => RelRepr [(w,r,w)] r w where
    asSuccs = asSuccs . fmap mapRep . mapRep . map f
        where f (w,r,w') = (r,(w,w'))

instance (Ord r, Ord w) => RelRepr (Set (w,r,w)) r w where
    asSuccs = asSuccs . toList



class ValRepr a p w | a -> p, a -> w where
    asVal :: a -> (p -> Set w)

instance ValRepr (p -> Set w) p w where
    asVal = id

instance Ord w => ValRepr (p -> [w]) p w where
    asVal v = Set.fromList . v

instance (Ord w, Ord p) => ValRepr (Map p (Set w)) p w where
    asVal v = \p -> fromMaybe Set.empty $ Map.lookup p v

instance (Ord w, Ord p) => ValRepr (Map p [w]) p w where
    asVal = asVal . fmap Set.fromList

instance (Ord w, Ord p, Foldable t) => ValRepr (t (p,w)) p w where
    asVal = asVal . mapRep


class ValNomRepr a n w | a -> n, a -> w where
    asValN :: a -> (n -> w)

instance ValNomRepr (n -> w) n w where
    asValN = id

instance Ord n => ValNomRepr (Map n w, n -> w) n w where
    asValN (m, def) = \n -> fromMaybe (def n) (Map.lookup n m)

instance Ord n => ValNomRepr (Map n w, w) n w where
    asValN (m, def) = asValN (m, const def :: n -> w)


mapRep :: (Ord k, Ord v, Foldable t) => t (k,v) -> Map k (Set v)
mapRep = fmap setify . Map.fromListWith (.) . map f . toList
    where f (a,b)  = (a, (b:))
          setify x = Set.fromList (x [])

namesOf :: Eq w => w -> Model w n p r -> [n]
namesOf w m = [n | n <- toList . nomSymbols . getSignature $ m, valN m n == w]

propsOf :: Ord w => w -> Model w n p r -> [p]
propsOf w m = [p | p <- toList . propSymbols . getSignature $ m,
                   w `Set.member` valP m p]

countInModel :: (Ord w, Ord n, Ord p, Ord r) => F.Formula n p r -> Model w n p r -> Int
countInModel f m = length [ w | w <- toList (worlds m), (m,w) |= f ]