-- | Free modules. 
--
-- This module should be imported qualified

{-# LANGUAGE BangPatterns, FlexibleInstances, TypeSynonymInstances #-}
module Math.RootLoci.Algebra.FreeMod where

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

import Prelude   hiding ( sum , product )
import Data.List hiding ( sum , product )

import Data.Monoid
import Data.Ratio
import Data.Maybe

import Math.Combinat.Sets ( choose )

import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)

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

-- | Free module over a coefficient ring with the given base. Internally a map
-- storing the coefficients. We maintain the invariant that the coefficients
-- are never zero.
newtype FreeMod coeff base = FreeMod { unFreeMod :: Map base coeff } deriving (Eq,Show)

-- | Free module with integer coefficients
type ZMod base = FreeMod Integer  base

-- | Free module with rational coefficients
type QMod base = FreeMod Rational base

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

instance (Monoid b, Ord b, Eq c, Num c) => Num (FreeMod c b) where
  (+)    = add
  (-)    = sub
  negate = neg
  (*)    = mul
  fromInteger = konst . fromInteger
  abs    = error "FreeMod/abs"
  signum = error "FreeMod/signum"

--------------------------------------------------------------------------------
-- * Sanity checking

-- | Should be the identity function
normalize :: (Ord b, Eq c, Num c) => FreeMod c b -> FreeMod c b
normalize = FreeMod . Map.filter (/=0) . unFreeMod

-- | Safe equality testing (should be identical to @==@)
safeEq :: (Ord b, Eq b, Eq c, Num c) => FreeMod c b -> FreeMod c b -> Bool
safeEq x y = normalize x == normalize y

--------------------------------------------------------------------------------
-- * Constructing and deconstructing

-- | The additive unit
zero :: FreeMod c b
zero = FreeMod $ Map.empty

-- | A module generator
generator :: Num c => b -> FreeMod c b 
generator x = FreeMod $ Map.singleton x 1

-- | A single generator with a coefficient
singleton :: (Ord b) => b -> c -> FreeMod c b
singleton b c = FreeMod $ Map.singleton b c

-- | Conversion from list. 
-- Note that we assume here that each generator appears at most once!
fromList :: (Eq c, Num c, Ord b) => [(b,c)] -> FreeMod c b
fromList = FreeMod . Map.fromList . filter cond where
  cond (b,x) = (x/=0)

-- | Conversion to list 
toList :: FreeMod c b -> [(b,c)]
toList = Map.toList . unFreeMod

-- | Extract the coefficient of a generator
coeffOf :: (Ord b, Num c) => b -> FreeMod c b -> c
coeffOf b (FreeMod x) = case Map.lookup b x of
  Just c  -> c
  Nothing -> 0

-- | Finds the term with the largest generator (in the natural ordering of the generators)
findMaxTerm :: (Ord b) => FreeMod c b -> Maybe (b,c)
findMaxTerm (FreeMod m) = if Map.null m
  then Nothing
  else Just (Map.findMax m)

-- | Finds the term with the smallest generator (in the natural ordering of the generators)
findMinTerm :: (Ord b) => FreeMod c b -> Maybe (b,c)
findMinTerm (FreeMod m) = if Map.null m
  then Nothing
  else Just (Map.findMin m)

--------------------------------------------------------------------------------
-- * Basic operations

-- | Negation
neg :: Num c => FreeMod c b -> FreeMod c b 
neg (FreeMod m) = FreeMod (Map.map negate m)

-- | Additions
add :: (Ord b, Eq c, Num c) => FreeMod c b -> FreeMod c b -> FreeMod c b
add (FreeMod m1) (FreeMod m2) = FreeMod (Map.mergeWithKey f id id m1 m2) where
  f _ x y = case x+y of { 0 -> Nothing ; z -> Just z }

-- | Subtraction
sub :: (Ord b, Eq c, Num c) => FreeMod c b -> FreeMod c b -> FreeMod c b
sub (FreeMod m1) (FreeMod m2) = FreeMod (Map.mergeWithKey f id (Map.map negate) m1 m2) where
  f _ x y = case x-y of { 0 -> Nothing ; z -> Just z }

-- | Scaling by a number
scale :: (Ord b, Eq c, Num c) => c -> FreeMod c b -> FreeMod c b
scale 0 _           = zero
scale x (FreeMod m) = FreeMod (Map.mapMaybe f m) where
  f y = case x*y of { 0 -> Nothing ; z -> Just z }

