Fun with Polynomials

gistfile1.hs

{-# LANGUAGE FlexibleContexts #-}
module Algebra.Monomial where
import Data.List
import Data.Function
import Text.Parsec hiding (parse)
import qualified Text.Parsec.Token as P
import qualified Text.Parsec.Language as P
import Control.Applicative ((<$>))
import Data.Functor.Identity

data Polynomial = T Term
                | Polynomial :+: Polynomial
                | Polynomial :*: Polynomial
                deriving(Show, Eq, Ord)

data Term = Term {
        coefficent :: Int,
        power   :: Int
    }
    deriving(Show, Eq, Ord)

simplify :: Polynomial -> Polynomial
simplify p = combinePowers $ go $ distributeMulti p where
    go x = combine x' where
       x' = case x of
                a :*: b -> go a :*: go  b
                a :+: b -> go a :+: go b
                _ -> x
    
    combine x = case x of 
        (T a) :*: (T b) -> T $ multi a b
        (T a) :+: (T b)
            | power a == power b -> T $ add a b
            | otherwise -> x
        _ -> x

distributeMulti :: Polynomial -> Polynomial
distributeMulti m = go m where
    go x = case x of
        a :*: (b :+: c) -> (a :*: b) :+: (go (a :*: c))
        a :*: b         -> go a :*: go b
        a :+: b         -> go a :+: go b
        _ -> x

-- only works when there is no multiplication
toList :: Polynomial -> [Term]
toList p = go p where
    go x = case x of
        a :+: b -> go a ++ go b
        a :*: b -> error "toList had a :*:"
        T x     -> [x]

fromList :: [Term] -> Polynomial
fromList e = case e of
    x : [] -> T x
    x:y:[] -> T x :+: T y
    x : xs -> T x :+: fromList xs

combinePowers :: Polynomial -> Polynomial
combinePowers = fromList . map (foldl1 add) . 
    groupBy ((==) `on` power) . sortBy (compare `on` power) . toList
    
add (Term x _) (Term y p)     = Term (x + y) p
multi (Term x px) (Term y py) = Term (x * y) (px + py)


ppr :: Polynomial -> String
ppr m = case m of
    a :+: b      -> ppr a  ++ " + " ++ ppr b
    a :*: b      -> ppr a  ++ " * " ++ ppr b
    T (Term c p) -- This is total shit but it is late
        | c > 1 && p == 0  -> show c
        | c == 0           -> "0"
        | c > 1 && p > 1   -> show c ++ "x^"  ++ show p
        | c > 1            -> show c ++ "x"
        | c == 1 && p == 1 -> "x"
        | p > 1            -> "x^" ++ show p
        
    
parse str = runParser pMonomial () "" str
    
pMonomial = foldl chainl1 (parens pMonomial <|> (T <$> pTerm)) [pMulti, pAdd]

pTerm = do
    spaces
    c <- option 1 (fromIntegral <$> natural)
    p <- try (string "x^" >> fromIntegral <$> integer) <|> (string "x" >> return 1)
    spaces
    return $ Term c p

pAdd   = binOp "+" (:+:)
pMulti = binOp "*" (:*:)

binOp :: String -> b -> ParsecT String u Identity b
binOp x rest = do
    string x
    return rest
    
parens  = P.parens  P.haskell
natural = P.natural P.haskell
integer = P.integer P.haskell