-- | 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 (Eq,Ord,Show,Functor)

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

{-# SPECIALIZE sum' :: [Int] -> Int #-}
sum' :: Num a => [a] -> a
sum' = foldl' (+) 0 

{-# SPECIALIZE unique :: [Partition] -> [Partition] #-}
unique :: Ord a => [a] -> [a]
unique = map head . group . sort

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

{-# SPECIALIZE histogram :: [Partition] -> Map Partition Integer #-}
histogram :: Ord b => [b] -> Map b Integer
histogram xs = foldl' f Map.empty xs where
  f old x = Map.insertWith (+) x 1 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 f g h = go where
  go (x:xs) (y:ys) = h x y : go xs ys
  go xs     []     = map f xs
  go []     ys     = map g ys

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

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

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

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

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

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

-- | Example usage: @buildMap (:[]) (:) ...@
buildMap :: Ord k => (b -> a) -> (b -> a -> a) -> [(k,b)] -> Map k a
buildMap f g xs = foldl' worker Map.empty xs where
  worker !old (k,y) = Map.alter h k old where
    h mb = case mb of
      Nothing -> Just (f y)
      Just x  -> Just (g y 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 part = product $ map factorial es where
  es = map snd $ toExponentialForm part   

-- | TODO: move this into combinat
exponentVector :: Partition -> [Int]
exponentVector p = go 1 (toExponentialForm p) where
  go _  []              = []
  go !i ef@((j,e):rest) = if i<j 
    then 0 : go (i+1) ef
    else e : go (i+1) 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 = SetPartition . linearIndices

-- | Produce linear indices from a partition @nu@ (to encode the diagonal map @Delta_nu@).
linearIndices :: Partition -> [[Int]]
linearIndices (Partition ps) = go 0 ps where
  go _  []     = []
  go !k (a:as) = [k+1..k+a] : go (k+a) 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 r = case denominator r of
  1 -> numerator r
  _ -> error "fromRat: not an integer"    

safeDiv :: Integer -> Integer -> Integer
safeDiv a b = case divMod a b of
  (q,0) -> q
  (q,r) -> error $ "saveDiv: " ++ show a ++ " = " ++ show b ++ " * " ++ show q ++ " + " ++ show r

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

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

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

-- | Put into parentheses
paren :: String -> String
paren s = '(' : s ++ ")"

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

-- | Exponential form of a partition
expFormString :: Partition -> String
expFormString p = "(" ++ intercalate "," (map f ies) ++ ")" where
  ies = toExponentialForm p
  f (i,e) = show i ++ "^" ++ show e

extendStringL :: Int -> String -> String
extendStringL k s = s ++ replicate (k - length s) ' '

extendStringR :: Int -> String -> String
extendStringR k s = replicate (k - length s) ' ' ++ s

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

class Mathematica a where
  mathematica :: a -> String

instance Mathematica Int where
  mathematica = show

instance Mathematica Integer where
  mathematica = show

instance Mathematica String where
  mathematica = show

instance Mathematica Partition where
  mathematica (Partition ps) = "{" ++ intercalate "," (map show ps) ++ "}"

data Indexed a = Indexed String a

instance Mathematica a => Mathematica (Indexed a) where
  mathematica (Indexed x sub) = "Subscript[" ++ x ++ "," ++ mathematica sub ++ "]"

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