import Text.ParserCombinators.Parsec
import Text.ParserCombinators.Parsec.Token
import Text.ParserCombinators.Parsec.Expr
import Text.ParserCombinators.Parsec.Language
import System.Environment
import Control.Monad

---------------------
-- Expression tree --
---------------------
data Expr = Number Integer
          | Variable Char
          | Term Op Expr Expr

data Op = Mul | Div | Sub | Add | Exp

------------------
-- Infix output --
------------------
instance Show Expr where
    show (Number x) = show x
    show (Variable a) = a : ""
    show (Term op x y) = '(' : (show op) ++ " " ++ show x ++ " " ++ show y ++ ")"

instance Show Op where
    show Mul = "*"
    show Div = "/"
    show Add = "+"
    show Sub = "-"
    show Exp = "^"

------------
-- Parser --
------------
lexer = makeTokenParser 
        (emptyDef { reservedOpNames = ["*","/","+","-"],
                    identStart = letter,
                    identLetter = (fail "")})

number  = do n <- (natural lexer)
             return $ Number n
variable = do v <- (identifier lexer)
              return $ Variable (head v)
term = (parens lexer) expr <|> number <|> variable

expr = buildExpressionParser [[op "*" Mul, op "/" Div],
                              [op "+" Add, op "-" Sub]]
                             term
    where op s o = Infix (do (reservedOp lexer s); return $ Term o) AssocLeft

parseExpr input = case (parse (do e <- expr
                                  eof
                                  return e)
                        "" input) of
                    Left _ -> error "Parse error."
                    Right x -> return x

----------------------------------------
-- Mini algebra system for reductions --
----------------------------------------
data Poly = Poly Char [Coof]
            deriving Show

data Coof = NumCoof Integer | PolyCoof Poly
            deriving Show
                   
-- Some utilities --
simplifyPolyCoof (Poly _ [NumCoof x]) = NumCoof x
simplifyPolyCoof p = PolyCoof p

complicatePolyCoof (NumCoof x) = Poly '_' [(NumCoof x)]
complicatePolyCoof (PolyCoof x) = x

zipExcess _ [] qs = qs
zipExcess _ ps [] = ps
zipExcess fn (p:ps) (q:qs) = (fn p q) : (zipExcess fn ps qs)

polyDegree (Poly _ cs) = toInteger $ length cs - 1
isPolySame (Poly v1 _) (Poly v2 _) = v1 == v2 && v1 /= '_'

-- Addition --
polyAdd p@(Poly _ [NumCoof _]) q = polyAddConst q p
polyAdd p q@(Poly _ [NumCoof _]) = polyAddConst p q
polyAdd p q
    | isPolySame p q = polyAddSame p q
polyAdd p@(Poly v1 _) q@(Poly v2 _)
    | v2 > v1 = polyAddConst p q
polyAdd p q = polyAddConst q p

polyAddConst p (Poly _ [NumCoof 0]) = p
polyAddConst (Poly _ [NumCoof x]) (Poly _ [NumCoof y]) = Poly '_' [NumCoof (x + y)]
polyAddConst (Poly v (c:cs)) k = Poly v ((simplifyPolyCoof (polyAdd (complicatePolyCoof c) k)):cs)

polyAddSame (Poly v cs1) (Poly _ cs2) = Poly v (map simplifyPolyCoof
                                                (zipExcess polyAdd 
                                                 (map complicatePolyCoof cs1)
                                                 (map complicatePolyCoof cs2)))

-- Multiplication --
polyMul p@(Poly _ [NumCoof _]) q = polyMulConst q p
polyMul p q@(Poly _ [NumCoof _]) = polyMulConst p q
polyMul p q
    | isPolySame p q = polyMulSame p q
polyMul p@(Poly v1 _) q@(Poly v2 _)
    | v2 > v1 = polyMulConst p q
polyMul p q = polyMulConst q p

