module Exercise_9 where
{-WETT-}

import           Data.List
import           Data.Function
import           Data.Maybe
import           Data.Ord
--import           Control.Monad
import           Control.Arrow
--import           Data.Array
import           Data.Graph
--import Data.Containers.ListUtils

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]

instance Ord Atom where
    compare T       _       = GT
    compare (Var _) T       = LT
    compare (Var n) (Var m) = compare m n

instance Ord Literal where
    compare (Pos _) (Neg _) = LT
    compare (Neg _) (Pos _) = GT
    compare (Pos a) (Pos b) = compare b a
    compare (Neg a) (Neg b) = compare b a

{-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
-- TODO: use maybe to accelerate subst by earily exit
substConj :: Valuation -> ConjForm -> ConjForm
substConj = map . substClause

{-H9.2.1-}
isTop :: Literal -> Bool
isTop (Pos T) = True
isTop _       = False

isBottom :: Literal -> Bool
isBottom (Neg T) = True
isBottom _       = False

noTop :: Clause -> Bool
noTop = all (not . isTop)

parseBottom :: Clause -> Clause
parseBottom = filter (not . isBottom)

simpConj :: ConjForm -> ConjForm
simpConj f = if any null g then [[]] else g
    where g = map parseBottom $ filter noTop f

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

{-H9.2.3-}
isVar :: Literal -> Bool
isVar (Pos T) = False
isVar (Neg T) = False
isVar _       = True

selectV :: ConjForm -> Maybe (Name, Bool)
selectV scf = case find isVar (concat scf) of
    Nothing               -> Nothing
    Just (Pos (Var name)) -> Just (name, True)
    Just (Neg (Var name)) -> Just (name, True)
    _                     -> undefined

{-H9.2.5-}
-------------------- Tester Functions --------------------
satConjNaive :: ConjForm -> Maybe Valuation
satConjNaive cf
    | null scf = Just []
    | any null scf = Nothing
    | otherwise = case selectV scf of
        Nothing     -> Nothing
        Just (n, b) -> case satConjNaive (substConj [(n, b)] scf) of
            Just val -> Just ((n, b) : val)
            Nothing  -> case satConjNaive (substConj [(n, not b)] scf) of
                Nothing  -> Nothing
                Just val -> Just ((n, not b) : val)
    where scf = simpConj cf


satConjOri :: ConjForm -> Maybe Valuation
satConjOri conjform = case satBase form of
    Nothing -> Nothing
    Just vs -> Just $ val ++ vs
    where (val, form) = preProcess conjform

-------------------- Main Function --------------------

satConj :: ConjForm -> Maybe Valuation
satConj conjform
    | null form = Just val
    | any null form = Nothing
    | otherwise = fmap (val ++)
        $   try (all ((2 ==) . length))  sat2CNF
        >>? try (all isJust . parseHorn) satHorn
        >>? try (all isJust . parseDual) satDual
        >>? satFreq form
        >>? Nothing
  where
    (val, form) = preProcess conjform
    -- use a custom arrow to avoid too much case..of
    Nothing >>? next = next
    Just a  >>? _    = Just a
    try test satStrategy = if test form then satStrategy form else Nothing

-- always handles the shortest clause first by sorting everytime
-- insertion sort may be faster in this case (?)
satBase :: ConjForm -> Maybe Valuation
satBase []              = Just []
satBase ([]       : _ ) = Nothing
satBase ((l : ls) : cf) = case sc (substConj [(name, bool)] cf) of
    Just val -> Just ((name, bool) : val)
    Nothing  -> case sc (substConj [(name, not bool)] (ls : cf)) of
        Just val -> Just ((name, not bool) : val)
        Nothing  -> Nothing
  where
    sc           = satBase . sortOn length . simplify
    (name, bool) = fromJust $ parseLiteral l

-------------------- Pure Literal --------------------

preProcess :: ConjForm -> (Valuation, ConjForm)
preProcess conjform = (puried, simplify $ substConj puried simple)
  where
    puried = purify simple
    simple = sortOn length $ simplify conjform

simplify :: ConjForm -> ConjForm
simplify form = fromMaybe [[]] $ simplify' form

simplify' :: ConjForm -> Maybe ConjForm
simplify' []       = Just []
simplify' ([] : _) = Nothing
simplify' (f : form)
    | Pos T `elem` f = simplify' form
    | otherwise = case simplify' form of
        Nothing -> Nothing
        Just fm -> Just $ filter (not . isBottom) f : fm

{-# INLINE purify #-}
purify :: ConjForm -> Valuation
purify =
    map toValuation
        . concat
        . filter ((1 ==) . length)
        . groupBy (on (==) getName)
        . sortOn getName
        . nub
        . concat

-------------------- Frequent Literal --------------------

{- This version always handles only the one most frequent element
-- to not influence performance too much -}
satFreq :: ConjForm -> Maybe Valuation
satFreq form = case satBase (simplify $ substConj [l] form) of
    Just v  -> Just $ l : v
    Nothing -> case satBase (simplify $ substConj [ol] form) of
        Just v  -> Just $ ol : v
        Nothing -> Nothing
  where
    l  = toValuation $ frequency form
    ol = opposite l

frequency :: ConjForm -> Literal
frequency = head . maximumBy (comparing length) . group . sort . concat

{- This version handles at most two frequent elements
-- too many elements will have exponential complexity 
-- due to the implementation of getBoolSeq -}

-- frequency :: ConjForm -> [(Literal, Int)]
-- frequency = map (head &&& length) . group . sort . concat

-- -- assume that frequent means occurs in at least half of the clauses
-- freqElemt :: ConjForm -> [Literal]
-- freqElemt form = map fst $ filter ((> lim) . snd) freq
--   where
--     freq = frequency form
--     lim  = length freq `div` 2

-- satFreq :: ConjForm -> Maybe Valuation
-- satFreq form = if null frequentElement
--     then Nothing
--     else findFreq form (parseFreq frequentElement)
--     where frequentElement = take 2 $ freqElemt form
--     -- take at most 2 to avoid exponential runtime

-- findFreq :: ConjForm -> [Valuation] -> Maybe Valuation
-- findFreq _    []       = Nothing
-- findFreq form (l : ls) = case satConj' (simplify $ substConj l form) of
--     Just v  -> Just $ l ++ v
--     Nothing -> findFreq form ls

-- -- assume no T
-- parseFreq :: [Literal] -> [Valuation]
-- parseFreq ls = zipWith zip (replicate (2 ^ len) (mapMaybe getName ls))
--     $ getBoolSeq len
--     where len = length ls

-- getBoolSeq :: Int -> [[Bool]]
-- getBoolSeq i = map (map toEnum) $ replicateM i [0, 1]

-------------------- (dual) Horn Clause --------------------

satHorn :: ConjForm -> Maybe Valuation
satHorn form = combine $ foldTillEq (Just [], sort $ catMaybes pf)
    where pf = parseHorn form

-- split the form into pos and neg
{-# INLINE parseHorn #-}
parseHorn :: ConjForm -> [Maybe (Maybe Literal, Clause)]
parseHorn = map
    (\form ->
        let (pos, neg) = partition isPos form
        in  case pos of
                []  -> Just (Nothing, sort $ map toPos $ nub neg)
                [p] -> Just (Just p, sort $ map toPos $ nub neg)
                _   -> Nothing -- not horn-sat
    )

-- just negate every literal and reuse horn works
satDual :: ConjForm -> Maybe Valuation
satDual form = combine $ foldTillEq (Just [], sort $ catMaybes pf)
    where pf = parseDual form

{-# INLINE parseDual #-}
parseDual :: ConjForm -> [Maybe (Maybe Literal, Clause)]
parseDual = map
    (\form ->
        let (neg, pos) = partition (not . isPos) form
        in  case neg of
                []  -> Just (Nothing, sort $ map toNeg $ nub pos)
                [n] -> Just (Just n, sort $ map toNeg $ nub pos)
                _   -> Nothing -- not dual-horn-sat
    )

foldTillEq
    :: (Maybe Clause, [(Maybe Literal, Clause)])
    -> (Maybe Clause, [(Maybe Literal, Clause)])
foldTillEq x = if hornFolder x == x then x else foldTillEq (hornFolder x)
  where
    hornFolder (Just val, xs) = case find (any (`elem` val) . snd) xs of
        Just (  Nothing, _) -> (Nothing, xs)
        Just c@(Just l , _) -> (Just (l : val), deleteBy (on (==) fst) c xs)
        Nothing ->
            if any (null . snd) xs then (Nothing, xs) else (Just val, xs)
            -- 
    hornFolder a = a -- Nothing

combine :: (Maybe Clause, [(Maybe Literal, Clause)]) -> Maybe Valuation
combine (Nothing, _) = Nothing
combine (Just val, xs) =
    Just $ map toValuation $ nub (val ++ concatMap (map negLiteral . snd) xs)

-------------------- 2-CNF --------------------

-- use graph theory to solve 2-cnf
sat2CNF :: ConjForm -> Maybe Valuation
sat2CNF form = case cnfSolver nodes graph of
    Nothing -> Nothing
    Just vs -> (vs ++) <$> satBase (simplify $ substConj vs form)
  where
    nodes = getNode form
    graph = buildG bound edge
    bound = getBounds edge
    edge  = parseForm form

cnfSolver :: [Vertex] -> Graph -> Maybe Valuation
cnfSolver [] _ = Just []
cnfSolver (v : vs) g | to && from = Nothing
                     | to         = ((show $ abs v, False) :) <$> cnfSolver vs g
                     | from       = ((show $ abs v, True) :) <$> cnfSolver vs g
                     | otherwise  = cnfSolver vs g
    --  | otherwise  = (vertexToValuation v :) <$> cnfSolver vs g
    -- have no idea how to handle the vertices not in scc
    -- so just solve them by basic approach



  where
    to   = path g v (negate v)
    from = path g (negate v) v

parseForm :: ConjForm -> [Edge]
parseForm [] = []
parseForm (f : form) =
    (negate $ getValueOfName a, getValueOfName b)
        : (negate $ getValueOfName b, getValueOfName a)
        : parseForm form
    where (a, b) = (head f, last f)

vertexToValuation :: Vertex -> (Name, Bool)
vertexToValuation v | v > 0     = (show v, True)
                    | otherwise = (show $ negate v, False)

getNode :: ConjForm -> [Vertex]
getNode = intNub . map (abs . getValueOfName) . concat

{-# INLINABLE getBounds #-}
getBounds :: [Edge] -> Bounds
getBounds = (minimum &&& maximum) . concatMap (\(a, b) -> [a, b])

-------------------- General Helper --------------------

parseLiteral :: Literal -> Maybe (Name, Bool)
parseLiteral (Pos (Var name)) = Just (name, True)
parseLiteral (Neg (Var name)) = Just (name, False)
parseLiteral _                = Nothing

toValuation :: Literal -> (Name, Bool)
toValuation (Pos (Var name)) = (name, True)
toValuation (Neg (Var name)) = (name, False)
toValuation _                = undefined

opposite :: (Name, Bool) -> (Name, Bool)
opposite (name, bool) = (name, not bool)

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

toPos :: Literal -> Literal
toPos (Pos l) = Pos l
toPos (Neg l) = Pos l

toNeg :: Literal -> Literal
toNeg (Pos l) = Neg l
toNeg (Neg l) = Neg l

negLiteral :: Literal -> Literal
negLiteral (Pos l) = Neg l
negLiteral (Neg l) = Pos l

getName :: Literal -> Maybe Name
getName (Pos T         ) = Nothing
getName (Neg T         ) = Nothing
getName (Pos (Var name)) = Just name
getName (Neg (Var name)) = Just name

getValueOfName :: Literal -> Int
getValueOfName (Pos (Var name)) = read name
getValueOfName (Neg (Var name)) = negate $ read name
getValueOfName _                = undefined

-- stolen from Exercise_5 :)
intNub :: [Int] -> [Int]
intNub = map head . group . sort


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

{-TTEW-}