module Logic.Statement where

import Logic.Parse

import Control.Applicative (Alternative((<|>), some))
import Data.Char (isAlphaNum)
import Data.List (intercalate)
import Data.Either (fromRight)
import Data.Set (singleton, union, toAscList)

data Statement
  = Atom String
  | Not Statement
  | And Statement Statement
  | Or Statement Statement
  | Implies Statement Statement
  | Iff Statement Statement
  deriving (Show, Eq)

stmtAtom :: Parser Statement
stmtAtom = Atom <$> parse
  where
    parse = some $ parseIf "statement variable" $ \char -> isAlphaNum char || char == '_'

stmtNot :: Parser Statement
stmtNot = Not <$> (parseToken "!" *> stmt)

stmtBinary :: Parser Statement
stmtBinary = parseToken "(" *> body <* parseToken ")"
  where
    body = (\s1 f s2 -> f s1 s2) <$> stmt <*> parseConnective <*> stmt

    parseConnective =
      fmap (const And) (parseToken "&")
      <|> fmap (const Or) (parseToken "|")
      <|> fmap (const Implies) (parseToken "->")
      <|> fmap (const Iff) (parseToken "<->")
      <|> fail

    fail = Parser $ \input -> Left $ expected "connective" input

stmt :: Parser Statement
stmt = Parser $ \input ->
  let
    parser =
      case inputSeq input of
        [] -> fail
        ('!':_) -> stmtNot
        ('(':_) -> stmtBinary
        _ -> stmtAtom <|> fail
  in runParser parser input
  where
    fail = Parser $ \input -> Left $ expected "statement" input

p :: Statement
p = fromRight undefined $ eof stmt $ mkInput "((a->b)<->!c)"

q :: Statement
q = fromRight undefined $ eof stmt $ mkInput "(!a<->((!(p->q)&x)|y))"

atoms :: Statement -> [String]
atoms = toAscList . mkSet
  where
    mkSet (Atom key)      = singleton key
    mkSet (Not s)         = mkSet s
    mkSet (And s1 s2)     = union (mkSet s1) (mkSet s2)
    mkSet (Or s1 s2)      = union (mkSet s1) (mkSet s2)
    mkSet (Implies s1 s2) = union (mkSet s1) (mkSet s2)
    mkSet (Iff s1 s2)     = union (mkSet s1) (mkSet s2)

eval :: [(String, Bool)] -> Statement -> Either String Bool
eval assignments = aux
  where
    aux (Atom key)      = maybe (Left key) Right $ lookup key assignments
    aux (Not s)         = not <$> aux s
    aux (And s1 s2)     = (&&) <$> aux s1 <*> aux s2
    aux (Or s1 s2)      = (||) <$> aux s1 <*> aux s2
    aux (Implies s1 s2) = not <$> ((&&) <$> aux s1 <*> (not <$> aux s2))
    aux (Iff s1 s2)     = (==) <$> aux s1 <*> aux s2

data Bucket
  = Tautology
  | Contradiction
  | Contingent
  deriving (Eq, Show)

bucket :: Statement -> Bucket
bucket s
  | and values = Tautology
  | all not values = Contradiction
  | otherwise = Contingent
  where
    values = [fromRight undefined $ eval assignments s | assignments <- enumerate $ atoms s]

