module Main where

import Random
import System

{--
 ** Functions for generating sequences of pseudo random numbers
 --}

randNumSeq :: Int -> (Int,Int) -> [Int]
randNumSeq n (min,max) = randomRs (min,max) (mkStdGen n)

rollDices = randNumSeq 0 (1,6)
flipCoins = randNumSeq 0 (0,1)


{-- 
 ** Some auxilary functions
 --}

evens (x0:x1:xs) = x0 : (evens xs)
odds (x0:x1:xs) = x1 : (odds xs)

compactify [] = []
compactify ((k,x):xs) = (s,x) : compactify ys
    where s = foldr (\(l,_) a->a+l) k (filter (\(_,y)->x==y) xs)
	  ys = filter (\(_,y)->x/=y) xs

nth 0 (x:xs) = x
nth n (x:xs) = nth (n-1) xs
nth _ _ = error "nth: cannot take the nth element of the list"

string2int :: String -> Int
string2int = read

{-- 
 ** Data types for representing formulas
 --}

data RelOp = EQUAL | NEQUAL | LESS | LESSEQUAL | GREATER | GREATEREQUAL
	     deriving Eq

instance Show RelOp where
    show EQUAL = "="
    show NEQUAL = "/="
    show LESS = "<"
    show LESSEQUAL = "<="
    show GREATER = ">"
    show GREATEREQUAL = ">="

toRelOp 0 = EQUAL 
toRelOp 1 = NEQUAL
toRelOp 2 = LESS
toRelOp 3 = LESSEQUAL
toRelOp 4 = GREATER
toRelOp 5 = GREATEREQUAL

data Variable = Var String
		deriving Eq 

instance Show Variable where
    show (Var x) = x 

data Formula a = TT | FF 
	       | Atom ([(Int,a)],RelOp,Int) 
	       | Not (Formula a)
	       | And (Formula a) (Formula a)
	       | Or (Formula a) (Formula a)
	       | Imp (Formula a) (Formula a)
	       | Iff (Formula a) (Formula a)
	       | Exists a (Formula a)
	       | Forall a (Formula a)
		 deriving Eq

{-- 
 ** Convert the formula into a string such that LASH can parse it.
 ** Note that the LASH parser is limited/poor/restrictive. So, the
 **  conversion is rather cumbersome.
 --}

instance (Show a, Eq a) => Show (Formula a) where
    show TT = "0*x=0"
    show FF = "0*x=1"
    show (Atom (xs,NEQUAL,c)) = show (Not (Atom (xs,EQUAL,c)))
    show (Atom ([],op,c)) = "0" ++ (show op) ++ (show c)
    show (Atom (xs,op,c)) = (showTerm (compactify xs)) ++ " " ++ (show op) ++ " " ++ (show c)
	where showTerm [(k,x)] 
		  | k>0 = (show k) ++ "*" ++ (show x)
		  | k<0 = "(" ++ (show k) ++ "*" ++ (show x) ++ ")"
		  | otherwise = "0"
	      showTerm ((k,x):xs) 
		  | k>0 = (show k) ++ "*" ++ (show x) ++ " + " ++ (showTerm xs) 
		  | k<0 = "(" ++ (show k) ++ "*" ++ (show x) ++ ")" ++ " + " ++ (showTerm xs) 
		  | otherwise = showTerm xs
    show (Not f) = "NOT " ++ "(" ++ (show f) ++ ")"
    show (And f g) = "(" ++ (show f) ++ " AND " ++ (show g) ++ ")"
    show (Or f g) = "(" ++ (show f) ++ " OR " ++ (show g) ++ ")"
    show (Imp f g) = show (Or (Not f) g)
    show (Iff f g) = show (And (Imp f g) (Imp g f))
    show (Exists x f) = "(EXISTS (" ++ (show x) ++ " : " ++ (show f) ++ "))"
    show (Forall x f) = "(FORALL (" ++ (show x) ++ " : " ++ (show f) ++ "))"


{--
 ** Simple rewrites 
 --}

elimNot (Not TT) = FF
elimNot (Not FF) = TT

elimNot (Not (Atom (t,EQUAL,c))) = Atom (t,NEQUAL,c)
elimNot (Not (Atom (t,NEQUAL,c))) = Atom (t,EQUAL,c)
elimNot (Not (Atom (t,LESS,c))) = Atom (t,GREATEREQUAL,c)
elimNot (Not (Atom (t,LESSEQUAL,c))) = Atom (t,GREATER,c)
elimNot (Not (Atom (t,GREATER,c))) = Atom (t,LESSEQUAL,c)
elimNot (Not (Atom (t,GREATEREQUAL,c))) = Atom (t,LESS,c)

