-- | Some auxilary functions

{-# LANGUAGE CPP, BangPatterns, TypeSynonymInstances, FlexibleInstances, DeriveFunctor #-}
module Math.RootLoci.Misc.Common where

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

import Data.List
import Data.Monoid
import Data.Ratio
import Data.Ord

import Control.Monad

import Math.Combinat.Numbers
import Math.Combinat.Sign
import Math.Combinat.Partitions.Integer 
import Math.Combinat.Partitions.Set
import Math.Combinat.Sets

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

--------------------------------------------------------------------------------
-- * Pairs

data Pair a 
  = Pair a a 
  deriving (Pair a -> Pair a -> Bool
(Pair a -> Pair a -> Bool)
-> (Pair a -> Pair a -> Bool) -> Eq (Pair a)
forall a. Eq a => Pair a -> Pair a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Pair a -> Pair a -> Bool
$c/= :: forall a. Eq a => Pair a -> Pair a -> Bool
== :: Pair a -> Pair a -> Bool
$c== :: forall a. Eq a => Pair a -> Pair a -> Bool
Eq,Eq (Pair a)
Eq (Pair a)
-> (Pair a -> Pair a -> Ordering)
-> (Pair a -> Pair a -> Bool)
-> (Pair a -> Pair a -> Bool)
-> (Pair a -> Pair a -> Bool)
-> (Pair a -> Pair a -> Bool)
-> (Pair a -> Pair a -> Pair a)
-> (Pair a -> Pair a -> Pair a)
-> Ord (Pair a)
Pair a -> Pair a -> Bool
Pair a -> Pair a -> Ordering
Pair a -> Pair a -> Pair a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (Pair a)
forall a. Ord a => Pair a -> Pair a -> Bool
forall a. Ord a => Pair a -> Pair a -> Ordering
forall a. Ord a => Pair a -> Pair a -> Pair a
min :: Pair a -> Pair a -> Pair a
$cmin :: forall a. Ord a => Pair a -> Pair a -> Pair a
max :: Pair a -> Pair a -> Pair a
$cmax :: forall a. Ord a => Pair a -> Pair a -> Pair a
>= :: Pair a -> Pair a -> Bool
$c>= :: forall a. Ord a => Pair a -> Pair a -> Bool
> :: Pair a -> Pair a -> Bool
$c> :: forall a. Ord a => Pair a -> Pair a -> Bool
<= :: Pair a -> Pair a -> Bool
$c<= :: forall a. Ord a => Pair a -> Pair a -> Bool
< :: Pair a -> Pair a -> Bool
$c< :: forall a. Ord a => Pair a -> Pair a -> Bool
compare :: Pair a -> Pair a -> Ordering
$ccompare :: forall a. Ord a => Pair a -> Pair a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (Pair a)
Ord,Int -> Pair a -> ShowS
[Pair a] -> ShowS
Pair a -> String
(Int -> Pair a -> ShowS)
-> (Pair a -> String) -> ([Pair a] -> ShowS) -> Show (Pair a)
forall a. Show a => Int -> Pair a -> ShowS
forall a. Show a => [Pair a] -> ShowS
forall a. Show a => Pair a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Pair a] -> ShowS
$cshowList :: forall a. Show a => [Pair a] -> ShowS
show :: Pair a -> String
$cshow :: forall a. Show a => Pair a -> String
showsPrec :: Int -> Pair a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Pair a -> ShowS
Show,a -> Pair b -> Pair a
(a -> b) -> Pair a -> Pair b
(forall a b. (a -> b) -> Pair a -> Pair b)
-> (forall a b. a -> Pair b -> Pair a) -> Functor Pair
forall a b. a -> Pair b -> Pair a
forall a b. (a -> b) -> Pair a -> Pair b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Pair b -> Pair a
$c<$ :: forall a b. a -> Pair b -> Pair a
fmap :: (a -> b) -> Pair a -> Pair b
$cfmap :: forall a b. (a -> b) -> Pair a -> Pair b
Functor)

--------------------------------------------------------------------------------
-- * Lists

