{-# LANGUAGE BangPatterns #-}
module Exercise_9 where
{-WETT-}

import Data.List
import Data.Maybe
import Data.Either
import Data.Functor ((<$))
import Control.Monad (foldM, guard)
import Data.IntMap.Strict (IntMap, (!?))
import qualified Data.IntMap.Strict as IM
import Data.Sequence (Seq, (><), (<|), (|>), ViewL((:<)))
import qualified Data.Sequence as SQ
import Data.IntSet (IntSet)
import qualified Data.IntSet as IS
import qualified Data.Foldable as F
import Data.Tuple
import Data.Map (Map)
import qualified Data.Map as M

type Name = String

type Valuation = [(Name,Bool)]

data Atom = T | Var Name
  deriving (Eq, Show)

data Literal = Pos Atom | Neg Atom
  deriving (Eq, Show)

data Form = L Literal | Form :&: Form | Form :|: Form
  deriving (Eq, Show)

type Clause = [Literal]

type ConjForm = [Clause]

{-T9.3.2-}
top :: Literal
top = Pos T

bottom :: Literal
bottom = Neg T

{-T9.3.3-}
clauseToForm :: Clause -> Form
clauseToForm [] = L bottom
clauseToForm ls = foldr ((:|:) . L) (L $ last ls) (init ls)

conjToForm :: ConjForm -> Form
conjToForm [] = L top
conjToForm ds = foldr ((:&:) . clauseToForm) (clauseToForm $ last ds) (init ds)

{-T9.3.4-}
substLiteral :: Valuation -> Literal -> Literal
substLiteral v l@(Pos (Var n)) = case lookup n v of
  Just b -> if b then top else bottom
  Nothing -> l
substLiteral v l@(Neg (Var n)) = case lookup n v of
  Just b -> if b then bottom else top
  Nothing -> l
substLiteral _ l = l

substClause :: Valuation -> Clause -> Clause
substClause = map . substLiteral

substConj :: Valuation -> ConjForm -> ConjForm
substConj = map . substClause

-- ignore this, it is terrible
{-H9.2.1-}
simpConj :: ConjForm -> ConjForm
simpConj c = reverse $ map reverse $ simpConj' c []
        where simpLit l = simpLit' l []
                where  simpLit' [] acc = acc
                       simpLit' (l:ls) acc
                           | l == top = [top]
                           | l == bottom = simpLit' ls acc
                           | otherwise = simpLit' ls $ l:acc
              simpConj' [] acc = acc
              simpConj' (c:cs) acc 
                  | null k = [[]]
                  | head k == top = simpConj' cs acc
                  | otherwise = simpConj' cs $ k:acc
                      where k = simpLit c

{-H9.2.2-}
cnf :: Form -> ConjForm
cnf (L l)  = [[l]]
cnf (a:&:b) = cnf a ++ cnf b
cnf (a:|:b) = [a'++ b' | a' <- cnf a, b' <- cnf b]

-- Just a protoype for testing the SAT-Solver
-- Absolutely no heutistics
{-H9.2.3-}
selectV :: ConjForm -> Maybe (Name, Bool)
selectV [] = Nothing
selectV (x:xs) = if null x' then selectV xs else return (getV $ head x', True)
        where x' = filter isV x

isV :: Literal -> Bool
isV l   = not $ l == top || l == bottom

getV :: Literal -> String
getV (Pos (Var k)) = k
getV (Neg (Var k)) = k
getV _             = undefined

{-H9.2.5-}
-- Most of the coding was just busywork, I must confess I only feel great about
-- simpEval because discovering guard made me happy. Apart from BCP the algo
-- is still very very naive (mostly because I want to enjoy time with family and friends).
-- Oh yeah - Merry Late-Christmas/Happy new year/Happy whatever day it is to 
-- anybody bothering to decipher the code below - everybody else is already having a 
-- better day by definition :D

-- The first int represents the length of the set. This is crucial to the functionality of
-- the algorithms below, so for a given MapClause(l,(t,f)) 
-- l == size t + size f
-- must always hold.
-- The IntSets represent the Variables that are True/False respectively
type MapClause = (Int, (IntSet, IntSet))

-- I think this is self-explanatory
type CNF =  Seq(MapClause)
type Assignment = (Int, Bool) 
type Eval = Seq(Assignment)
type VarMapping = IntMap Name

-- First, construct the Bijection, then run the DP-Algo
-- and finally remap the Eval to a Valuation
satConj :: ConjForm -> Maybe Valuation
satConj cj = case constructBijection cj of
               Nothing -> Nothing
               Just (cnf, vm) ->  case dp cnf SQ.empty of
                           Nothing -> Nothing
                           Just ev -> return $ reconstruct ev vm

-- Davis-Putnam Algorithm, really basic implementation...
dp :: CNF -> Eval -> Maybe Eval
dp cnf ev = case doBCP cnf SQ.empty of
              Nothing -> Nothing
              Just (cnf', ev') -> if SQ.null cnf' 
                                     then return $ ev >< ev'
                                     else let   
                                              (v,b) = unsafeSelect cnf'
                                              cnfA = simpEval cnf' (v,b)
                                              cnfB = simpEval cnf' (v, not b)
                                              tryT = case cnfA of
                                                       Nothing -> Nothing
                                                       Just cnfA -> dp cnfA $ SQ.singleton (v,b)
                                              tryF = case cnfB of
                                                       Nothing -> Nothing
                                                       Just cnfB -> dp cnfB $ SQ.singleton (v,not b)
                                          in
                                     if isJust tryT then return $ ev >< ev' >< fromJust tryT
                                     else if isJust tryF then return $ ev >< ev' >< fromJust tryF
                                     else Nothing