elimNot (Not (Not f)) = elimNot f
elimNot (Not (And f g)) = Or (elimNot (Not f)) (elimNot (Not g)) 
elimNot (Not (Or f g)) = And (elimNot (Not f)) (elimNot (Not g)) 
elimNot (Not (Imp f g)) = elimNot (Not (Or (Not f) g))
elimNot (Not (Iff f g)) = elimNot (Not (And (Imp f g) (Imp g f)))
elimNot (Not (Exists x f)) = Forall x (elimNot (Not f))
elimNot (Not (Forall x f)) = Exists x (elimNot (Not f))

elimNot (And f g) = And (elimNot f) (elimNot g) 
elimNot (Or f g) = Or (elimNot f) (elimNot g) 
elimNot (Imp f g) = elimNot (Or (Not f) g)
elimNot (Iff f g) = elimNot (And (Imp f g) (Imp g f))
elimNot (Exists x f) = Exists x (elimNot f)
elimNot (Forall x f) = Forall x (elimNot f)

elimNot f = f

simplify (Atom (t,op,c))
    | t'==[] && (toRel op 0 c) = TT
    | t'==[] = FF
    | otherwise = Atom (t',op,c)
    where t' = [ (k,x) | (k,x)<-compactify t, k/=0 ] 
	  toRel EQUAL = (==)
	  toRel NEQUAL = (/=)
	  toRel LESS  = (<)
	  toRel LESSEQUAL = (<=)
	  toRel GREATER = (>)
	  toRel GREATEREQUAL = (>=)

simplify (Not f) 
    | f'==TT = FF
    | f'==FF = TT
    | otherwise = Not f'
    where f' = simplify f
simplify (And f g)
    | f'==FF || g'==FF = FF
    | f'==TT = g'
    | g'==TT = f'
    | otherwise = And f' g'
    where f' = simplify f
	  g' = simplify g
simplify (Or f g)
    | f'==TT || g'==TT = TT
    | f'==FF = g'
    | g'==FF = f'
    | otherwise = Or f' g'
    where f' = simplify f
	  g' = simplify g

simplify (Imp f g) = simplify (Or (Not f) g)
simplify (Iff f g) = simplify (And (Imp f g) (Imp g f))

simplify (Exists x f)
    | f'==TT = TT
    | f'==FF = FF
    | otherwise = Exists x f'
    where f' = simplify f
simplify (Forall x f)
    | f'==TT = TT
    | f'==FF = FF
    | otherwise = Forall x f'
    where f' = simplify f


{--
 ** Functions for some statistics
 --}

numAtoms (Atom _) = 1
numAtoms (Not f) = numAtoms f
numAtoms (And f g) = (numAtoms f) + (numAtoms g)
numAtoms (Or f g) = (numAtoms f) + (numAtoms g)
numAtoms (Imp f g) = (numAtoms f) + (numAtoms g)
numAtoms (Iff f g) = (numAtoms f) + (numAtoms g)
numAtoms (Exists _ f) = (numAtoms f)
numAtoms (Forall _ f) = (numAtoms f)

-- number of quantifiers
qn (Atom _) = 0
qn (Not f) = qn f
qn (And f g) = (qn f) + (qn g)
qn (Or f g) = (qn f) + (qn g)
qn (Imp f g) = (qn f) + (qn g)
qn (Iff f g) = (qn f) + (qn g)
qn (Exists _ f) = 1 + (qn f)
qn (Forall _ f) = 1 + (qn f)

-- quantifier rank
qr (Atom _) = 0
qr (Not f) = qr f
qr (And f g) = max (qr f) (qr g)
qr (Or f g) = max (qr f) (qr g)
qr (Imp f g) = max (qr f) (qr g)
qr (Iff f g) = qr (And (Imp f g) (Imp g f))
qr (Exists _ f) = 1 + (qr f)
qr (Forall _ f) = 1 + (qr f)

-- quantifier alternations
qa f = max (aux 0 f) (aux 1 f)
    where aux _ (Atom _) = 0
	  aux q (Not f) = aux (1-q) f
	  aux q (And f g) = max (aux q f) (aux q g)
	  aux q (Or f g) = max (aux q f) (aux q g)
	  aux q (Imp f g) = aux q (Or (Not f) g)
	  aux q (Iff f g) = aux q (And (Imp f g) (Imp g f))
	  aux 0 (Exists _ f) = aux 0 f
	  aux 1 (Exists _ f) = 1 + aux 0 f
	  aux 0 (Forall _ f) = 1 + aux 1 f
	  aux 1 (Forall _ f) = aux 1 f

statistics f = (numAtoms f, qn f, qr f, qa f)

{--
 ** Functions for randomly generating formulas
 --}

randFormula vars bounds md seed = generate (ints,ops,connects) vars (0,md) 
    where ints = zipWith (\x y-> if y==0 then 0 else x) (randNumSeq seed bounds) (randNumSeq seed (0,3))
	  ops = map toRelOp (randNumSeq seed (0,5))
	  connects = randNumSeq seed (0,4)

generate rands (xs,[]) depth = generateQF rands xs depth
generate (ints,ops,c:cs) (xs,y:ys) (d,md)
    | d<md && c==0 = Not (generate (ints,ops,cs) (xs,ys) (d+1,md))
    | d<md && c==1 = And f g
    | d<md && c==2 = Or f g
--    | d<md && c==5 = Imp f g
--    | d<md && c==6 = Iff f g
    | d<md && c==3 = Exists y h
    | d<md && c==4 = Forall y h
    | otherwise = Atom (zip (tail ints) xs, head ops, head ints)
    where f = generate (evens ints, evens ops, evens cs) (xs,y:ys) (d+1,md)
	  g = generate (odds ints, odds ops, odds cs) (xs,y:ys) (d+1,md)
	  h = generate (ints,ops,cs) (y:xs,ys) (d+1,md)

generateQF (ints,ops,c:cs) xs (d,md)
    | d<md && c==0 = Not (generateQF (ints,ops,cs) xs (d+1,md))
    | d<md && (c==1||c==3) = And f g
    | d<md && (c==2||c==4) = Or f g
--    | d<md && c==3 = Imp f g
--    | d<md && c==4 = Iff f g
    | otherwise = Atom (zip (tail ints) xs, head ops, head ints)
    where f = generateQF (evens ints, evens ops, evens cs) xs (d+1,md)
	  g = generateQF (odds ints, odds ops, odds cs) xs (d+1,md)


{--
 ** Main (compile with 'ghc --make -o randprsb randPrsb.hs')
 --}

withProgArgs :: (String -> [Int] -> IO ()) -> IO ()
withProgArgs main = do p <- getProgName
                       args <- getArgs
                       main p (map string2int args)

usage p = "USAGE: " ++ p ++ " #freeVars #boundVars #minInt #maxInt #maxDepth #seed [s]" 
--	  ++ "  where #freeVars+#boundVars>0 and #minInt<#maxInt and #maxDeptph>0"
	  

main = withProgArgs main'

main' p args 
    | (length args)==6 && okay = print f
    | (length args)==7 && okay = print (statistics f)
    | otherwise = print (usage p)
    where fv = nth 0 args
	  bv = nth 1 args
	  m = nth 2 args
	  n = nth 3 args
	  maxdepth = nth 4 args
	  seed = nth 5 args
	  f = simplify (elimNot (randFormula vars (m,n) maxdepth seed))
	  vars = ([ Var ("x"++(show i)) | i<-[1..fv] ], 
		  [ Var ("y"++(show i)) | i<-[1..bv] ])
	  okay = and [fv+bv>0, m<n, maxdepth>0]




{-- 
 ** Used for testing
 --}
freeVars = [Var "x1", Var "x2", Var "x3", Var "x4"]
boundVars = [Var "y1", Var "y2", Var "y3", Var "y4"]
vars = (freeVars,boundVars)

f1 = elimNot . (randFormula vars (-10,10) 5)
f2 = elimNot . (randFormula vars (-20,20) 5)
f3 = elimNot . (randFormula vars (-30,30) 5)

g1 = elimNot . (randFormula vars (-10,10) 7)
g2 = elimNot . (randFormula vars (-20,20) 7)
g3 = elimNot . (randFormula vars (-30,30) 7)

h1 = elimNot . (randFormula vars (-10,10) 9)
h2 = elimNot . (randFormula vars (-20,20) 9)
h3 = elimNot . (randFormula vars (-30,30) 9)

a1 = elimNot . (randFormula (freeVars,[]) (-10,10) 5)
a2 = elimNot . (randFormula (freeVars,[]) (-20,20) 5)
a3 = elimNot . (randFormula (freeVars,[]) (-30,30) 5)

b1 = elimNot . (randFormula (freeVars,[]) (-10,10) 7)
b2 = elimNot . (randFormula (freeVars,[]) (-20,20) 7)
b3 = elimNot . (randFormula (freeVars,[]) (-30,30) 7)