polyMulConst p (Poly _ [NumCoof 0]) = Poly '_' [NumCoof 0]
polyMulConst p (Poly _ [NumCoof 1]) = p
polyMulConst (Poly _ [NumCoof x]) (Poly _ [NumCoof y]) = Poly '_' [NumCoof (x * y)]
polyMulConst (Poly v cs) k = Poly v (map (\c->simplifyPolyCoof (polyMul (complicatePolyCoof c) k)) cs)

polyMulSame p@(Poly v cs1) q@(Poly _ cs2) = Poly v (addCoofs
                                                    (replicate (fromInteger (polyDegree p + polyDegree q + 1))
                                                               (NumCoof 0))
                                                    (liftM2 (\(d1, c1) (d2, c2)
                                                                 ->(d1 + d2,
                                                                           (polyMul
                                                                            (complicatePolyCoof c1)
                                                                            (complicatePolyCoof c2))))
                                                     (zip [0..] cs1) (zip [0..] cs2)))
    where addCoofs :: [Coof] -> [(Integer, Poly)] -> [Coof]
          addCoofs cs [] = cs
          addCoofs cs ((d, poly):ps) = addCoofs (setAt cs (fromInteger d)
                                                           (simplifyPolyCoof
                                                            (polyAdd poly (complicatePolyCoof
                                                                           (cs !! (fromInteger d)))))) ps
          setAt xs d v = (take d xs) ++ [v] ++ (drop (d + 1) xs)

-- Subtraction --
polySub p q = polyAdd p (polyMul q (Poly '_' [NumCoof (-1)]))

-- Expression to canonical representation --
canon (Number x) = Poly '_' [NumCoof x]
canon (Variable x) = Poly x [NumCoof 0, NumCoof 1]
canon (Term Add t1 t2) = (canon t1) `polyAdd` (canon t2)
canon (Term Sub t1 t2) = (canon t1) `polySub` (canon t2)
canon (Term Mul t1 t2) = (canon t1) `polyMul` (canon t2)
canon (Term Div _ _) = error "Sorry, the algebra system can't handle division."

-- Canonical representation to expression --
polyToExpr (Poly v cs) = polyToExprAux (buildTerm (head revCs) v (len - 1)) (len - 2) (Poly v (tail revCs))
    where revCs = reverse cs
          len = toInteger (length cs)
          buildExp _ 0 = (Number 1)
          buildExp v 1 = (Variable v)
          buildExp v d = (Term Exp (Variable v) (Number d))
          buildMul e (NumCoof 0) = (Number 0)
          buildMul e (NumCoof 1) = e
          buildMul (Number 1) (NumCoof c) = (Number c)
          buildMul e (NumCoof c) = (Term Mul (Number c) e)
          buildMul e (PolyCoof p) = (Term Mul e (polyToExpr p))
          buildTerm (NumCoof 0) _ 0 = (Number 0)
          buildTerm (PolyCoof p) _ 0 = polyToExpr p
          buildTerm c v d = buildMul (buildExp v d) c
          polyToExprAux expr _ (Poly _ []) = expr
          polyToExprAux expr d (Poly v (c:cs))
              = let term = (buildTerm c v d)
                in case term of
                     (Number 0) -> polyToExprAux expr (d - 1) (Poly v cs)
                     term -> case expr of
                               (Number 0) -> polyToExprAux term (d - 1) (Poly v cs)
                               expr -> polyToExprAux (Term Add expr term) (d - 1) (Poly v cs)
                                           
------------
-- Driver --
------------
main = (do args <- getArgs
           let argCount = (length args)
           if (argCount > 2) || (argCount == 2 && head args /= "-r") || argCount < 1
              then putStrLn "Usage: runhaskell algebra.hs [-r] (FILE_NAME)"
              else do let fileName = if argCount == 1
                                        then head args
                                        else head $ tail args
                      input <- readFile fileName
                      parsedExpr <- parseExpr input
                      let outExpr = if (argCount == 2)
                                       then polyToExpr $ canon parsedExpr
                                       else parsedExpr
                      putStrLn $ show outExpr
                      return ())
       `catch` (\_ -> putStrLn "Error reading input file")