-- | Geometry of the degree @n!@ finite map @pi@, which just forgets the order points:
--
-- > pi : Q^n = P^1 x P^1 x ... x P^1  ->  P^n = P(Sym^n C^2)
--
-- It's clear that the degree of @pi@ restricted to an open stratum corresponding to
-- a partition @mu@ is the multinomial coefficient corresponding to @n `choose` mu@.
--
-- It is also not hard to see that the degree of @pi@ restricted to the intersection
-- of the open stratum corresponding to @mu@ with the image of the diagonal map 
-- corresponding to @nu@ equals the number of \"automorphisms\" @aut(mu) = prod (e_i!)@
-- where @mu = (1^e1 2^e2 ... k^ek)@ and the number of ways @nu@ is refinement of @mu@.
--
-- Note that for @nu=(1,1...1)@ the multinomial agrees with the number of refinements.
--
-- This module contains functions to compute these numbers.
--

{-# LANGUAGE BangPatterns, TypeSynonymInstances, FlexibleInstances #-}
module Math.RootLoci.Geometry.Forget where

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

import Data.List

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.Strict (Map)

import qualified Data.Map.Lazy as LMap

import qualified Data.Set as Set
import Data.Set (Set)

import Data.Array.IArray
import Data.Array.Unboxed
import Data.Array (Array)

import Math.RootLoci.Misc.Common
import Math.RootLoci.Misc.PTable
-- import Math.RootLoci.Geometry.Mobius

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

-- | Given a partition, we list all coarser partitions together
-- with the number of ways the input is a refinement of the
-- coarser partition.
--
-- TODO: at the moment this is just a synonym for 'countCoarseningsNaive' ...
--
countCoarsenings :: Partition -> Map Partition Integer
countCoarsenings = countCoarseningsNaive

-- | Count coarsenings (with multiplicities) which are shorter by just 1.
countDirectCoarsenings :: Partition -> Map Partition Integer
countDirectCoarsenings part = Map.fromListWith (+) list where
  list =  
    [ ( fromExponentialForm ((i1+i2,1):(i1,e1-1):(i2,e2-1):rest) , fromIntegral (e1*e2) )
    | ( (i1,e1):(i2,e2):[] , rest ) <- choose' 2 ies
    ] ++
    [ ( fromExponentialForm ((2*i,1):(i,e-2):rest) , binomial e 2 )
    | ( (i,e):[] , rest ) <- choose' 1 ies
    , e >= 2
    ]
  ies = toExponentialForm part 
  
--------------------------------------------------------------------------------

-- | Naive (very slow) implementation of 'countCoarsenings'.
countCoarseningsNaive :: Partition -> Map Partition Integer
countCoarseningsNaive = pcache count where

  count (Partition ps) = histogram (map f setps) where
    d     = length ps
    setps = map fromSetPartition $ setPartitions d :: [[[Int]]]
    arr   = listArray (1,d) ps :: UArray Int Int
    f iss = mkPartition [ sum [ arr ! k | k <- is ] | is <- iss ]

-- | Given a partition @nu@, we stratify the image of the 
-- corresponding diagonal @Delta_nu@ as usual, and list
-- the degree of @pi@ restricted to these strata
--
-- This is just counting the coarsenings, multiplied by
-- the number of \"automorphisms\" of the partition.
--
countPreimage :: Partition -> Map Partition Integer
countPreimage = pcache compute where
  compute part = Map.mapWithKey f (countCoarsenings part) 
  f q c = c * aut q

-- | The preimage counts, but the partition itself is separated out.
preimageView :: Partition -> (Integer, Map Partition Integer)
preimageView part = unsafeDeleteLookup part (countPreimage part) 

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

-- | The preimage @pi^-1(x)@ of a point under the map 
-- @pi : Q^n -> P^n@ is just a multinomial coefficient
countFullPreimage :: Partition -> Integer
countFullPreimage part@(Partition ps) = multinomial ps 
  
--------------------------------------------------------------------------------