enumerate :: [a] -> [[(a, Bool)]]
enumerate keys = map reverse $ aux start
  where
    aux assignments = (assignments:) $
      case next assignments of
        Nothing -> []
        Just (assignments') -> aux assignments'

    start = map (, False) keys

    next [] = Nothing
    next ((key, False):rest) = Just $ (key, True):rest
    next ((key, True):rest) = ((key, False):) <$> (next rest)

data SerializeFormat
  = Ascii
  | Latex
  deriving (Show, Eq)

data Cell
  = Main String
  | NotMain String
  deriving (Show, Eq)

fromCell :: Cell -> String
fromCell (Main x) = x
fromCell (NotMain x) = x

fromCellLatex :: Cell -> String
fromCellLatex (Main x) = "\\textbf " <> x
fromCellLatex (NotMain x) = x

isMain :: Cell -> Bool
isMain (Main _) = True
isMain _ = False

serializeCells :: SerializeFormat -> Statement -> [Cell]
serializeCells = aux Main
  where
    aux mkCell Ascii = ascii mkCell
    aux mkCell Latex = latex mkCell

    ascii' = ascii NotMain
    latex' = latex NotMain

    ascii mkCell (Atom key)      = [mkCell key]
    ascii mkCell (Not s)         = [mkCell "!"] ++ ascii' s
    ascii mkCell (And s1 s2)     = connective (mkCell "&") (ascii' s1) (ascii' s2)
    ascii mkCell (Or s1 s2)      = connective (mkCell "|") (ascii' s1) (ascii' s2)
    ascii mkCell (Implies s1 s2) = connective (mkCell "->") (ascii' s1) (ascii' s2)
    ascii mkCell (Iff s1 s2)     = connective (mkCell "<->") (ascii' s1) (ascii' s2)

    latex mkCell (Atom key)      = [mkCell key]
    latex mkCell (Not s)         = [mkCell "\\neg "] ++ latex' s
    latex mkCell (And s1 s2)     = connective (mkCell "\\land ") (latex' s1) (latex' s2)
    latex mkCell (Or s1 s2)      = connective (mkCell "\\lor ") (latex' s1) (latex' s2)
    latex mkCell (Implies s1 s2) = connective (mkCell "\\to ") (latex' s1) (latex' s2)
    latex mkCell (Iff s1 s2)     = connective (mkCell "\\leftrightarrow ") (latex' s1) (latex' s2)

    connective middle s1 s2 = [NotMain "("] ++ s1 ++ [middle] ++ s2 ++ [NotMain ")"]

serialize :: SerializeFormat -> Statement -> String
serialize fmt s = concat $ map fromCell $ serializeCells fmt s

latexTruthTableMainColumnIndex :: Statement -> Int
latexTruthTableMainColumnIndex s =
  fst $ head $ filter (\(i, cell) -> isMain cell) $ zip [0..] $ serializeCells Latex s

serializeLatexTruthTableRow :: [(String, Bool)] -> Statement -> Either String [String]
serializeLatexTruthTableRow ass s = map fromCellLatex <$> latexRow Main ass s
  where
    latexRow mkCell ass (Atom key)      = list <$> mkCell <$> toInt <$> eval ass (Atom key)
    latexRow mkCell ass (Not s)         = (list <$> mkCell <$> toInt <$> not <$> eval ass s) <> latexRow NotMain ass s
    latexRow mkCell ass (And s1 s2)     = latexRowBinaryConnective (&&) mkCell ass s1 s2
    latexRow mkCell ass (Or s1 s2)      = latexRowBinaryConnective (||) mkCell ass s1 s2
    latexRow mkCell ass (Implies s1 s2) = latexRowBinaryConnective implies mkCell ass s1 s2
    latexRow mkCell ass (Iff s1 s2)     = latexRowBinaryConnective (==) mkCell ass s1 s2

    latexRowBinaryConnective op mkCell ass s1 s2 =
      (
        \subrow1 subrow2 result ->
          [NotMain " "] ++ subrow1 ++ [mkCell $ toInt result] ++ subrow2 ++ [NotMain " "]
      ) <$>
      latexRow NotMain ass s1 <*> latexRow NotMain ass s2 <*>
      (op <$> eval ass s1 <*> eval ass s2)

    toInt :: Bool -> String
    toInt False = "0"
    toInt True = "1"

    list :: a -> [a]
    list x = [x]

    implies :: Bool -> Bool -> Bool
    implies b1 b2 = not b1 || b2

truthtable :: Statement -> String
truthtable s = open <> header <> "\\hline\n" <> body <> close
  where
    mainIndex = latexTruthTableMainColumnIndex s

    cellsSpec
      | mainIndex == 0 && length serial <= 1 = "c"
      | mainIndex == 0 = "c|" <> replicate (length serial - 1) 'c'
      | mainIndex == length serial - 1 = replicate (length serial - 1) 'c' <> "|c"
      | otherwise = replicate mainIndex 'c' <> "|c|" <> replicate (length serial - mainIndex - 1) 'c'

    open =
      "\\begin{tabular}{" <>
      replicate (length $ atoms s) 'c' <>
      "||" <>
      cellsSpec <>
      "}\n"

    close = "\\end{tabular}\n\n"

    serial :: [Cell]
    serial = serializeCells Latex s

    header :: String
    header =
      intercalate " & " (map dollars $ atoms s) <>
      " & " <>
      intercalate " & " (map dollars $ map fromCell serial) <>
      " \\\\\n"

    dollars :: String -> String
    dollars string = "$" <> string <> "$"

    body :: String
    body = concat $ map line $ enumerate $ atoms s

    line assignments =
      intercalate " & " (bools assignments) <>
      " &" <>
      intercalate "&" (cells assignments) <>
      "\\\\\n"

    bools assignments = [if bool then "1" else "0" | (key, bool) <- assignments]

    cells assignments =
      (\xs -> [if x /= " " then " " <> x <> " " else " " | x <- xs]) $
      fromRight undefined $ serializeLatexTruthTableRow assignments s