import Control.Applicative (Applicative, Alternative(..))
import Data.Char (isAlphaNum)
import Data.Functor (Functor)
import Data.List (intercalate)
import Data.Maybe (fromJust)
import qualified Data.Set as S

newtype Parser symbol output = Parser
  { runParser :: [symbol] -> Maybe (output, [symbol])
  }

instance Functor (Parser s) where
  fmap :: (a -> b) -> Parser s a -> Parser s b
  fmap f (Parser p) = Parser 
    { runParser = \xs -> do
      (result, xs') <- p xs
      return (f result, xs')
    }

instance Applicative (Parser s) where
  pure :: a -> Parser s a
  pure result = Parser $ \xs -> Just (result, xs)

  (Parser p1) <*> (Parser p2) =
    Parser $ \xs -> do
      (f, xs') <- p1 xs
      (result, xs'') <- p2 xs'
      return (f result, xs'')

instance Alternative (Parser s) where
  empty :: Parser s a
  empty = Parser $ const Nothing

  (Parser p1) <|> (Parser p2) =
    Parser $ \xs -> p1 xs <|> p2 xs

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

parseToken :: String -> Parser Char String
parseToken token = Parser parse
  where
    n = length token
    parse xs
      | token == take n xs = Just (token, drop n xs)
      | otherwise = Nothing

parseWhile :: (Char -> Bool) -> Parser Char String
parseWhile check = Parser parse
  where
    parse [] = Nothing
    parse (token:rest)
      | check token = 
        case parse rest of
          Nothing -> Just ([token], rest)
          Just (tokens, rest') -> Just (token:tokens, rest')
      | otherwise = Nothing

parseBrackets :: Parser Char a -> Parser Char a
parseBrackets p = parseToken "(" *> p <* parseToken ")"

parseInfix :: String -> Parser Char a -> Parser Char a -> Parser Char (a, a)
parseInfix token p1 p2 = (,) <$> (p1 <* parseToken token) <*> p2

parseConnective :: String -> Parser Char a -> Parser Char a -> Parser Char (a, a)
parseConnective token p1 p2 = parseBrackets $ parseInfix token p1 p2

stmtAtom :: Parser Char Statement
stmtAtom = Atom <$> parse
  where
    parse = parseWhile $ \char -> isAlphaNum char || char == '_'

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

stmtAnd :: Parser Char Statement
stmtAnd = uncurry And <$> parseConnective "&" stmt stmt

stmtOr :: Parser Char Statement
stmtOr = uncurry Or <$> parseConnective "|" stmt stmt

stmtImplies :: Parser Char Statement
stmtImplies = uncurry Implies <$> parseConnective "->" stmt stmt

stmtIff :: Parser Char Statement
stmtIff = uncurry Iff <$> parseConnective "<->" stmt stmt

stmt :: Parser Char Statement
stmt = stmtAtom <|> stmtNot <|> stmtAnd <|> stmtOr <|> stmtImplies <|> stmtIff

{-
ghci> runParser stmt "(!a<->((!(p->q)&x)|y))"
Just (Iff (Not (Atom "a")) (Or (And (Not (Implies (Atom "p") (Atom "q"))) (Atom "x")) (Atom "y")),"")
-}

p :: Statement
p = fromJust $ fst <$> runParser stmt "((a->b)<->!c)"

q :: Statement
q = fromJust $ fst <$> runParser stmt "(!a<->((!(p->q)&x)|y))"

atoms :: Statement -> S.Set String
atoms (Atom key)      = S.singleton key
atoms (Not s)         = atoms s
atoms (And s1 s2)     = S.union (atoms s1) (atoms s2)
atoms (Or s1 s2)      = S.union (atoms s1) (atoms s2)
atoms (Implies s1 s2) = S.union (atoms s1) (atoms s2)
atoms (Iff s1 s2)     = S.union (atoms s1) (atoms s2)

eval :: [(String, Bool)] -> Statement -> Maybe Bool
eval assignments = aux
  where
    aux (Atom key)      = 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
    atomsList = S.toList $ atoms s
    values = [fromJust $ eval assignments s | assignments <- enumerate $ atomsList]

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

    start = [(key, False) | key <- keys]

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

{-
 ghci> bucket $ fromJust $ fst <$> runParser stmt "(p<->p)"
Tautology
ghci> bucket p
Contingent
ghci> bucket q
Contingent
ghci> bucket $ fromJust $ fst <$> runParser stmt "(p<->!p)"
Contradiction
ghci> bucket $ fromJust $ fst <$> runParser stmt "(p|!p)"
Tautology
ghci> bucket $ fromJust $ fst <$> runParser stmt "(p->p)"
Tautology
ghci> bucket $ fromJust $ fst <$> runParser stmt "(p&!p)"
Contradiction
ghci> bucket $ fromJust $ fst <$> runParser stmt "(!(p->q)&!p)"
Contradiction
ghci> bucket $ fromJust $ fst <$> runParser stmt "(p|(p->q))"
Tautology
ghci> bucket $ fromJust $ fst <$> runParser stmt "((p->q)->(!p->r))"
Contingent
ghci> bucket $ fromJust $ fst <$> runParser stmt "(!(a&b)<->(!a|!b))"
Tautology
ghci> bucket $ fromJust $ fst <$> runParser stmt "(!(a|b)<->(!a&!b))"
Tautology
-}

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

serialize :: SerializeFormat -> Statement -> String
serialize = aux
  where
    aux Ascii = ascii
    aux Latex = latex

    ascii (Atom string)   = string
    ascii (Not s)         = "!" <> ascii s
    ascii (And s1 s2)     = connective "&" (ascii s1) (ascii s2)
    ascii (Or s1 s2)      = connective "|" (ascii s1) (ascii s2)
    ascii (Implies s1 s2) = connective "->" (ascii s1) (ascii s2)
    ascii (Iff s1 s2)     = connective "<->" (ascii s1) (ascii s2)

    latex (Atom string)   = string
    latex (Not s)         = "\\neg " <> latex s
    latex (And s1 s2)     = connective "\\cap " (latex s1) (latex s2)
    latex (Or s1 s2)      = connective "\\cup " (latex s1) (latex s2)
    latex (Implies s1 s2) = connective "\\to " (latex s1) (latex s2)
    latex (Iff s1 s2)     = connective "\\leftrightarrow " (latex s1) (latex s2)

    connective token s1 s2 = "(" <> s1 <> token <> s2 <> ")"

truthtable :: Statement -> String
truthtable s = open <> header <> "\\hline\n" <> body <> close
  where
    open =
      "\\begin{tabular}{" <>
      replicate (length atomsList) 'c' <>
      "|" <>
      replicate (length serial) 'c' <>
      "}\n"

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

    serial = serialize Latex s

    atomsList :: [String]
    atomsList = S.toAscList $ atoms s

    header :: String
    header = intercalate "&" atomsList <> " \\\\\n"

    body = concat $ map line $ enumerate atomsList

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

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

    parts assignments = [""]