{-# LANGUAGE NoMonomorphismRestriction, NoImplicitPrelude, ConstraintKinds,
      TypeFamilies, RebindableSyntax, DeriveFunctor, DeriveFoldable #-}
module Knots.Graded where

import Knots.Prelude hiding ((.))
import Control.Category
import qualified Data.Map.Lazy as Map

data Graded a = Graded
    { grade      :: Int
    , components :: Map.Map Int a
    }
    deriving (Eq,Show,Read,Functor,Foldable)

graded :: Int -> [ (Int, a) ] -> Graded a
graded g = Graded g . Map.fromList

instance Default (Graded a) where
    def = Graded 0 mempty

instance (AbelianGroup a) => AbelianGroup (Graded a) where
    zero    = Graded 0 mempty
    negate  = fmap negate
    Graded i x + Graded j y
        | i /= j    = error "Graded: +: cannot add differently graded objects"
        | otherwise = Graded i (Map.unionWith (+) x y)

-- | This is practically multiplication of polynomials in one variable, if we
-- interpret the keys of the `components' map as exponents.  If we write
-- @r[X]@ for the ring of polynomials over @r@, then the `grade' can be
-- interpreted as the degree of the variable X.  This should justify how the
-- instance is defined.
instance (Ring r) => Ring (Graded r) where
    fromInteger n = Graded 0 (Map.singleton 0 $ fromInteger n)
    Graded g x * Graded h y = Graded
        { grade      = g + h
        , components = Map.fromListWith (+)
                        [ (i+j, m*n) | (i,m) <- Map.toList x, (j,n) <- Map.toList y ]
        }

instance (Monoid a, AbelianGroup a) => Monoid (Graded a) where
    mempty = Graded 0 (Map.singleton 0 $ mempty)
    Graded g x `mappend` Graded h y = Graded
        { grade = g + h
        , components =
            let x_shifted = Map.mapKeys (subtract h) x
            in  Map.intersectionWith (mappend) x_shifted y
                  `Map.union` fmap (zero `mappend`) (y Map.\\ x_shifted)
                  `Map.union` fmap (`mappend` zero) (x_shifted Map.\\ y)
        }