-- BooleanConstraintPropagation, pretty straight-forward...
-- Returns Nothing in case of a conflict and just the updated cnf
-- along with the corresponding assignments after fully
-- propagating a boolean constraint.
-- The way the cnf is processed ensures unit-clauses are always at the beginning of
-- the cnf, if any exist ==> if the first element of cnf is not a unit clause
-- there sumply aren't any.
doBCP :: CNF -> Eval -> Maybe (CNF, Eval)
doBCP cnf eval = if SQ.null cnf then return (cnf, eval) -- handle empty clause
                                else let (l, (t,f)) = unsafeHead $ SQ.viewl cnf in
                                        case l of
                                          -- unit clause
                                          1 -> let ass = if IS.null t 
                                                            then (IS.findMin f, False)
                                                            else (IS.findMin t, True)
                                                   cnf' = simpEval cnf ass in
                                                  case cnf' of
                                                    Nothing -> Nothing
                                                    Just cnf' -> doBCP cnf' $ eval |> ass
                                          -- no unit clause :(
                                          _ -> return (cnf, eval)
        where unsafeHead (a:<_) = a -- pattern-match against ViewL seems to be faster than index 0

-- Use foldM to implement a filter-like function that will terminate
-- if it finds an empty clause to avoid computing irrelevant values 
-- in case of a conflict. Unit clauses are prepended in the accumulator,
-- so their existence can be tested in O(1). Thank god for doubly-linked
-- lists.
--
-- We have three possible results from foldM:
-- > Nothing        ==> CNF contains CONFLICT; return Nothing
-- > Just empty     ==> CNF is SAT;   return Just CNF'
-- > Just MapClause ==> CNF is UNSAT; return Just CNF'
-- Calling functions are responsible for checking for empty/SAT
--
-- Note: I am pretty sure (>>=) is strict for maybe, so there
-- is no need to make this strict by hand..?
-- Also, this function looks huge, but it's realy just an
-- really ugly bunch of conditions D:
simpEval :: CNF -> Assignment -> Maybe CNF
simpEval cnf (var,b) = foldM rplEv SQ.empty cnf
              where rplEv accCNF (l,(t,f)) = case cl of
                                  Left cl@(ct,cf) -> return $ if IS.null ct && IS.null cf
                                                           then accCNF -- throw out True-Clause
                                                           else if l == 1 -- is unit clause?
                                                           then (l, cl) <| accCNF -- prepend
                                                           else accCNF |> (l, cl) -- append
                                  Right cl        -> if l == 2 -- is unit clause?
                                                        then return $ (pred l, cl) <| accCNF
                                                        else (accCNF |> (pred l, cl)) 
                                                              <$ guard(l > 1) -- conflict-guard
                      where cl = case b of
                                   True ->  if IS.member var t 
                                               then Left (IS.empty, IS.empty) -- clause =^= True
                                               else if IS.member var f
                                               then Right (t, IS.delete var f) -- remove False
                                            else Left (t,f) -- unchanged
                                   False -> if IS.member var f
                                               then Left (IS.empty, IS.empty) -- clause =^= True
                                               else if IS.member var t
                                               then Right (IS.delete var t, f) -- remove False
                                            else Left (t, f) -- unchanged
                
-- CAREFUL: This WILL throw an error when being passed the empty sequence
-- AND if there is an element where both IntSets are empty
-- No heuristics for now, that could change, but it probably won't :D
unsafeSelect :: CNF -> Assignment
unsafeSelect cnf = if IS.null t then (IS.findMin f, False) else (IS.findMin t, True)
        where unsafeHead (a:<_) = a -- pattern-match against ViewL seems to be faster than index 0
              (_,(t,f))            = unsafeHead $ SQ.viewl cnf

-- Use IntSets instead of Sets really just because I have to parse the ConjForm anyways.
-- (but also because I lowkey hate the fact that variables are Strings :D )
-- Use Maybe to filter clauses of length 0 and Either to partition clauses from unit-clauses.
-- This way, only one pass is needed instead of two (first partition, then filter, or vice versa).
constructBijection :: ConjForm -> Maybe (CNF, VarMapping)
constructBijection cj = if null cj' then return (SQ.empty, IM.empty) 
                                    else if cj' == [[]] then Nothing 
                                    else return (cnf, mapIN)
        where cj' = simpConj cj -- remove instances of bottom and stuff like that :D
              vars= zip [1..] $ nub $ map getV $ concat cj'
              mapIN  = IM.fromAscList vars -- used for reconstruction
              mapNI = M.fromList $ map swap vars -- for MapClause creation
              (units,clauses) = partitionEithers $ catMaybes $ map (\x -> toMC x mapNI) cj'
              cnf = SQ.fromList units >< SQ.fromList clauses -- faster than concatinating lists

toMC :: Clause -> Map Name Int -> Maybe (Either MapClause MapClause)
toMC cl rvm = case length c of
                0 -> Nothing
                1 -> return $ Left (1, (t', f'))
                k -> return $ Right (k, (t',f'))
        where c = nub cl -- if an element is contained twice pos/neg the length-invariant breaks :,)
              t = filter getB c  
              f = filter (not.getB) c
              t' = IS.fromList $ map (\x -> (fromJust $ M.lookup (getV x) rvm)) t
              f' = IS.fromList $ map (\x -> (fromJust $ M.lookup (getV x) rvm)) f

getB :: Literal -> Bool
getB (Pos _) = True
getB (Neg _) = False

-- This assumes that the VarMapping is correct, so throw an error
-- if any variable is not contained
reconstruct :: Eval -> VarMapping -> Valuation
reconstruct ev vm = map toEv $ F.toList ev
        where toEv (i,b) = (fromJust $ vm !? i, b)


{-H9.2.6-}
sat :: Form -> Maybe Valuation
sat = satConj.cnf 

{-TTEW-}