module Exercise09 where

import Types

import Data.List (sort, sortOn, nub, nubBy, delete, (\\))
import Data.Ord (Down(Down))
import Data.Maybe

import Data.Set (Set)
import qualified Data.Set as Set
-- ~ import Data.Map (Map)
-- ~ import qualified Data.Map as Map
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

-- ~ import Debug.Trace (traceShowId)
traceShowId = id

-- things from the tutorial

instance Show Polarity where
  show Pos = ""
  show Neg = "~"

instance Show ELiteral where
  show (Literal p a) = show p ++ a

instance Show MPolarity where
  show MPos = ""
  show MNeg = "~"

instance Show MLiteral where
  show (MLiteral p a) = show p ++ a

lName :: Literal -> Name
lName (MLiteral _ n) = n

lNameE :: ELiteral -> EName
lNameE (Literal _ n) = n

lPos :: Name -> Literal
lPos = MLiteral MPos

lNeg :: Name -> Literal
lNeg = MLiteral MNeg

lPosE :: EName -> ELiteral
lPosE = Literal Pos

lNegE :: EName -> ELiteral
lNegE = Literal Neg

lIsPos :: Literal -> Bool
lIsPos (MLiteral p _) = p == MPos

lIsPosE :: ELiteral -> Bool
lIsPosE (Literal p _) = p == Pos

lIsNeg :: Literal -> Bool
lIsNeg = not . lIsPos

lNegateE :: ELiteral -> ELiteral
lNegateE (Literal Pos n) = Literal Neg n
lNegateE (Literal Neg n) = Literal Pos n

lNegate :: MLiteral -> MLiteral
lNegate (MLiteral MPos n) = MLiteral MNeg n
lNegate (MLiteral MNeg n) = MLiteral MPos n

complements :: Literal -> Literal -> Bool
complements (MLiteral p n) (MLiteral p' n') = p /= p' && n == n'

complementsE :: ELiteral -> ELiteral -> Bool
complementsE (Literal p n) (Literal p' n') = p /= p' && n == n'

evalLiteral :: Valuation -> Literal -> Bool
evalLiteral val l@(MLiteral _ n)
  | lIsPos l  = n `elem` val
  | otherwise = n `notElem` val

evalLiteralE :: EValuation -> ELiteral -> Bool
evalLiteralE val l@(Literal _ n)
  | lIsPosE l  = n `elem` val
  | otherwise = n `notElem` val

evalClause :: Valuation -> Clause -> Bool
evalClause = any . evalLiteral

evalClauseE :: EValuation -> EClause -> Bool
evalClauseE = any . evalLiteralE

eval :: Valuation -> ConjForm -> Bool
eval = all . evalClause

evalE :: EValuation -> EConjForm -> Bool
evalE = all . evalClauseE

clauseIsTauto :: Clause -> Bool
clauseIsTauto c = or [lNegate l `Set.member` c | l <- Set.toList c]

-- homework starts here

-- resolves a new clause from two clauses
-- assuming that
-- a positive literal with given name is in the first and
-- a negative literal with same name is in the second
-- feel free to change behaviour and type
resolve :: Name -> Clause -> Clause -> Clause
resolve n cp cn = Set.union (Set.delete (MLiteral MPos n) cp) (Set.delete (MLiteral MNeg n) cn)

-- gets two clauses with keys
-- and returns all resolved clauses together with the Resolve thing
-- feel free to change behaviour and type
resolvants :: KeyClause -> KeyClause -> [(Resolve, Clause)]
resolvants kc1@(k1, c1) kc2@(k2, c2) = [
    resolvant n (if p == MPos then (kc1, kc2) else (kc2, kc1))
    | l@(MLiteral p n) <- Set.toList c1, lNegate l `Set.member` c2] -- TODO: use sorted merge join
  where
    -- creates one resolvant from positive and negative key and clause
    resolvant n ((kp, cp), (kn, cn)) = (MResolve n kp kn, resolve n cp cn)

-- proof checking

-- don't change the type nor expected behaviour
proofCheck :: EConjForm -> Proof -> Bool
proofCheck f p = proofCheckM f (toEProof p)

proofCheckM :: EConjForm -> EProof -> Bool
proofCheckM formula (Refutation rs) = refuteCheck (zip [0..] formula) rs
  where
    refuteCheck :: [EKeyClause] -> [EResolve] -> Bool
    refuteCheck kcs [] = True
    refuteCheck kcs ((Resolve n kp kn):rs)
      | kp >= len || kn >= len = False
      | Literal Pos n `notElem` cp = False
      | Literal Neg n `notElem` cn = False
      | otherwise = refuteCheck ((len, new_clause):kcs) rs
      where
        len = length kcs
        Just cp = lookup kp kcs
        Just cn = lookup kn kcs
        new_clause = toEClause $ resolve n (toClause cp) (toClause cn)
proofCheckM formula (Model vals) = evalE vals formula

-- select clause with fewest literals
-- feel free to change behaviour and type
selClause :: SelClauseStrategy
selClause kcs_set
  | null kcs_set = Nothing
  | otherwise = Just $ head shortests
  where
    shortests = [kc | kc@(_, c) <- kcs, length c == min_len]
    min_len = minimum [length c | (_, c) <- kcs]
    kcs = Map.foldrWithKey (\c k kcs -> (k, c) : kcs) [] kcs_set

