SOP.hs

module SOP where

import Control.Applicative
import Data.Function
import Data.List
import Data.Maybe
import Data.Either
import Debug.Trace

data Expr
  = Lit Integer
  | Var String
  | Add Expr Expr
  | Sub Expr Expr
  | Mul Expr Expr
  | Exp Expr Expr
  deriving (Show,Eq)

data Symbol = I Integer
            | S String
            | E { base :: SOP, ex :: Product }
  deriving (Eq,Ord,Show)

type Product = [Symbol]
type SOP     = [Product]

renderSOP     = concat . intersperse (" + ") . map renderProduct
renderProduct = concat . intersperse ("*") . map renderSymbol
renderSymbol (I i) = show i
renderSymbol (S s) = s
renderSymbol (E b e) = case (renderSimple b, renderSimple [e]) of
                         (bS,eS) -> bS ++ "^" ++ eS
renderSimple [[I i]] = show i
renderSimple [[S s]] = s
renderSimple sop     = "(" ++ renderSOP sop ++ ")"

simplify :: (a -> a -> Either a a) -> [a] -> [a]
simplify _ []      = []
simplify op (f:fs) = case partitionEithers $ map (`op` f) fs of
                       ([],_)              -> f : simplify op fs
                       (updated,untouched) -> simplify op (updated ++ untouched)

isSimple :: Symbol -> Bool
isSimple (I _)       = True
isSimple (S _)       = True
isSimple (E [[_]] _) = True
isSimple _           = False

reduceSymbol :: Symbol -> Symbol
reduceSymbol (E [[I 0]]     _      ) = I 0
reduceSymbol (E _           [(I 0)]) = I 1
reduceSymbol (E [[(I i)]]   [(I j)]) = I (i ^ j)
reduceSymbol (E [[(E k i)]] j      ) = E k (sort . map reduceSymbol  $ simplify mergeS (i ++ j))
reduceSymbol s                       = s

mergeS :: Symbol -> Symbol -> Either Symbol Symbol
mergeS (I i)            (I j)            = Left (I (i * j))
mergeS (I 1)            r                = Left r
mergeS l                (I 1)            = Left l
mergeS (I 0)            r                = Left (I 0)
mergeS l                (I 0)            = Left (I 0)
mergeS s                (E [[s']] [I i])
  | s == s'                              = Left (E [[s']] [I (i + 1)])
mergeS (E [[s']] [I i]) s        
  | s == s'                              = Left (E [[s']] [I (i + 1)]  )
mergeS l        r 
  | l == r && isSimple l                 = Left (E [[l]] [I 2])
mergeS l        _                        = Right l

mergeP :: Product -> Product -> Either Product Product
mergeP ((I i):is) ((I j):js)  
  | is == js = Left $ (I (i + j)) : is
mergeP is js
 | is == js  = Left $ (I 2) : is
 | otherwise = Right is

expandExp :: SOP -> SOP -> SOP
expandExp b         [[(I 1)]]   = b
expandExp b@([[_]]) [e@(_:_)]   = [[reduceSymbol (E b e)]]
expandExp b         e@[[(I i)]] = foldr1 mergeSOPMul (replicate (fromInteger i) b)
expandExp b         [e@[_]]     = [[reduceSymbol (E b e)]]
expandExp b         e           = foldr1 mergeSOPMul (map (expandExp b . (:[])) e)

toSOP :: Expr -> SOP
toSOP (Lit i)     = [[I i]]
toSOP (Var s)     = [[S s]]
toSOP (Add e1 e2) = mergeSOPAdd (toSOP e1) (toSOP e2)
toSOP (Sub e1 e2) = mergeSOPAdd (toSOP e1) (mergeSOPMul [[I (-1)]] (toSOP e2))
toSOP (Mul e1 e2) = mergeSOPMul (toSOP e1) (toSOP e2)
toSOP (Exp e1 e2) = expandExp   (toSOP e1) (toSOP e2)

simplifySOP :: SOP -> SOP
simplifySOP = sort . simplify mergeP . filter (/= [I 0]) . map (sort . map reduceSymbol  . simplify mergeS)

mergeSOPAdd :: SOP -> SOP -> SOP
mergeSOPAdd sop1 sop2 = simplifySOP $ sop1 ++ sop2

mergeSOPMul :: SOP -> SOP -> SOP
mergeSOPMul sop1 sop2 = simplifySOP $ concatMap (zipWith (++) sop1 . repeat) sop2

expr1 = Mul (Lit 4)   (Var "k")
expr2 = Mul (Var "x") (Lit 8)
expr3 = Add expr1 expr2
expr4 = Mul (Lit 1)   (Var "y")
expr5 = Lit 4
expr6 = Add expr4 expr5
expr7 = Mul expr3 expr6
expr8 = Mul (Lit 7) (Var "x")
expr9 = Mul expr7 expr8

exprK = Mul (Add (Var "y") (Lit 2)) (Exp (Add (Var "y") (Lit 2)) (Add (Var "x") (Lit 1)))
exprL = Add (Exp (Add (Var "y") (Lit 2)) (Var "x")) (Exp (Add (Var "y") (Lit 2)) (Var "x")) 

exprT1 = Add (Lit 3) (Lit 4)
exprT2 = Mul (Sub (Lit 224) (Mul (Lit 56) (Var "y"))) (Exp (Var "x") (Lit 2))
exprT3 = Mul (Sub (Mul (Lit 56) (Var "y")) (Lit 224)) (Exp (Var "x") (Lit 2))
exprT4 = Sub (Sub (Var "x") (Var "y")) (Var "z")
exprT5 = Sub (Var "z") (Sub (Var "x") (Var "y"))

exprT6 = Mul (Exp (Var "x") (Add (Var "y") (Lit 1))) (Exp (Var "x") (Add (Var "y") (Lit 4)))
exprT7 = Exp (Add (Var "x") (Lit 2)) (Var "x")

exprT8 = Add (Exp (Lit 2) (Var "x")) (Exp (Lit 2) (Var "x"))
exprT9 = Exp (Lit 2) (Add (Lit 1) (Var "x"))

exprT10 = Exp (Exp (Var "x") (Var "y")) (Var "z")

equalSOP :: Expr -> Expr -> Bool
equalSOP = (==) `on` toSOP