-- | Dividing by a number (assuming that the coefficient ring is integral, and each coefficient
-- is divisible by the given number)
invScale :: (Ord b, Eq c, Integral c, Show c) => c -> FreeMod c b -> FreeMod c b
invScale d (FreeMod m) = FreeMod (Map.mapMaybe f m) where
  f a = case divMod a d of
    (b,0) -> case b of { 0 -> Nothing ; z -> Just z }
    _     -> error $ "FreeMod/invScale: not divisible by " ++ show d

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

-- | Summation
sum :: (Ord b, Eq c, Num c) => [FreeMod c b] -> FreeMod c b
sum []  = zero
sum zms = (foldl1' add) zms

-- | Linear combination
linComb :: (Ord b, Eq c, Num c) => [(c, FreeMod c b)] -> FreeMod c b
linComb = sumWith where

   sumWith :: (Ord b, Eq c, Num c) => [(c, FreeMod c b)] -> FreeMod c b
   sumWith []  = zero
   sumWith zms = sum [ scale c zm | (c,zm) <- zms ]

-- | Expand each generator into a term in another module and then sum the results
flatMap :: (Ord b1, Ord b2, Eq c, Num c) => (b1 -> FreeMod c b2) -> FreeMod c b1 -> FreeMod c b2
flatMap f = sum . map g . Map.assocs . unFreeMod where
  g (x,c) = scale c (f x)

flatMap' :: (Ord b1, Ord b2, Eq c2, Num c2) => (c1 -> c2) -> (b1 -> FreeMod c2 b2) -> FreeMod c1 b1 -> FreeMod c2 b2
flatMap' embed f = sum . map g . Map.assocs . unFreeMod where
  g (x,c) = scale (embed c) (f x)

-- | The histogram of a multiset of generators is naturally an element of the given Z-module.
{-# SPECIALIZE histogram :: Ord b => [b] -> ZMod b #-} 
histogram :: (Ord b, Num c) => [b] -> FreeMod c b
histogram xs = FreeMod $ foldl' f Map.empty xs where
  f old x = Map.insertWith (+) x 1 old
  
--------------------------------------------------------------------------------
-- * Rings

-- | The multiplicative unit
one :: (Monoid b, Num c) => FreeMod c b
one = konst 1

-- | A constant
konst :: (Monoid b) => c -> FreeMod c b
konst c = FreeMod (Map.singleton mempty c)

-- | Multiplying two ring elements
mul :: (Ord b, Monoid b, Eq c, Num c) => FreeMod c b -> FreeMod c b -> FreeMod c b
mul xx yy = sum [ (f x c) | (x,c) <- toList xx ] where
  f x c = FreeMod $ Map.fromList [ (x<>y, cd) | (y,d) <- ylist , let cd = c*d , cd /= 0 ]
  ylist = toList yy

-- | Product
product :: (Ord b, Monoid b, Eq c, Num c) => [FreeMod c b] -> FreeMod c b
product []  = generator mempty
product xs  = foldl1' mul xs

-- | Multiplies by a monomial
mulMonom :: (Ord b, Monoid b) => b -> FreeMod c b -> FreeMod c b
mulMonom monom = FreeMod . Map.mapKeys (mappend monom) . unFreeMod

--------------------------------------------------------------------------------
-- * Misc

-- | A symmetric polynomial of some generators
symPoly :: (Ord a, Monoid a) => Int -> [a] -> ZMod a
symPoly k xs = fromList $ map (\x -> (x,1)) $ (map mconcat $ choose k xs) 

-- | Changing the base set
mapBase :: (Ord a, Ord b) => (a -> b) -> FreeMod c a -> FreeMod c b
mapBase f = onFreeMod (Map.mapKeys f)

-- | Changing the coefficient ring
mapCoeff :: (Ord b) => (c1 -> c2) -> FreeMod c1 b -> FreeMod c2 b
mapCoeff f = onFreeMod' (Map.map f)

-- | Extract a subset of terms
filterBase :: (Ord a, Ord b) => (a -> Maybe b) -> FreeMod c a -> FreeMod c b
filterBase f = onFreeMod (Map.fromList . mapMaybe g . Map.toList) where
  g (k,x) = case f k of { Just k' -> Just (k',x) ; Nothing -> Nothing }

onFreeMod :: (Ord a, Ord b) => (Map a c -> Map b c) -> FreeMod c a -> FreeMod c b
onFreeMod f = FreeMod . f . unFreeMod

onFreeMod' :: (Ord a, Ord b) => (Map a c -> Map b d) -> FreeMod c a -> FreeMod d b
onFreeMod' f = FreeMod . f . unFreeMod

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