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

where

import Prelude hiding ( (!!) )

import Test.QuickCheck     ( Arbitrary(..), Gen, Property,
                             elements, forAll, (==>) )
import HyLo.Test           ( UnitTest, runTest )
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.Map.Arbitrary()

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 ]

-- -----------------------------
-- QuickCheck stuff
-- -----------------------------

instance (Ord w, Ord n, Ord p, Ord r,
          Arbitrary w, Arbitrary n, Arbitrary p, Arbitrary r)
      => Arbitrary (Model w n p r) where
    arbitrary = do (w,ws)       <- arbitrary -- carrier set must be non-empty
                   (rs,ps,ns,s) <- arbitrary
                   return $ Model (w `Set.insert` ws) rs ps ns s
    --
    coarbitrary (Model w r p n s) = coarbitrary w .
                                    coarbitrary r .
                                    coarbitrary p .
                                    coarbitrary n .
                                    coarbitrary s

infix 1 `iff`
iff :: Bool -> Bool -> Bool
iff = (==)

modelFor :: (HasSignature x,Ord (NomsOf x), Ord (RelsOf x), Ord (PropsOf x))
         => x
         -> Gen (Model Int (NomsOf x) (PropsOf x) (RelsOf x))
modelFor x =
    do (w,ws) <- arbitrary -- carrier set must be non empty
       let dom     =  w:ws
           s       = getSignature x
       rs <- go (relSymbols  s) (\r -> sublist [(v,r,v') | v <- dom, v' <- dom])
       ps <- go (propSymbols s) (\p -> map (\a -> (p,a)) <$> sublist dom)
       ns <- go (nomSymbols  s) (\i -> (,) i <$> elements dom)
       --
       return $ model (Set.fromList dom)
                      (concat rs)
                      (concat ps)
                      (Map.fromList ns, w)
                      s
  where go ::  Set a -> (a -> Gen b) -> Gen [b]
        go c p = mapM p (Set.toList c)
        --
        sublist :: [a] -> Gen [a]
        sublist = filterM (const arbitrary)

worldOf :: Model w n p r -> Gen w
worldOf m = elements (Set.toList $ worlds m)

type TestModel     = Model Int NomSymbol PropSymbol RelSymbol
type TestFormula   = F.Formula NomSymbol PropSymbol RelSymbol
type TestSignature = Signature NomSymbol PropSymbol RelSymbol

prop_equiv_ref :: TestModel -> Bool
prop_equiv_ref m = m `equiv` m

prop_read_equiv :: TestModel -> Bool
prop_read_equiv m = (read . show $ m) `equiv` m

prop_lem ::  TestFormula -> Property
prop_lem f = forAll (modelFor f) $ \m ->
               forAll (worldOf m) $ \w ->
                 (m,w) |= f || (m,w) |= Neg f

prop_neg :: TestModel ->  TestFormula -> Property
prop_neg m f = forAll (worldOf m) $ \w ->
                 (m,w) |/= f || (m,w) |/= Neg f

prop_nnfPreservesSat :: TestFormula -> Property
prop_nnfPreservesSat f = forAll (modelFor f) $ \m ->
                           forAll (worldOf m) $ \w ->
                             (m,w) |= f == (m,w) |= F.nnf f

prop_expandDoesNotChangePrev :: TestFormula -> TestSignature -> Property
prop_expandDoesNotChangePrev f s = forAll (modelFor f) $ \m ->
                                     forAll (worldOf m) $ \w ->
                                       (m,w) |= f == (expand s m,w) |= f

prop_removeWorlds :: TestModel -> Property
prop_removeWorlds m =
  Set.size (worlds m) > 1 ==> forAll (worldOf m) $ \w ->
    let m' = removeWorld w m in
     and [s == sig m',
          not $ w `Set.member` worlds m',
          all (not . (w `Set.member`) . valP m') (toList . propSymbols $ s),
          all (not . (w ==) . valN m') (toList . nomSymbols $ s),
          all (not . (w `Set.member`) . uncurry (succs m'))
              [(r,w') | r <- toList (relSymbols s), w' <- toList (worlds m')]
         ]
    where s = sig m

prop_axK :: RelSymbol -> TestFormula -> TestFormula -> Property
prop_axK r p q = mp_valid $ (Box r $ p :-->: q) :-->: (Box r p) :-->: (Box r q)

prop_axKAt :: NomSymbol -> TestFormula -> TestFormula -> Property
prop_axKAt i p q = mp_valid $ (At i $ p :-->: q) :-->: (At i p) :-->: (At i q)

prop_axKDown :: NomSymbol -> TestFormula -> TestFormula -> Property
prop_axKDown x p q = mp_valid $ (Down x $ p :-->: q) :-->:
                                (Down x p) :-->: (Down x q)

prop_axKA:: TestFormula -> TestFormula -> Property
prop_axKA p q = mp_valid $ (A $ p :-->: q) :-->: (A p) :-->: (A q)

prop_axKB:: TestFormula -> TestFormula -> Property
prop_axKB p q = mp_valid $ (B $ p :-->: q) :-->: (B p) :-->: (B q)

prop_axDualBoxDiam :: RelSymbol -> TestFormula -> Property
prop_axDualBoxDiam r p = mp_valid $ (Box r p) :<-->: (Neg $ Diam r (Neg p))


prop_axSelfDualAt :: NomSymbol -> TestFormula -> Property
prop_axSelfDualAt i p = mp_valid $ (At i p) :<-->: (Neg $ At i (Neg p))

prop_axSelfDualDown :: NomSymbol -> TestFormula -> Property
prop_axSelfDualDown x p = mp_valid $ (Down x p) :<-->: (Neg $ Down x (Neg p))

prop_axDualAE :: TestFormula -> Property
prop_axDualAE p = mp_valid $ (A p) :<-->: (Neg $ E (Neg p))

prop_axDualBD :: TestFormula -> Property
prop_axDualBD p = mp_valid $ (B p) :<-->: (Neg $ D (Neg p))

prop_axPast :: RelSymbol -> TestFormula -> Property
prop_axPast r p = mp_valid $ p :-->: (Box r $ IDiam r p)

prop_count :: TestFormula -> Int -> Property
prop_count f i =
     forAll (modelFor f) $ \m ->
         forAll (worldOf m) $ \w ->
            ( ((f `countInModel` m) == i) `iff` ((m,w) |= Count (:=:) Global i f ) )


mp_valid :: TestFormula -> Property
mp_valid f = forAll (modelFor f) $ \m ->
               forAll (worldOf m) $ \w ->
                  (m,w) |= f


unit_tests :: UnitTest
unit_tests = [
    ("equiv is reflexive",     runTest prop_equiv_ref),
    ("read/show modulo equiv", runTest prop_read_equiv),
    ("law of excluded middle", runTest prop_lem),
    ("closed under negation",  runTest prop_neg),
    ("nnf preserves sat",      runTest prop_nnfPreservesSat),
    ("expand preserves prev",  runTest prop_expandDoesNotChangePrev),
    ("removeWorlds works",     runTest prop_removeWorlds),
    ("axiom K",                runTest prop_axK),
    ("axiom K@",               runTest prop_axKAt),
    ("axiom KDown",            runTest prop_axKDown),
    ("axiom KA",               runTest prop_axKA),
    ("axiom KB",               runTest prop_axKB),
    ("axiom dual [] - <>",     runTest prop_axDualBoxDiam),
    ("axiom self-dual-@",      runTest prop_axSelfDualAt),
    ("axiom self-dual-down",   runTest prop_axSelfDualDown),
    ("axiom dual A - E",       runTest prop_axDualAE),
    ("axiom dual B - D",       runTest prop_axDualBD),
    ("axiom past",             runTest prop_axPast),
    ("count worlds",           runTest prop_count)
  ]