content: {-
      Copyright 2018 Google LLC
  
      Licensed under the Apache License, Version 2.0 (the "License");
      you may not use this file except in compliance with the License.
      You may obtain a copy of the License at
  
          https://www.apache.org/licenses/LICENSE-2.0
  
      Unless required by applicable law or agreed to in writing, software
      distributed under the License is distributed on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      See the License for the specific language governing permissions and
      limitations under the License.
  -}    
  
  module Main where
  
  import Control.Exception.Base
  import Control.Monad
  import Data.Array.IO
  import qualified Data.Char as Char
  import qualified Data.Logic.Propositional as Prop
  import qualified Data.Logic.Propositional.NormalForms as NormalForms
  import qualified Data.List as List
  import qualified Data.Map as Map
  import qualified Data.Maybe as Maybe
  import qualified Data.Set as Set
  import qualified Data.Tuple as Tuple
  import MiniSat
  import System.Environment
  import System.IO
  import System.Process
  import System.Random
  import Text.ParserCombinators.Parsec
  import Text.Parsec.Combinator
  
  ------------------------------------ Types ------------------------------------
  
  {-
  An Entailment is a struct containing: 
     two propositions A and B
     whether A entails B (written A ⊧ B)
     three statistics (used as simple heuristic estimates):
       whether length(A) >= length(B)
       whether vars(B) ⊆ vars(A)
       whether literals(B*) ⊆ literals(A*)
         where A* is negation-normal-form version of A
  -}       
  type Entailment = (Prop.Expr, Prop.Expr, Int, Int, Int, Int)
  
  {-
  A FourTuple is four propositions (A, B, A*, B*) such that:
    A ⊧ B
    A* ⊧ B*
    A ⊭ B*
    A* ⊭ B
  -}
  type FourTuple = (Prop.Expr, Prop.Expr, Prop.Expr, Prop.Expr)
  
  ------------------------------------ main -----------------------------------
  
  {-
  Usage:
    first argument: number of pairs to generate
    second argument: max number of variables in a formula
    third argument: max complexity of a formula
  
  NB four-tuples are extracted from generated pairs.
  But many pairs are unsuitable for four-tuples.
  So if you want 100 entailment lines, you need 100/4 = 25 four-tuples,
    and you might need 200 pairs to produce 25 four-tuples.
    (The ratio of pairs-to-four-tuples depends on the complexity of the formulas)
  -}
  main :: IO ()
  main = do
      args <- getArgs
      case args of
        [] -> make_four_tuples 100 20 30
        [ns] -> do
          let n = read ns :: Int
          make_four_tuples n 20 30
        [ns, vs] -> do
          let n = read ns :: Int
          let v = read vs :: Int
          make_four_tuples n v 30
        [ns, vs, ks] -> do
          let n = read ns :: Int
          let v = read vs :: Int
          let k = read ks :: Int
          make_four_tuples n v k
  
  ------------------------------------ Generation -------------------------------
  
  gen_formula :: RandomGen r => r -> Int -> [Char]  -> (Prop.Expr, r) 
  gen_formula r complexity vs = gen_f r complexity vs
  
  gen_f :: RandomGen r => r -> Int -> [Char] -> (Prop.Expr, r)
  gen_f r 0 vs = gen_atom r vs
  gen_f r n vs = f r3 where
    n1 = n - 1
    (na, r2) = pick [0 .. n1] r
    nb = n1 - na
    (f, r3) = pick [gen_unary Prop.Negation n1 vs, 
                    gen_binary Prop.Conjunction na nb vs, 
                    gen_binary Prop.Disjunction na nb vs, 
                    gen_binary Prop.Conditional na nb vs] r2
  
  gen_atom :: RandomGen r => r -> [Char] -> (Prop.Expr, r)
  gen_atom r vs = (Prop.Variable (Prop.Var v), r2) where
    (v, r2) = pick vs r
  
  gen_unary :: RandomGen r => (Prop.Expr -> Prop.Expr) -> Int -> [Char] -> r -> (Prop.Expr, r)
  gen_unary f n vs r = (f p, r') where
    (p, r') = gen_f r n vs
  
  gen_binary :: RandomGen r => (Prop.Expr -> Prop.Expr -> Prop.Expr) -> Int -> Int -> [Char] -> r -> (Prop.Expr, r)
  gen_binary f na nb vs r = (f p q, r3) where
    (p, r2) = gen_f r na vs
    (q, r3) = gen_f r2 nb vs
  
  alphabet :: [Char]
  alphabet = ['a'..'z']
  
  make_exprs :: Int -> Int -> Int -> Int -> IO (Prop.Expr, Prop.Expr)
  make_exprs c1 c2 num_vars n = do
    r <- newStdGen
    let (vs, _, r2) = take_n alphabet num_vars r
    let (p1, r3) = gen_formula r2 c1 vs
    let (p2, _) = gen_formula r3 c2 vs
    return (p1, p2)
  
  pick :: RandomGen r => [a] -> r -> (a, r)
  pick [] _ = error "pick wrongly given empty list"
  pick as r =
      let indices = (0, length as-1)
          (i, r') = randomR indices r
      in  (as !! i, r')
  
  pick_freq :: RandomGen r => [(Int, a)] -> r -> (a, r)   
  pick_freq [] _ = error "pick_freq wrongly given empty list"
  pick_freq as r = (pick_n n as, r') where
    (n, r') = randomR (0, tot-1) r
    tot = List.foldl' (\b -> (\(a,_) -> a + b)) 0 as
  
  pick_n :: Int -> [(Int, a)] -> a
  pick_n _ [] = error "Empty list"
  pick_n x ((y,a):_) | x < y = a
  pick_n x ((y,_):ps) | otherwise = pick_n (x-y) ps
  
  pick_take :: RandomGen r => [a] -> r -> (a, [a], r)
  pick_take [] _ = error "pick_take wrongly given empty list"
  pick_take as r = 
      let indices = (0, length as-1)
          (i, r') = randomR indices r
      in  (as !! i, take i as ++ drop (i+1) as, r')
  
  take_n :: RandomGen r => [a] -> Int -> r -> ([a], [a], r)
  take_n as n r = take_n2 as n r []
  
  take_n2 :: RandomGen r => [a] -> Int -> r -> [a] -> ([a], [a], r)
  take_n2 as n r acc | length acc == n = (acc, as, r)
  take_n2 as n r acc | otherwise = take_n2 as' n r' (a:acc) where
    (a, as', r') = pick_take as r
  
  ------------------------------------ Entails ----------------------------------
  
  entails :: Prop.Expr -> Prop.Expr -> IO Bool
  entails a b = do
    let p = Prop.Conjunction a (Prop.Negation b)
    b <- sat p
    return (not b)
  
  equiv :: Prop.Expr -> Prop.Expr -> IO Bool
  equiv a b = do
    let p = Prop.Conjunction a (Prop.Negation b)
    let q = Prop.Conjunction b (Prop.Negation a)  
    b <- sat (Prop.Disjunction p q)
    return (not b)
  
  sat :: Prop.Expr -> IO Bool  
  sat p = do
    let (n, cs) = convert_to_dimacs p
    s <- make_solver cs n
    b <- solve s []
    deleteSolver s
    return b
  
  convert_to_dimacs :: Prop.Expr -> (Int, [[Int]])
  convert_to_dimacs = expr_to_dimacs . toCNF
  
  toCNF :: Prop.Expr -> Prop.Expr  
  toCNF = fixed_point NormalForms.toCNF
  
  fixed_point :: Eq a => (a->a) -> a -> a
  fixed_point f a = 
      if f a == a then a
      else fixed_point f (f a)
  
  expr_to_dimacs :: Prop.Expr -> (Int, [[Int]])
  expr_to_dimacs e = to_dimacs e (zip (extract_vars e) [1..])
  
  to_dimacs :: Prop.Expr -> [(Char, Int)] -> (Int, [[Int]])
  to_dimacs e m = (length m, to_dimacs_c e m)
  
  to_dimacs_c :: Prop.Expr -> [(Char, Int)] -> [[Int]]
  to_dimacs_c (Prop.Variable (Prop.Var c)) m = case lookup c m of
    Just i -> [[i]]
  to_dimacs_c (Prop.Negation (Prop.Variable (Prop.Var c))) m = case lookup c m of
    Just i -> [[-i]]
  to_dimacs_c (Prop.Conjunction p q) m = to_dimacs_c p m ++ to_dimacs_c q m
  to_dimacs_c (Prop.Disjunction p q) m = [to_dimacs_d p m ++ to_dimacs_d q m]
  
  to_dimacs_d :: Prop.Expr -> [(Char, Int)] -> [Int]
  to_dimacs_d (Prop.Variable (Prop.Var c)) m = case lookup c m of
    Just i -> [i]
  to_dimacs_d (Prop.Negation (Prop.Variable (Prop.Var c))) m = case lookup c m of
    Just i -> [-i]
  to_dimacs_d (Prop.Disjunction p q) m = to_dimacs_d p m ++ to_dimacs_d q m
  
  extract_vars :: Prop.Expr -> [Char]
  extract_vars e = List.sort (List.nub (f e)) where
    f (Prop.Variable (Prop.Var c)) = [c]
    f (Prop.Negation p) = f p
    f (Prop.Conjunction p q ) = f p ++ f q
    f (Prop.Disjunction p q ) = f p ++ f q
    f (Prop.Conditional p q ) = f p ++ f q
  
  make_solver:: [[Int]] -> Int -> (IO Solver)
  make_solver baseClauses vars = do 
    s <- newSolver
    literals <- sequence [newLit s | i <- [1..vars]]
    sequence_ $ map ( (addClause s).(toLits)) baseClauses
    return s
  
  toLits :: [Int] -> [Lit] 
  toLits xs = map lit' xs where
    lit' x | x > 0 = minisat_mkLit $ MkVar (fromIntegral (x-1)) 
    lit' x | x < 0 = neg $ minisat_mkLit $ MkVar (fromIntegral (abs x - 1))
  
  ------------------------------------ show -------------------------------------
  
  showF :: Prop.Expr -> String  
  showF (Prop.Variable (Prop.Var c)) = [c]
  showF (Prop.Negation p) = "~(" ++ showF p ++ ")"
  showF (Prop.Conjunction p q) = "(" ++ showF p ++ "&" ++ showF q ++ ")"
  showF (Prop.Disjunction p q) = "(" ++ showF p ++ "|" ++ showF q ++ ")"
  showF (Prop.Conditional p q) = "(" ++ showF p ++ ">" ++ showF q ++ ")"
  
  ------------------------------------ check ------------------------------------
  
  verify_dataset :: String -> IO ()
  verify_dataset f = do
    ts <- parse_file f
    forM_ ts $ \(a, b, e, _, _, _) -> do
      let not_e = e == 0
      not_entails <- sat (Prop.Conjunction a (Prop.Negation b))
      putStrLn $ "Verifying " ++ show a ++ ", " ++ show b ++ ", " ++ show e
      case not_entails == not_e of
        True -> return ()
        False -> do
          putStrLn $ "Problem with " ++ show a ++ ", " ++ show b ++ ", " ++ show e
          error "Error"
    putStrLn "Ok"
  
  ------------------------------------ Parsing ----------------------------------
  
  parse_file :: String -> IO [Entailment]
  parse_file f = do
      x <- readFile f
      case (parse file f x) of
          Left e -> do
              putStrLn $ "Parse error: " ++ show e
              return []
          Right rs -> do
              return rs
  
  file :: Parser [Entailment]
  file = spaces >> parse_triples
  
  parse_triples :: Parser [Entailment]
  parse_triples = many1 parse_triple
  
  parse_triple :: Parser Entailment
  parse_triple = do
      p <- parse_formula
      lexeme ","
      q <- parse_formula
      lexeme ","
      n <- parse_int
      lexeme ","
      h1 <- parse_int
      lexeme ","
      h2 <- parse_int
      lexeme ","
      h3 <- parse_int
      return (p, q, n, h1, h2, h3)
  
  parse_formula :: Parser Prop.Expr
  parse_formula = do
    e <-  Text.ParserCombinators.Parsec.try parse_negation <|> 
          Text.ParserCombinators.Parsec.try parse_conjunction <|> 
          Text.ParserCombinators.Parsec.try parse_disjunction <|> 
          Text.ParserCombinators.Parsec.try parse_implication <|> 
          parse_atom
    return e
  
  comma = lexeme ","
        
  parse_atom :: Parser Prop.Expr
  parse_atom = do
      a <- identifierSymbol
      return $ Prop.Variable (Prop.Var a)
      
  identifierSymbol :: Parser Char
  identifierSymbol = alphaNum <|> (char '_')
  
  parse_negation :: Parser Prop.Expr
  parse_negation = do
    lexeme "~"
    lexeme "("
    p <- parse_formula
    lexeme ")"
    return $ Prop.Negation p
  
  parse_conjunction :: Parser Prop.Expr
  parse_conjunction = do
    lexeme "("
    p <- parse_formula
    lexeme "&"
    q <- parse_formula
    lexeme ")"
    return $ Prop.Conjunction p q
  
  parse_disjunction :: Parser Prop.Expr
  parse_disjunction = do
    lexeme "("
    p <- parse_formula
    lexeme "|"
    q <- parse_formula
    lexeme ")"
    return $ Prop.Disjunction p q
  
  parse_implication :: Parser Prop.Expr
  parse_implication = do
    lexeme "("
    p <- parse_formula
    lexeme ">"
    q <- parse_formula
    lexeme ")"
    return $ Prop.Conditional p q
  
  parse_int :: Parser Int
  parse_int = do
    e <- Text.ParserCombinators.Parsec.try parse_1 <|> parse_0
    return e
  
  parse_1 :: Parser Int
  parse_1 = do
    lexeme "1"
    return 1
  
  parse_0 :: Parser Int
  parse_0 = do
    lexeme "0"
    return 0
  
  lexeme :: String -> Parser String
  lexeme s = do
      string s
      spaces
      return s
  
  ------------------------------------ de Bruijn --------------------------------
  
  type SymbolMap = [(Prop.Var, Int)]
  
  deBruijn :: (Prop.Expr, SymbolMap) -> (Prop.Expr, SymbolMap)
  deBruijn (Prop.Negation p, m) = (Prop.Negation p', m') where (p', m') = deBruijn (p, m)
  deBruijn (Prop.Conjunction p q, m) = (Prop.Conjunction p' q', m'') where 
    (p', m') = deBruijn (p, m)
    (q', m'') = deBruijn (q, m')
  deBruijn (Prop.Disjunction p q, m) = (Prop.Disjunction p' q', m'') where 
    (p', m') = deBruijn (p, m)
    (q', m'') = deBruijn (q, m')
  deBruijn (Prop.Conditional p q, m) = (Prop.Conditional p' q', m'') where 
    (p', m') = deBruijn (p, m)
    (q', m'') = deBruijn (q, m')
  deBruijn (Prop.Variable p, m) = case lookup p m of
    Nothing -> (Prop.Variable p', m') where
      n = length m + 1
      p' = Prop.Var (Char.chr n)
      m' = (p, n) : m
    Just n -> (Prop.Variable p', m) where
      p' = Prop.Var (Char.chr n)
  
  ------------------------------------ Alpha-Equivalence ------------------------
  
  build_set :: String -> IO (Set.Set (Prop.Expr, Prop.Expr))
  build_set f = do
    ts <- parse_file f
    let fs = List.foldl' add_triple Set.empty ts
    return fs
  
  add_triple :: Set.Set (Prop.Expr, Prop.Expr) -> Entailment -> Set.Set (Prop.Expr, Prop.Expr)
  add_triple s (p,q,_,_,_,_) = Set.insert (p',q') s where
    (p', m) = deBruijn (p, [])
    (q', _) = deBruijn (q, m)
  
  check_triple :: Set.Set (Prop.Expr, Prop.Expr) -> String -> Entailment -> IO ()
  check_triple s f (a,b,n,h1,h2,h3) = do
    let (a', m) = deBruijn (a, [])
    let (b', _) = deBruijn (b, m)
    case Set.member (a',b') s of
      True -> putStrLn $ "Redundant: " ++ showF a ++ "," ++ showF b
      False -> appendFile f $ showF a ++ "," ++ showF b ++ "," ++ show n ++ "," ++ show h1 ++ "," ++ show h2 ++ "," ++ show h3 ++ "\n"
  
  filter_alpha_equivalences :: String -> String -> String -> IO ()  
  filter_alpha_equivalences training test results = do
    s <- build_set training
    ts <- parse_file test
    forM_ ts (check_triple s results)
  
  instance Ord Prop.Expr where
    compare (Prop.Variable p) (Prop.Variable q) = compare p q
    compare (Prop.Variable _) (Prop.Negation _) = LT
    compare (Prop.Variable _) (Prop.Conjunction _ _) = LT
    compare (Prop.Variable _) (Prop.Disjunction _ _) = LT
    compare (Prop.Variable _) (Prop.Conditional _ _) = LT
  
    compare (Prop.Negation _) (Prop.Variable _) = GT
    compare (Prop.Negation p) (Prop.Negation q) = compare p q
    compare (Prop.Negation _) (Prop.Conjunction _ _) = LT
    compare (Prop.Negation _) (Prop.Disjunction _ _) = LT
    compare (Prop.Negation _) (Prop.Conditional _ _) = LT
  
    compare (Prop.Conjunction _ _) (Prop.Variable _) = GT
    compare (Prop.Conjunction _ _) (Prop.Negation _) = GT
    compare (Prop.Conjunction p q) (Prop.Conjunction p' q') = compare (p,q) (p',q')
    compare (Prop.Conjunction _ _) (Prop.Disjunction _ _) = LT
    compare (Prop.Conjunction _ _) (Prop.Conditional _ _) = LT
  
    compare (Prop.Disjunction _ _) (Prop.Variable _) = GT
    compare (Prop.Disjunction _ _) (Prop.Negation _) = GT
    compare (Prop.Disjunction _ _) (Prop.Conjunction _ _) = GT
    compare (Prop.Disjunction p q) (Prop.Disjunction p' q') = compare (p,q) (p',q')
    compare (Prop.Disjunction _ _) (Prop.Conditional _ _) = LT
  
    compare (Prop.Conditional _ _) (Prop.Variable _) = GT
    compare (Prop.Conditional _ _) (Prop.Negation _) = GT
    compare (Prop.Conditional _ _) (Prop.Conjunction _ _) = GT
    compare (Prop.Conditional _ _) (Prop.Disjunction _ _) = GT
    compare (Prop.Conditional p q) (Prop.Conditional p' q') = compare (p,q) (p',q')
  
  ------------------------------------ Simple Baselines -------------------------
  
  subset :: (Eq a) => [a] -> [a] -> Bool
  subset a b = all (`elem` b) a
  
  prop_length :: Prop.Expr -> Int
  prop_length (Prop.Variable _) = 1
  prop_length (Prop.Negation p) = prop_length p + 1
  prop_length (Prop.Conjunction p q) = prop_length p + prop_length q + 1
  prop_length (Prop.Disjunction p q) = prop_length p + prop_length q + 1
  prop_length (Prop.Conditional p q) = prop_length p + prop_length q + 1
  prop_length (Prop.Biconditional p q) = prop_length p + prop_length q + 1
  
  heuristic_1 :: Prop.Expr -> Prop.Expr -> Int
  heuristic_1 a b = case prop_length a >= prop_length b of
    True -> 1
    False -> 0
  
  heuristic_2 :: Prop.Expr -> Prop.Expr -> Int
  heuristic_2 a b = case Prop.variables b `subset` Prop.variables a of
    True -> 1
    False -> 0
  
  heuristic_3 :: Prop.Expr -> Prop.Expr -> Int
  heuristic_3 a b = case ls_b `subset` ls_a of
    True -> 1
    False -> 0
    where
    ls_a = literals (fixed_point NormalForms.toNNF a)
    ls_b = literals (fixed_point NormalForms.toNNF b)
  
  literals :: Prop.Expr -> [(Prop.Var, Bool)]
  literals (Prop.Variable p) = [(p, True)]
  literals (Prop.Negation (Prop.Variable p)) = [(p, False)]
  literals (Prop.Conjunction p q) = List.nub $ literals p ++ literals q
  literals (Prop.Disjunction p q) = List.nub $ literals p ++ literals q
  literals (Prop.Conditional p q) = List.nub $ literals p ++ literals q
  literals (Prop.Biconditional p q) = List.nub $ literals p ++ literals q
  
  ------------------------------------ Baselines Stats -------------------------
  
  get_stats :: String -> IO ()
  get_stats f = do
    ts <- parse_file f
    (stat1, stat2, stat3, stat4, stat5) <- foldM accumulate_stat (0, 0, 0, 0, 0) ts
    putStrLn $ "Heuristic 1: " ++ show stat1 ++ " / " ++ show (length ts)
    putStrLn $ "Heuristic 2: " ++ show stat2 ++ " / " ++ show (length ts)
    putStrLn $ "Heuristic 3: " ++ show stat3 ++ " / " ++ show (length ts)
    putStrLn $ "Proportion of tautologies: " ++ show stat4 ++ " / " ++ show (length ts)
    putStrLn $ "Proportion of satisfiable Bs: " ++ show stat5 ++ " / " ++ show (length ts)
  
  accumulate_stat :: (Int, Int, Int, Int, Int) -> Entailment -> IO (Int, Int, Int, Int, Int)
  accumulate_stat (s1, s2, s3, s4, s5) (_, b, entails, h1, h2, h3) = do
    let s1' = if h1 == entails then s1 + 1 else s1
    let s2' = if h2 == entails then s2 + 1 else s2
    let s3' = if h3 == entails then s3 + 1 else s3
    not_is_tautology_b <- sat (Prop.Negation b)
    let s4' = if not_is_tautology_b then s4 else s4 + 1
    satisfiable_b <- sat b
    let s5' = if satisfiable_b then s5 + 1 else s5
    return (s1', s2', s3', s4', s5')
  
  ------------------------------------ Bias Stats -------------------------------
  
  bias_stats :: String -> IO ()
  bias_stats f = do
    ts <- parse_file f
    print_stats ts prop_length "length"
    print_stats ts num_vars "num vars"
    print_stats2 ts num_new_vars "new vars"
    print_stats ts num_negs "num negs"
    print_stats ts num_conjs "num conjs"
    print_stats ts num_disjs "num disjs"
    print_stats ts num_imps "num imps"
    print_stats ts count_sats "num satisfying TVAs"
    print_stats ts (num_negs_at_level 0) "negs at level 0"
    print_stats ts (num_negs_at_level 1) "negs at level 1"
    print_stats ts (num_negs_at_level 2) "negs at level 2"
    print_stats ts (num_conjs_at_level 0) "conjs at level 0"
    print_stats ts (num_conjs_at_level 1) "conjs at level 1"
    print_stats ts (num_conjs_at_level 2) "conjs at level 2"
    print_stats ts (num_disjs_at_level 0) "disjs at level 0"
    print_stats ts (num_disjs_at_level 1) "disjs at level 1"
    print_stats ts (num_disjs_at_level 2) "disjs at level 2"
    
  print_stats :: [Entailment] -> (Prop.Expr -> Int) -> String -> IO ()
  print_stats ts f s = do
    let (dist1, dist0) = List.foldl' (accumulate_bias_a f) (Map.empty, Map.empty) ts
    let (m1, m2) = (mean dist1, mean dist0)
    putStrLn $ "Mean " ++ s ++ " for A: " ++ show m1
    putStrLn $ "Mean " ++ s ++ " for A*: " ++ show m2
    putStrLn $ "Chi-squared for " ++ s ++ " is: " ++ show (chi_squared dist1 dist0)
    let (dist1, dist0) = List.foldl' (accumulate_bias_b f) (Map.empty, Map.empty) ts
    let (m1, m2) = (mean dist1, mean dist0)
    putStrLn $ "Mean " ++ s ++ " for B: " ++ show m1
    putStrLn $ "Mean " ++ s ++ " for B*: " ++ show m2
    putStrLn $ "Chi-squared for " ++ s ++ " is: " ++ show (chi_squared dist1 dist0)
  
  print_stats2 :: [Entailment] -> (Prop.Expr -> Prop.Expr -> Int) -> String -> IO ()
  print_stats2 ts f s = do
    let (dist1, dist0) = List.foldl' (accumulate_bias2 f) (Map.empty, Map.empty) ts
    let (m1, m2) = (mean dist1, mean dist0)
    putStrLn $ "Mean " ++ s ++ " for B: " ++ show m1
    putStrLn $ "Mean " ++ s ++ " for B*: " ++ show m2
    putStrLn $ "Chi-squared for " ++ s ++ " is: " ++ show (chi_squared dist1 dist0)
  
  chi_squared :: Map.Map Int Int -> Map.Map Int Int -> (Float, Int)
  chi_squared m1 m2 = (sum (map f all_keys), length all_keys) where
    all_keys = List.sort (List.nub (Map.keys m1 ++ Map.keys m2))
    f i = case (Map.lookup i m1, Map.lookup i m2) of
      (Nothing, Nothing) -> 0.0
      (Nothing, Just s_i) -> fromIntegral s_i 
      (Just r_i, Nothing) -> fromIntegral r_i
      (Just r_i, Just s_i) -> g (fromIntegral r_i) (fromIntegral s_i)
    g x y = (x - y)**2 / (x+y)
  
  mean :: Map.Map Int Int -> Float  
  mean m = total / count where
    m_l = Map.toList m
    count = sum (map (fromIntegral . snd) m_l) :: Float
    total = sum (map (\(a,b) -> (fromIntegral (a * b))) m_l) :: Float
  
  accumulate_bias_a :: (Prop.Expr -> Int) -> (Map.Map Int Int, Map.Map Int Int) -> Entailment -> (Map.Map Int Int, Map.Map Int Int)
  accumulate_bias_a f (pos_dist, neg_dist) (a, _, 1, _, _, _) = (update_count f pos_dist a, neg_dist)
  accumulate_bias_a f (pos_dist, neg_dist) (a, _, 0, _, _, _) = (pos_dist, update_count f neg_dist a)
  
  accumulate_bias_b :: (Prop.Expr -> Int) -> (Map.Map Int Int, Map.Map Int Int) -> Entailment -> (Map.Map Int Int, Map.Map Int Int)
  accumulate_bias_b f (pos_dist, neg_dist) (_, b, 1, _, _, _) = (update_count f pos_dist b, neg_dist)
  accumulate_bias_b f (pos_dist, neg_dist) (_, b, 0, _, _, _) = (pos_dist, update_count f neg_dist b)
  
  accumulate_bias2 :: (Prop.Expr -> Prop.Expr -> Int) -> (Map.Map Int Int, Map.Map Int Int) -> Entailment -> (Map.Map Int Int, Map.Map Int Int)
  accumulate_bias2 f (pos_dist, neg_dist) (a, b, 1, _, _, _) = (update_count2 f pos_dist a b, neg_dist)
  accumulate_bias2 f (pos_dist, neg_dist) (a, b, 0, _, _, _) = (pos_dist, update_count2 f neg_dist a b)
  
  update_count :: (Prop.Expr -> Int) -> Map.Map Int Int -> Prop.Expr -> Map.Map Int Int
  update_count f m b = case Map.lookup l m of
    Nothing -> Map.insert l 1 m
    Just n -> Map.insert l (n+1) m
    where l =  f b
  
  update_count2 :: (Prop.Expr -> Prop.Expr -> Int) -> Map.Map Int Int -> Prop.Expr -> Prop.Expr -> Map.Map Int Int
  update_count2 f m a b = case Map.lookup l m of
    Nothing -> Map.insert l 1 m
    Just n -> Map.insert l (n+1) m
    where l =  f a b
  
  num_vars :: Prop.Expr -> Int
  num_vars = length . extract_vars
  
  num_new_vars :: Prop.Expr -> Prop.Expr -> Int
  num_new_vars a b = length (extract_vars b List.\\ extract_vars a)
  
  num_negs :: Prop.Expr -> Int
  num_negs (Prop.Variable _) = 0
  num_negs (Prop.Negation p) = num_negs p + 1
  num_negs (Prop.Conjunction p q) = num_negs p + num_negs q
  num_negs (Prop.Disjunction p q) = num_negs p + num_negs q
  num_negs (Prop.Conditional p q) = num_negs p + num_negs q
  num_negs (Prop.Biconditional p q) = num_negs p + num_negs q
  
  num_conjs :: Prop.Expr -> Int
  num_conjs (Prop.Variable _) = 0
  num_conjs (Prop.Negation p) = num_conjs p
  num_conjs (Prop.Conjunction p q) = num_conjs p + num_conjs q + 1
  num_conjs (Prop.Disjunction p q) = num_conjs p + num_conjs q
  num_conjs (Prop.Conditional p q) = num_conjs p + num_conjs q
  num_conjs (Prop.Biconditional p q) = num_conjs p + num_conjs q
  
  num_disjs :: Prop.Expr -> Int
  num_disjs (Prop.Variable _) = 0
  num_disjs (Prop.Negation p) = num_disjs p
  num_disjs (Prop.Conjunction p q) = num_disjs p + num_disjs q
  num_disjs (Prop.Disjunction p q) = num_disjs p + num_disjs q + 1
  num_disjs (Prop.Conditional p q) = num_disjs p + num_disjs q
  num_disjs (Prop.Biconditional p q) = num_disjs p + num_disjs q
  
  num_imps :: Prop.Expr -> Int
  num_imps (Prop.Variable _) = 0
  num_imps (Prop.Negation p) = num_imps p
  num_imps (Prop.Conjunction p q) = num_imps p + num_imps q
  num_imps (Prop.Disjunction p q) = num_imps p + num_imps q
  num_imps (Prop.Conditional p q) = num_imps p + num_imps q + 1
  num_imps (Prop.Biconditional p q) = num_imps p + num_imps q
  
  num_negs_at_level :: Int -> Prop.Expr -> Int
  num_negs_at_level level p = num_negs_at_level2 p level 0
  
  num_negs_at_level2 :: Prop.Expr -> Int -> Int -> Int
  num_negs_at_level2 (Prop.Variable _) _ _ = 0
  num_negs_at_level2 (Prop.Negation _) level current | level == current = 1
  num_negs_at_level2 (Prop.Negation _) level current | level < current = 0
  num_negs_at_level2 (Prop.Negation p) level current | level > current = num_negs_at_level2 p level (current+1)
  num_negs_at_level2 (Prop.Conjunction _ _) level current | level <= current = 0
  num_negs_at_level2 (Prop.Conjunction p q) level current | level > current = 
    num_negs_at_level2 p level (current+1) + num_negs_at_level2 q level (current+1)
  num_negs_at_level2 (Prop.Disjunction _ _) level current | level <= current = 0
  num_negs_at_level2 (Prop.Disjunction p q) level current | level > current = 
    num_negs_at_level2 p level (current+1) + num_negs_at_level2 q level (current+1)
  num_negs_at_level2 (Prop.Conditional _ _) level current | level <= current = 0
  num_negs_at_level2 (Prop.Conditional p q) level current | level > current = 
    num_negs_at_level2 p level (current+1) + num_negs_at_level2 q level (current+1)
  
  num_conjs_at_level :: Int -> Prop.Expr -> Int
  num_conjs_at_level level p = num_conjs_at_level2 p level 0
  
  num_conjs_at_level2 :: Prop.Expr -> Int -> Int -> Int
  num_conjs_at_level2 (Prop.Variable _) _ _ = 0
  num_conjs_at_level2 (Prop.Negation _) level current | level <= current = 0
  num_conjs_at_level2 (Prop.Negation p) level current | level > current = num_conjs_at_level2 p level (current+1)
  num_conjs_at_level2 (Prop.Conjunction _ _) level current | level == current = 1
  num_conjs_at_level2 (Prop.Conjunction _ _) level current | level < current = 0
  num_conjs_at_level2 (Prop.Conjunction p q) level current | level > current = 
    num_conjs_at_level2 p level (current+1) + num_conjs_at_level2 q level (current+1)
  num_conjs_at_level2 (Prop.Disjunction _ _) level current | level <= current = 0
  num_conjs_at_level2 (Prop.Disjunction p q) level current | level > current = 
    num_conjs_at_level2 p level (current+1) + num_conjs_at_level2 q level (current+1)
  num_conjs_at_level2 (Prop.Conditional _ _) level current | level <= current = 0
  num_conjs_at_level2 (Prop.Conditional p q) level current | level > current = 
    num_conjs_at_level2 p level (current+1) + num_conjs_at_level2 q level (current+1)
  
  num_disjs_at_level :: Int -> Prop.Expr -> Int
  num_disjs_at_level level p = num_disjs_at_level2 p level 0
  
  num_disjs_at_level2 :: Prop.Expr -> Int -> Int -> Int
  num_disjs_at_level2 (Prop.Variable _) _ _ = 0
  num_disjs_at_level2 (Prop.Negation _) level current | level <= current = 0
  num_disjs_at_level2 (Prop.Negation p) level current | level > current = num_disjs_at_level2 p level (current+1)
  num_disjs_at_level2 (Prop.Disjunction _ _) level current | level == current = 1
  num_disjs_at_level2 (Prop.Disjunction _ _) level current | level < current = 0
  num_disjs_at_level2 (Prop.Disjunction p q) level current | level > current = 
    num_disjs_at_level2 p level (current+1) + num_disjs_at_level2 q level (current+1)
  num_disjs_at_level2 (Prop.Conjunction _ _) level current | level <= current = 0
  num_disjs_at_level2 (Prop.Conjunction p q) level current | level > current = 
    num_disjs_at_level2 p level (current+1) + num_disjs_at_level2 q level (current+1)
  num_disjs_at_level2 (Prop.Conditional _ _) level current | level <= current = 0
  num_disjs_at_level2 (Prop.Conditional p q) level current | level > current = 
    num_disjs_at_level2 p level (current+1) + num_disjs_at_level2 q level (current+1)
  
  ------------------------------------ Counting Sats ----------------------------
  
  type TVA = [(Char, Bool)]
  
  getTVAs :: [Char] -> [TVA]
  getTVAs vars = 
      let bs = getBools (length vars)
          f b = zip vars b
      in  map f bs
  
  getBools :: Int -> [[Bool]]
  getBools 0 = [[]]
  getBools n = 
      let p = getBools (n-1)
          lhs = map (False :) p
          rhs = map (True :) p
      in  lhs ++ rhs
  
  getTable :: Prop.Expr -> [TVA]
  getTable p =
      let vs = extract_vars p
      in  getTVAs vs
  
  satProp :: TVA -> Prop.Expr -> Bool
  satProp tva (Prop.Variable (Prop.Var p)) = case lookup p tva of
      Nothing -> error "Unexpected"
      Just b -> b
  satProp tva (Prop.Negation p) = not (satProp tva p)
  satProp tva (Prop.Conjunction p q) = (satProp tva p) && (satProp tva q)
  satProp tva (Prop.Disjunction p q) = (satProp tva p) || (satProp tva q)
  satProp tva (Prop.Conditional p q) = (not $ satProp tva p) || (satProp tva q)
  
  count_sats :: Prop.Expr -> Int
  count_sats p = length (filter (\tva -> satProp tva p) (getTable p))
  
  ------------------------------------ four-tuples ------------------------------
  
  create_four_tuples :: Int -> Int -> Int -> IO [FourTuple]
  create_four_tuples n v k = do
    pairs <- make_pairs n v k
    extract_four_tuples pairs []
  
  make_pairs :: Int -> Int -> Int -> IO [(Prop.Expr, Prop.Expr)]
  make_pairs n v k = do
    putStrLn $ "Creating " ++ show n ++ " pairs..."
    r <- newStdGen
    let (num_vars, r2) = randomR (1,v) r
    let (vs, _, r2) = take_n alphabet num_vars r
    make_pairs2 n k vs []
  
  make_pairs2 :: Int -> Int -> [Char] -> [(Prop.Expr, Prop.Expr)] -> IO [(Prop.Expr, Prop.Expr)]
  make_pairs2 n _ _ xs | length xs >= n = return xs
  make_pairs2 n k vs acc = do
    r <- newStdGen
    let (c1, r2) = randomR (k `div` 3, k) r
    let (c2, r3) = randomR (k `div` 3, k) r2
    (e1, e2) <- make_pair c1 c2 vs
    is_entails <- e1 `entails` e2
    case is_entails of
      True -> make_pairs2 n k vs ((e1, e2) : acc)
      False -> make_pairs2 n k vs acc
  
  extract_four_tuples :: [(Prop.Expr, Prop.Expr)] -> [FourTuple] -> IO [FourTuple]
  extract_four_tuples [] acc = return acc
  extract_four_tuples ((a,b) : ps) acc = do
    x <- find_other_pair (a,b) ps []
    case x of
      Just ((a',b'), ps') -> extract_four_tuples ps' ((a, b, a', b'):acc)
      Nothing -> extract_four_tuples ps acc
  
  find_other_pair :: (Prop.Expr, Prop.Expr) -> [(Prop.Expr, Prop.Expr)] -> [(Prop.Expr, Prop.Expr)] -> IO (Maybe ((Prop.Expr, Prop.Expr), [(Prop.Expr, Prop.Expr)]))
  find_other_pair _ [] acc = return Nothing
  find_other_pair (a,b) ((a', b') : ps) acc = do
    x <- valid_four_tuple (a,b) (a', b')
    case x of
      True -> return $ Just ((a', b'), acc ++ ps)
      False -> find_other_pair (a, b) ps ((a', b') : acc)
  
  valid_four_tuple :: (Prop.Expr, Prop.Expr) -> (Prop.Expr, Prop.Expr) -> IO Bool
  valid_four_tuple (a, b) (a', b') = do
    let p1 = Prop.Conjunction a (Prop.Negation b')
    let p2 = Prop.Conjunction a' (Prop.Negation b)
    s1 <- sat p1
    s2 <- sat p2
    return (s1 && s2)
  
  output_file :: String  
  output_file = "four_tuples.txt"
  
  write_four_tuples :: [FourTuple] -> IO ()
  write_four_tuples fts = forM_ fts $ \(a, b, a2, b2) -> do
    let a' = showF a
    let b' = showF b
    let a2' = showF a2
    let b2' = showF b2
    --putStrLn "Four-tuple found"
    --putStrLn a'
    --putStrLn b'
    --putStrLn a2'
    --putStrLn b2'
    --putStrLn "---"  
    let file = output_file
    let baselines_ab = show (heuristic_1 a b) ++ "," ++ show (heuristic_2 a b) ++ "," ++ show (heuristic_3 a b)
    appendFile file $ a' ++ "," ++ b' ++ ",1," ++ baselines_ab ++ "\n" 
    let baselines_a2b2 = show (heuristic_1 a2 b2) ++ "," ++ show (heuristic_2 a2 b2) ++ "," ++ show (heuristic_3 a2 b2)
    appendFile file $ a2' ++ "," ++ b2' ++ ",1," ++ baselines_a2b2 ++ "\n"
    let baselines_ab2 = show (heuristic_1 a b2) ++ "," ++ show (heuristic_2 a b2) ++ "," ++ show (heuristic_3 a b2)
    appendFile file $ a' ++ "," ++ b2' ++ ",0," ++ baselines_ab2 ++ "\n" 
    let baselines_a2b = show (heuristic_1 a2 b) ++ "," ++ show (heuristic_2 a2 b) ++ "," ++ show (heuristic_3 a2 b)
    appendFile file $ a2' ++ "," ++ b' ++ ",0," ++ baselines_a2b ++ "\n" 
  
  make_four_tuples :: Int -> Int -> Int -> IO ()
  make_four_tuples n v k = do
    fts <- create_four_tuples n v k
    putStrLn $ "There were " ++ show (length fts) ++ " four-tuples found."
    putStrLn $ "They produce " ++ show (4 * length fts) ++ " lines."
    write_four_tuples fts
    putStrLn $ "Output written (appended) to " ++ output_file
  
  make_pair :: Int -> Int -> [Char] -> IO (Prop.Expr, Prop.Expr)
  make_pair c1 c2 vs = do
    r2 <- newStdGen
    let (p1, r3) = gen_formula r2 c1 vs
    let (p2, _) = gen_formula r3 c2 vs
    return (p1, p2)
  
  average_length = average_stats length_triple
  average_num_vars = average_stats num_vars_triple
  average_num_ops = average_stats num_ops_triple
  average_power_two_num_vars = average_stats power_two_num_vars_triple
  
  average_stats :: (Int -> Entailment -> Int) -> String -> IO Float
  average_stats g f = do
    ts <- parse_file f
    let x = List.foldl' g 0 ts
    let x' = fromIntegral x :: Float
    let l = fromIntegral (length ts) :: Float
    let y = x' / l :: Float
    return y
  
  length_triple :: Int -> Entailment -> Int
  length_triple n (a, b, _, _, _, _) = n + prop_length a + prop_length b
  
  num_vars_triple :: Int -> Entailment -> Int
  num_vars_triple n (a, b, _, _, _, _) = n + num_vars (Prop.Conditional a b)
  
  power_two_num_vars_triple :: Int -> Entailment -> Int
  power_two_num_vars_triple n (a, b, _, _, _, _) = n + 2 ^ (num_vars (Prop.Conditional a b))
  
  num_ops_triple :: Int -> Entailment -> Int
  num_ops_triple n (a, b, _, _, _, _) = n + num_ops a + num_ops b
  
  num_ops :: Prop.Expr -> Int
  num_ops (Prop.Variable _) = 0
  num_ops (Prop.Negation p) = num_ops p + 1
  num_ops (Prop.Conjunction p q) = num_ops p + num_ops q + 1
  num_ops (Prop.Disjunction p q) = num_ops p + num_ops q + 1
  num_ops (Prop.Conditional p q) = num_ops p + num_ops q + 1
  num_ops (Prop.Biconditional p q) = num_ops p + num_ops q + 1
  
  