{-# SPECIALIZE sum' :: [Int] -> Int #-}
sum' :: Num a => [a] -> a
sum' :: [a] -> a
sum' = (a -> a -> a) -> a -> [a] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0 

{-# SPECIALIZE unique :: [Partition] -> [Partition] #-}
unique :: Ord a => [a] -> [a]
unique :: [a] -> [a]
unique = ([a] -> a) -> [[a]] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> a
forall a. [a] -> a
head ([[a]] -> [a]) -> ([a] -> [[a]]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [[a]]
forall a. Eq a => [a] -> [[a]]
group ([a] -> [[a]]) -> ([a] -> [a]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. Ord a => [a] -> [a]
sort

-- | Synonym for histogram
count :: Ord b => [b] -> Map b Integer
count :: [b] -> Map b Integer
count = [b] -> Map b Integer
forall b. Ord b => [b] -> Map b Integer
histogram

{-# SPECIALIZE histogram :: [Partition] -> Map Partition Integer #-}
histogram :: Ord b => [b] -> Map b Integer
histogram :: [b] -> Map b Integer
histogram [b]
xs = (Map b Integer -> b -> Map b Integer)
-> Map b Integer -> [b] -> Map b Integer
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map b Integer -> b -> Map b Integer
forall k a. (Ord k, Num a) => Map k a -> k -> Map k a
f Map b Integer
forall k a. Map k a
Map.empty [b]
xs where
  f :: Map k a -> k -> Map k a
f Map k a
old k
x = (a -> a -> a) -> k -> a -> Map k a -> Map k a
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith a -> a -> a
forall a. Num a => a -> a -> a
(+) k
x a
1 Map k a
old

#if MIN_VERSION_base(4,8,0)
-- sortOn already in base, nothing to do
#else
-- sortOn not yet in base, let's define it
sortOn :: Ord b => (a -> b) -> [a] -> [a]
sortOn f = sortBy (comparing f)
#endif

longZipWith :: (a -> c) -> (b -> c) -> (a -> b -> c) -> [a] -> [b] -> [c]
longZipWith :: (a -> c) -> (b -> c) -> (a -> b -> c) -> [a] -> [b] -> [c]
longZipWith a -> c
f b -> c
g a -> b -> c
h = [a] -> [b] -> [c]
go where
  go :: [a] -> [b] -> [c]
go (a
x:[a]
xs) (b
y:[b]
ys) = a -> b -> c
h a
x b
y c -> [c] -> [c]
forall a. a -> [a] -> [a]
: [a] -> [b] -> [c]
go [a]
xs [b]
ys
  go [a]
xs     []     = (a -> c) -> [a] -> [c]
forall a b. (a -> b) -> [a] -> [b]
map a -> c
f [a]
xs
  go []     [b]
ys     = (b -> c) -> [b] -> [c]
forall a b. (a -> b) -> [a] -> [b]
map b -> c
g [b]
ys

evens :: [a] -> [a]
evens :: [a] -> [a]
evens (a
x:[a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
forall a. [a] -> [a]
odds [a]
xs
evens []     = []

odds :: [a] -> [a]
odds :: [a] -> [a]
odds (a
_:[a]
xs) = [a] -> [a]
forall a. [a] -> [a]
evens [a]
xs
odds []     = []

interleave :: [a] -> [a] -> [a]
interleave :: [a] -> [a] -> [a]
interleave = [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
go where 
  go :: [a] -> [a] -> [a]
go (a
x:[a]
xs) (a
y:[a]
ys) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
go [a]
xs [a]
ys
  go []     []     = []
  go [a]
_      [a]
_      = String -> [a]
forall a. HasCallStack => String -> a
error String
"interleave: input lists do not have the same length"

--------------------------------------------------------------------------------
-- * Maps
  
deleteLookup :: Ord a => a -> Map a b -> (Maybe b, Map a b)
deleteLookup :: a -> Map a b -> (Maybe b, Map a b)
deleteLookup a
k Map a b
table = (a -> Map a b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
k Map a b
table, a -> Map a b -> Map a b
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete a
k Map a b
table)  

unsafeDeleteLookup :: Ord a => a -> Map a b -> (b, Map a b)
unsafeDeleteLookup :: a -> Map a b -> (b, Map a b)
unsafeDeleteLookup a
k Map a b
table = (Maybe b -> b
forall p. Maybe p -> p
fromJust (a -> Map a b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
k Map a b
table), a -> Map a b -> Map a b
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete a
k Map a b
table) where
  fromJust :: Maybe p -> p
fromJust Maybe p
mb = case Maybe p
mb of
    Just p
y  -> p
y
    Maybe p
Nothing -> String -> p
forall a. HasCallStack => String -> a
error String
"unsafeDeleteLookup: key not found"

-- | Example usage: @insertMap (:[]) (:) ...@
insertMap :: Ord k => (b -> a) -> (b -> a -> a) -> k -> b -> Map k a -> Map k a
insertMap :: (b -> a) -> (b -> a -> a) -> k -> b -> Map k a -> Map k a
insertMap b -> a
f b -> a -> a
g k
k b
y = (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter Maybe a -> Maybe a
h k
k where
  h :: Maybe a -> Maybe a
h Maybe a
mb = case Maybe a
mb of
    Maybe a
Nothing -> a -> Maybe a
forall a. a -> Maybe a
Just (b -> a
f b
y)
    Just a
x  -> a -> Maybe a
forall a. a -> Maybe a
Just (b -> a -> a
g b
y a
x)    

-- | Example usage: @buildMap (:[]) (:) ...@
buildMap :: Ord k => (b -> a) -> (b -> a -> a) -> [(k,b)] -> Map k a
buildMap :: (b -> a) -> (b -> a -> a) -> [(k, b)] -> Map k a
buildMap b -> a
f b -> a -> a
g [(k, b)]
xs = (Map k a -> (k, b) -> Map k a) -> Map k a -> [(k, b)] -> Map k a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map k a -> (k, b) -> Map k a
forall k. Ord k => Map k a -> (k, b) -> Map k a
worker Map k a
forall k a. Map k a
Map.empty [(k, b)]
xs where
  worker :: Map k a -> (k, b) -> Map k a
worker !Map k a
old (k
k,b
y) = (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter Maybe a -> Maybe a
h k
k Map k a
old where
    h :: Maybe a -> Maybe a
h Maybe a
mb = case Maybe a
mb of
      Maybe a
Nothing -> a -> Maybe a
forall a. a -> Maybe a
Just (b -> a
f b
y)
      Just a
x  -> a -> Maybe a
forall a. a -> Maybe a
Just (b -> a -> a
g b
y a
x)    

--------------------------------------------------------------------------------
-- * Partitions

-- | @aut(mu)@ is the number of symmetries of the partition mu:
--
-- > aut(mu) = prod_r (e_r)!
--
-- where @mu = (1^e1 2^e2 .. k^ek)@
aut :: Partition -> Integer
aut :: Partition -> Integer
aut Partition
part = [Integer] -> Integer
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
factorial [Int]
es where
  es :: [Int]
es = ((Int, Int) -> Int) -> [(Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Int
forall a b. (a, b) -> b
snd ([(Int, Int)] -> [Int]) -> [(Int, Int)] -> [Int]
forall a b. (a -> b) -> a -> b
$ Partition -> [(Int, Int)]
toExponentialForm Partition
part   

-- | TODO: move this into combinat
exponentVector :: Partition -> [Int]
exponentVector :: Partition -> [Int]
exponentVector Partition
p = Int -> [(Int, Int)] -> [Int]
forall a a. (Ord a, Num a, Num a) => a -> [(a, a)] -> [a]
go Int
1 (Partition -> [(Int, Int)]
toExponentialForm Partition
p) where
  go :: a -> [(a, a)] -> [a]
go a
_  []              = []
  go !a
i ef :: [(a, a)]
ef@((a
j,a
e):[(a, a)]
rest) = if a
ia -> a -> Bool
forall a. Ord a => a -> a -> Bool
<a
j 
    then a
0 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [(a, a)] -> [a]
go (a
ia -> a -> a
forall a. Num a => a -> a -> a
+a
1) [(a, a)]
ef
    else a
e a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [(a, a)] -> [a]
go (a
ia -> a -> a
forall a. Num a => a -> a -> a
+a
1) [(a, a)]
rest

--------------------------------------------------------------------------------
-- * Set partitions
 
-- | Makes set partition from a partition (simply filling up from left to right)
-- with the shape giving back the input partition
defaultSetPartition :: Partition -> SetPartition
defaultSetPartition :: Partition -> SetPartition
defaultSetPartition = [[Int]] -> SetPartition
SetPartition ([[Int]] -> SetPartition)
-> (Partition -> [[Int]]) -> Partition -> SetPartition
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Partition -> [[Int]]
linearIndices

-- | Produce linear indices from a partition @nu@ (to encode the diagonal map @Delta_nu@).
linearIndices :: Partition -> [[Int]]
linearIndices :: Partition -> [[Int]]
linearIndices (Partition [Int]
ps) = Int -> [Int] -> [[Int]]
forall t. (Num t, Enum t) => t -> [t] -> [[t]]
go Int
0 [Int]
ps where
  go :: t -> [t] -> [[t]]
go t
_  []     = []
  go !t
k (t
a:[t]
as) = [t
kt -> t -> t
forall a. Num a => a -> a -> a
+t
1..t
kt -> t -> t
forall a. Num a => a -> a -> a
+t
a] [t] -> [[t]] -> [[t]]
forall a. a -> [a] -> [a]
: t -> [t] -> [[t]]
go (t
kt -> t -> t
forall a. Num a => a -> a -> a
+t
a) [t]
as

--------------------------------------------------------------------------------
-- * Signs

{-
class IsSigned a where
  signOf :: a -> Maybe Sign

signOfNum :: (Ord a, Num a) => a -> Maybe Sign 
signOfNum x = case compare x 0 of
  LT -> Just Minus
  GT -> Just Plus
  EQ -> Nothing

instance IsSigned Int      where signOf = signOfNum
instance IsSigned Integer  where signOf = signOfNum
instance IsSigned Rational where signOf = signOfNum
-}

--------------------------------------------------------------------------------
-- * Numbers

fromRat :: Rational -> Integer
fromRat :: Rational -> Integer
fromRat Rational
r = case Rational -> Integer
forall a. Ratio a -> a
denominator Rational
r of
  Integer
1 -> Rational -> Integer
forall a. Ratio a -> a
numerator Rational
r
  Integer
_ -> String -> Integer
forall a. HasCallStack => String -> a
error String
"fromRat: not an integer"    

safeDiv :: Integer -> Integer -> Integer
safeDiv :: Integer -> Integer -> Integer
safeDiv Integer
a Integer
b = case Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
divMod Integer
a Integer
b of
  (Integer
q,Integer
0) -> Integer
q
  (Integer
q,Integer
r) -> String -> Integer
forall a. HasCallStack => String -> a
error (String -> Integer) -> String -> Integer
forall a b. (a -> b) -> a -> b
$ String
"saveDiv: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" = " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
b String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" * " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
q String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" + " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
r

--------------------------------------------------------------------------------
-- * Combinatorics

-- | Chooses (n-1) elements out of n
chooseN1 :: [a] -> [[a]]
chooseN1 :: [a] -> [[a]]
chooseN1 = [a] -> [[a]]
forall a. [a] -> [[a]]
go where
  go :: [a] -> [[a]]
go (a
x:[a]
xs) = [a]
xs [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [[a]]
go [a]
xs)
  go []     = []
  
symPolyNum :: Num a => Int -> [a] -> a
symPolyNum :: Int -> [a] -> a
symPolyNum Int
k [a]
xs = [a] -> a
sum' (([a] -> a) -> [[a]] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> a
prod' ([[a]] -> [a]) -> [[a]] -> [a]
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
choose Int
k [a]
xs) where
  sum' :: [a] -> a
sum'  = (a -> a -> a) -> a -> [a] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0
  prod' :: [a] -> a
prod' = (a -> a -> a) -> a -> [a] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' a -> a -> a
forall a. Num a => a -> a -> a
(*) a
1

--------------------------------------------------------------------------------
-- * Utility

-- | Put into parentheses
paren :: String -> String
paren :: ShowS
paren String
s = Char
'(' Char -> ShowS
forall a. a -> [a] -> [a]
: String
s String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

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

-- | Exponential form of a partition
expFormString :: Partition -> String
expFormString :: Partition -> String
expFormString Partition
p = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"," (((Int, Int) -> String) -> [(Int, Int)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> String
forall a a. (Show a, Show a) => (a, a) -> String
f [(Int, Int)]
ies) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")" where
  ies :: [(Int, Int)]
ies = Partition -> [(Int, Int)]
toExponentialForm Partition
p
  f :: (a, a) -> String
f (a
i,a
e) = a -> String
forall a. Show a => a -> String
show a
i String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"^" String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
e

extendStringL :: Int -> String -> String
extendStringL :: Int -> ShowS
extendStringL Int
k String
s = String
s String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
s) Char
' '

extendStringR :: Int -> String -> String
extendStringR :: Int -> ShowS
extendStringR Int
k String
s = Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
s) Char
' ' String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
s

--------------------------------------------------------------------------------
-- * Mathematica-formatted output

class Mathematica a where
  mathematica :: a -> String

instance Mathematica Int where
  mathematica :: Int -> String
mathematica = Int -> String
forall a. Show a => a -> String
show

instance Mathematica Integer where
  mathematica :: Integer -> String
mathematica = Integer -> String
forall a. Show a => a -> String
show

instance Mathematica String where
  mathematica :: ShowS
mathematica = ShowS
forall a. Show a => a -> String
show

instance Mathematica Partition where
  mathematica :: Partition -> String
mathematica (Partition [Int]
ps) = String
"{" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"," ((Int -> String) -> [Int] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Int -> String
forall a. Show a => a -> String
show [Int]
ps) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"}"

data Indexed a = Indexed String a

instance Mathematica a => Mathematica (Indexed a) where
  mathematica :: Indexed a -> String
mathematica (Indexed String
x a
sub) = String
"Subscript[" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
x String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"," String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Mathematica a => a -> String
mathematica a
sub String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"]"

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