{-# LANGUAGE DeriveDataTypeable   #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE TypeSynonymInstances #-}

module Common.Arith where

import Control.Monad
import Data.Ratio
import Data.Maybe
import Data.Typeable
import Test.QuickCheck
import System.Random
import Common.GuardedRewriting

propGen :: Int -> Gen (Prop (Equation Expr))
propGen n
   | n == 0 = frequency 
        [ (5, liftM VarP (sized eqGen))
        , (1, return T)
        , (1, return F)
        ]
   | otherwise = oneof 
        [ propGen 0
        , liftM2 (:\/:) rec rec
        , liftM2 (:/\:) rec rec
        , liftM  Not    rec 
        ]
 where
   rec = propGen (n `div` 2)

eqGen :: Int -> Gen (Equation Expr)
eqGen n = liftM2 (:==:) (exprGen n) constGen

-- Expression with exactly one variable
exprGen :: Int -> Gen Expr
exprGen 0 = return (Varia "x")
exprGen n = oneof
   [ return (Varia "x")
   , liftM2 (:+:) rec constGen, liftM2 (:+:) constGen rec
   , liftM2 (:-:) rec constGen, liftM2 (:-:) constGen rec
   , liftM2 (:**:) rec constGen, liftM2 (:**:) constGen rec
   , liftM2 (:/:) rec constGen, liftM2 (:/:) constGen rec
   , liftM2 (\a n -> a :^: Const (fromInteger (succ $ abs n))) rec arbitrary
   ]
 where
   rec = exprGen (n `div` 2)
   
constGen :: Gen Expr
constGen = liftM (Const . fromInteger) arbitrary

isSolved :: Prop (Equation Expr) -> Bool
isSolved prop =
   case prop of
      p :\/: q -> isSolved p && isSolved q
      p :/\: q -> isSolved p && isSolved q
      Not p    -> isSolved p
      VarP p    -> isSolvedEq p 
      _        -> True

isSolvedEq :: Equation Expr -> Bool
isSolvedEq (Varia _ :==: b) = noVaria b
isSolvedEq (a       :==: b) = noVaria a && noVaria b

repeatM :: Monad m => m a -> m [a]
repeatM m = liftM2 (:) m (repeatM m)

formula :: [Prop (Equation Expr)]
formula = generate 100 (mkStdGen 280578) (repeatM (sized propGen))

formulas :: [Prop (Equation Expr)]
formulas = take 10000 formula

-----------------------------------

data Prop a = VarP a | T | F 
            | Prop a :\/: Prop a 
            | Prop a :/\: Prop a 
            | Not (Prop a)
   deriving (Show, Typeable, Eq)

data Equation a = a :==: a
   deriving (Show, Typeable, Eq)

infix  2 :==: 
infixl 6 :+:, :-: 
infixl 7 :**:, :/: 
infixr 8 :^:

data Expr = Const Rational
          | Varia String
          | Expr :+: Expr
          | Expr :-: Expr
          | Expr :**: Expr
          | Expr :/: Expr
          | Expr :^: Expr
   deriving (Eq, Show, Typeable)

instance Functor Prop where
   fmap f (VarP a)    = VarP (f a)
   fmap f T          = T
   fmap f F          = F
   fmap f (p :\/: q) = fmap f p :\/: fmap f q
   fmap f (p :/\: q) = fmap f p :/\: fmap f q
   fmap f (Not p)    = Not (fmap f p)
   
instance Functor Equation where
   fmap f (a :==: b) = f a :==: f b


isEven, isOdd :: Rational -> Bool
isEven r = denominator r == 1 && even (numerator r)
isOdd  r = denominator r == 1 && odd  (numerator r)

hasDivisionByZero :: Expr -> Bool
hasDivisionByZero expr =
   case expr of
      _ :/: Const 0 -> True
      a :+: b  -> hasDivisionByZero a || hasDivisionByZero b
      a :-: b  -> hasDivisionByZero a || hasDivisionByZero b
      a :**: b -> hasDivisionByZero a || hasDivisionByZero b
      a :/: b  -> hasDivisionByZero a || hasDivisionByZero b
      a :^: b  -> hasDivisionByZero a || hasDivisionByZero b
      _ -> False

noVaria, hasVaria :: Expr -> Bool
noVaria = not . hasVaria
hasVaria expr = 
   case expr of
      Const _  -> False
      Varia _  -> True
      a :+: b  -> hasVaria a || hasVaria b
      a :-: b  -> hasVaria a || hasVaria b
      a :**: b -> hasVaria a || hasVaria b
      a :/: b  -> hasVaria a || hasVaria b
      a :^: b  -> hasVaria a || hasVaria b

applyD :: (a -> Maybe a) -> a -> a
applyD f a = fromMaybe a (f a)

applyOne :: [a -> Maybe a] -> a -> Maybe a
applyOne fs a = case mapMaybe ($ a) fs of
                   hd:_ -> Just hd
                   _    -> Nothing
                   

applyBin :: (a -> Maybe a) -> (a -> a -> b) -> a -> a -> Maybe b
applyBin rec op a b = 
   case (rec a, rec b) of
      (Nothing, Nothing) -> Nothing
      (ma, mb) -> Just (fromMaybe a ma `op` fromMaybe b mb)

somewhereProp :: (Prop a -> Maybe (Prop a)) -> Prop a -> Maybe (Prop a)
somewhereProp f = rec
 where
   rec prop = 
      let make op a b = case applyBin rec op a b of
                           Just c  -> Just (applyD f c)
                           Nothing -> f prop
      in case prop of
            p :\/: q -> make (:\/:) p q
            p :/\: q -> make (:/\:) p q
            Not p    -> case rec p of
                           Just p2 -> Just $ applyD f p2
                           Nothing -> f prop
            _        -> f prop

somewhereExpr :: (Expr -> Maybe Expr) -> Expr -> Maybe Expr
somewhereExpr f = rec 
 where
   rec expr = 
      let make2 op a b = case applyBin rec op a b of
                           Just c  -> Just (applyD f c)
                           Nothing -> f expr
      in case expr of
            a :+: b  -> make2 (:+:) a b
            a :-: b  -> make2 (:-:) a b
            a :**: b -> make2 (:**:) a b
            a :/: b  -> make2 (:/:) a b
            a :^: b  -> make2 (:^:) a b
            _        -> f expr

liftToEq :: (a -> Maybe a) -> Equation a -> Maybe (Equation a)
liftToEq f (a :==: b) =
   case (f a, f b) of
      (Nothing, Nothing) -> Nothing
      (ma, mb) -> Just $ fromMaybe a ma :==: fromMaybe b mb

liftToProp :: (a -> Maybe a) -> Prop a -> Maybe (Prop a)
liftToProp f = rec 
 where
   rec prop =
      case prop of
         p :\/: q  -> applyBin rec (:\/:) p q
         p :/\: q  -> applyBin rec (:/\:) p q
         Not p     -> liftM Not (rec p)
         VarP a    -> liftM VarP (f a)
         _         -> Nothing

transformExpr :: (Expr -> Expr) -> Expr -> Expr
transformExpr f = rec 
 where
   rec expr = f $
      case expr of
         a :+: b  -> rec a :+: rec b
         a :-: b  -> rec a :-: rec b
         a :**: b -> rec a :**: rec b
         a :/: b  -> rec a :/: rec b
         a :^: b  -> rec a :^: rec b
         _        -> expr

simplify :: Expr -> Expr
simplify = transformExpr f
 where
   f (Const a :+: Const b)  = Const (a+b)
   f (Const a :-: Const b)  = Const (a-b)
   f (Const a :**: Const b) = Const (a*b)
   f (Const a :/: Const b)  = Const (a/b)
   f (Const a :^: Const b) | denominator b == 1 
      = Const (a^numerator b)
   f a = a

-- Also counts the number of rules that hvae been applied (successfully)
solve :: [Prop (Equation Expr) -> Maybe (Prop (Equation Expr))]
         -> Prop (Equation Expr) -> Prop (Equation Expr)
solve rules = rec where
   rec p = 
      case applyOne rules p of
         Just a  -> rec a
         Nothing -> p

-- Examples
ex :: Prop (Equation Expr)
ex = VarP p :\/: VarP q
 where
   x = Varia "x"
   p = (((x :-: Const 3) :^: Const 2) :/: Const 4) :==: Const 25
   q = (Const 1 :-: (Const 2 :/: (Const 1 :+: (x :**: Const 5)))) :==: Const 4 
ex2 = ((VarP (Const (5 % 1) :/: Varia "x" :==: Const (4 % 1)) :\/: VarP ((Const (6 % 1)
 :/: (Const ((-4) % 1) :+: Varia "x")) :^: Const (3 % 1) :==: Const ((-6) % 1)))
 :\/: VarP (Const ((-7) % 1) :**: (Const ((-5) % 1) :**: Varia "x") :==: Const (4 %
 1))) :\/: (VarP (Varia "x" :==: Const ((-7) % 1)) :\/: VarP (Const (0 % 1) :/: ((Varia "x" :-: Const (0 % 1)) :+: Const (2 % 1)) :==: Const ((-3) % 1)))
 
-- Generic representation
instance (Representable a) => Representable (Prop a) where
  type Rep (Prop a) =      Var a
                      :+: U
                      :+: U
                      :+: Rec (Prop a) :*: Rec (Prop a)
                      :+: Rec (Prop a) :*: Rec (Prop a)
                      :+: Rec (Prop a)
  
  from (VarP x)   = L (Var x)
  from T          = R (L U)
  from F          = R (R (L U))
  from (p :\/: q) = R (R (R (L (Rec p :*: Rec q))))
  from (p :/\: q) = R (R (R (R (L (Rec p :*: Rec q)))))
  from (Not p)    = R (R (R (R (R (Rec p)))))
  
  to (L (Var x))                           = VarP x
  to (R (L U))                             = T
  to (R (R (L U)))                         = F
  to (R (R (R (L (Rec p :*: Rec q)))))     = p :\/: q
  to (R (R (R (R (L (Rec p :*: Rec q)))))) = p :/\: q
  to (R (R (R (R (R (Rec p))))))           = p

instance (Representable a) => Representable (Equation a) where
  type Rep (Equation a) = Var a :*: Var a
  
  from (x :==: y) = Var x :*: Var y
  to (Var x :*: Var y) = x :==: y

instance Representable Expr where
  type Rep Expr =      Var Rational
                  :+: Var String
                  :+: Rec Expr :*: Rec Expr
                  :+: Rec Expr :*: Rec Expr
                  :+: Rec Expr :*: Rec Expr
                  :+: Rec Expr :*: Rec Expr
                  :+: Rec Expr :*: Rec Expr
  
  from (Const r) = L (Var r)
  from (Varia s) = R (L (Var s))
  from (e1 :+: e2)  = R (R (L (Rec e1 :*: Rec e2)))
  from (e1 :-: e2)  = R (R (R (L (Rec e1 :*: Rec e2))))
  from (e1 :**: e2) = R (R (R (R (L (Rec e1 :*: Rec e2)))))
  from (e1 :/: e2)  = R (R (R (R (R (L (Rec e1 :*: Rec e2))))))
  from (e1 :^: e2)  = R (R (R (R (R (R (Rec e1 :*: Rec e2))))))
  
  to (L (Var r)) = Const r
  to (R (L (Var s))) = Varia s
  to (R (R (L (Rec e1 :*: Rec e2)))) = e1 :+: e2
  to (R (R (R (L (Rec e1 :*: Rec e2))))) = e1 :-: e2
  to (R (R (R (R (L (Rec e1 :*: Rec e2)))))) = e1 :**: e2
  to (R (R (R (R (R (L (Rec e1 :*: Rec e2))))))) = e1 :/: e2
  to (R (R (R (R (R (R (Rec e1 :*: Rec e2))))))) = e1 :^: e2

 
instance (Rewritable a) => Rewritable (Prop a)
instance (Rewritable a) => Rewritable (Equation a)
instance Rewritable Expr
instance Rewritable String