-- feel free to change behaviour and type
resolutionParam :: SelClauseStrategy -> EConjForm -> (Proof, KeyConjForm)
resolutionParam sellKlaus cs = resolutionParamDo u_kcs Map.empty given_kcs [] (length cs)
  where
    u_kcs = Map.filterWithKey (\c _ -> not $ clauseIsTauto c) given_kcs
    given_kcs = Map.fromList $ zip (map toClause cs) [0..]
    -- unprocessed -> processed -> all_kcs -> resolve_steps_reversed -> num_clauses -> (proof, kcs)
    -- iterate over unprocessed
    -- num_clauses is used to give new numbers
    resolutionParamDo :: KeyConjForm -> KeyConjForm -> KeyConjForm -> [MResolve] -> Int -> (Proof, KeyConjForm)
    resolutionParamDo u_kcs p_kcs all_kcs rrs cl_count
      | Set.empty `Map.member` u_kcs = (MRefutation $ reverse rrs, all_kcs) -- found empty clause
      | Map.null u_kcs = (MModel $ extractModel $ Map.keysSet p_kcs, all_kcs) -- nothing to process => must be true
      | otherwise = resolutionParamDo u_kcs_2 p_kcs_1 all_kcs_1 rrs_1 cl_count_1
      where
        Just u_kc@(u_k, u_c) = sellKlaus u_kcs
        u_kcs_1 = Map.delete u_c u_kcs
        (u_kcs_2, all_kcs_1, rrs_1, cl_count_1) = allResolvants u_kc p_kcs_l u_kcs_1 all_kcs rrs cl_count
        p_kcs_l = Map.foldrWithKey (\c k kcs -> (k, c) : kcs) [] p_kcs
        p_kcs_1 = Map.insert u_c u_k p_kcs

        -- u_kc -> p_kcs_l -> u_kcs_1 -> all_kcs -> rrs -> cl_count -> (new_u_kcs, new_all_kcs, new_rrs, new_cl_count)
        -- new kcs are numbered with [cl_count..=(new_cl_count-1)]
        -- new clauses and co. are added to the given containers, the modified containers are returned (they are all accumulators)
        -- we are iterating over p_kcs
        allResolvants :: KeyClause -> [KeyClause] -> KeyConjForm -> KeyConjForm -> [MResolve] -> Int -> (KeyConjForm, KeyConjForm, [MResolve], Int)
        allResolvants u_kc [] u_kcs all_kcs rrs cl_count = (u_kcs, all_kcs, rrs, cl_count)
        allResolvants u_kc (p_kc:p_kcs) u_kcs all_kcs rrs cl_count = allResolvants u_kc p_kcs new_u_kcs new_all_kcs new_rrs new_cl_count
          where
            the_resolvants = resolvants u_kc p_kc
            (new_u_kcs, new_all_kcs, new_rrs, new_cl_count) = addWithKey the_resolvants u_kcs all_kcs rrs cl_count

            -- adds the resolvats with keys in [cl_count..=(new_cl_count-1)]
            addWithKey :: [(Resolve, Clause)] -> KeyConjForm -> KeyConjForm -> [MResolve] -> Int -> (KeyConjForm, KeyConjForm, [MResolve], Int)
            addWithKey [] u_kcs all_kcs rrs cl_count = (u_kcs, all_kcs, rrs, cl_count)
            addWithKey ((r,c):rcs) u_kcs all_kcs rrs cl_count
              | clauseIsTauto c = addWithKey rcs u_kcs all_kcs rrs cl_count -- donot add tautos
              | c `Map.member` all_kcs = addWithKey rcs u_kcs all_kcs rrs cl_count -- donot add duplicates
              | is_subsumed = addWithKey rcs u_kcs all_kcs rrs cl_count -- new clause is subsumed by existing
              | otherwise = addWithKey rcs new_u_kcs new_all_kcs new_rrs new_cl_count
              where
                new_u_kcs = Map.insert c cl_count new_u_kcs_1
                new_all_kcs = Map.insert c cl_count all_kcs
                new_rrs = r:rrs
                new_cl_count = cl_count + 1
                -- ~ is_subsumed = or [c' `Set.isSubsetOf` c | c' <- Map.keys u_kcs]
                is_subsumed = False -- probably only happens rarely and is expensive
                -- subsume existing
                new_u_kcs_1 = Map.filterWithKey (\c' _ -> not (c `Set.isSubsetOf` c')) u_kcs

{-WETT-}

-- don't change the type; instantiate your best resolution prover here
-- Please leave some comments for the MC Jr to understand your ideas and code :')
resolution :: EConjForm -> (Proof, KeyConjForm)
resolution = resolutionParam selClause

{-TTEW-}

-- extract a model as described on https://lara.epfl.ch/w/_media/sav08/gbtalk.pdf

instance Ord Polarity where
  Neg <= Pos = False
  _ <= _ = True

instance Ord ELiteral where
  (Literal p a) <= (Literal p' a') = a < a' || a == a' && p <= p'

extractModel :: ConjForm -> Valuation
extractModel = toValuation . extractModelE . toEConjForm

extractModelE :: EConjForm -> EValuation
extractModelE = run [] . sort . map (sortOn Down)
  where
    run :: EValuation -> EConjForm -> EValuation
    run val [] = val
    run val (c:cs)
      | evalClauseE val c = run val cs
      | otherwise = run (lNameE (head c):val) cs